Refactor NodeStack and NodePath (#275)

This PR adds new `NodePath` helper functions and pushes some of the responsibilities from the `NodeStack` to the `NodePath`. The former is designed to be a mutable stack that updates as we traverse the tree, whereas the latter is designed to hold an immutable path of nodes finishing in a specific typed node.

packages/renderers-js/src/getRenderMapVisitor.ts

@@ -16,6 +16,7 @@ import {
 import { RenderMap } from '@codama/renderers-core';
 import {
     extendVisitor,
+    findProgramNodeFromPath,
     getResolvedInstructionInputsVisitor,
     LinkableDictionary,
     NodeStack,
@@ -141,13 +142,14 @@ export function getRenderMapVisitor(options: GetRenderMapOptions = {}) {
         v =>
             extendVisitor(v, {
                 visitAccount(node) {
-                    if (!stack.getProgram()) {
+                    const accountPath = stack.getPath('accountNode');
+                    if (!findProgramNodeFromPath(accountPath)) {
                         throw new Error('Account must be visited inside a program.');
                     }
 
                     const scope = {
                         ...globalScope,
-                        accountPath: stack.getPath('accountNode'),
+                        accountPath,
                         typeManifest: visit(node, typeManifestVisitor),
                     };
 
@@ -221,7 +223,8 @@ export function getRenderMapVisitor(options: GetRenderMapOptions = {}) {
                 },
 
                 visitInstruction(node) {
-                    if (!stack.getProgram()) {
+                    const instructionPath = stack.getPath('instructionNode');
+                    if (!findProgramNodeFromPath(instructionPath)) {
                         throw new Error('Instruction must be visited inside a program.');
                     }
 
@@ -236,7 +239,7 @@ export function getRenderMapVisitor(options: GetRenderMapOptions = {}) {
                                 strict: nameApi.dataType(instructionExtraName),
                             }),
                         ),
-                        instructionPath: stack.getPath('instructionNode'),
+                        instructionPath,
                         renamedArgs: getRenamedArgsMap(node),
                         resolvedInputs: visit(node, resolvedInstructionInputVisitor),
                     };
@@ -289,11 +292,12 @@ export function getRenderMapVisitor(options: GetRenderMapOptions = {}) {
                 },
 
                 visitPda(node) {
-                    if (!stack.getProgram()) {
+                    const pdaPath = stack.getPath('pdaNode');
+                    if (!findProgramNodeFromPath(pdaPath)) {
                         throw new Error('Account must be visited inside a program.');
                     }
 
-                    const scope = { ...globalScope, pdaPath: stack.getPath('pdaNode') };
+                    const scope = { ...globalScope, pdaPath };
                     const pdaFunctionFragment = getPdaFunctionFragment(scope);
                     const imports = new ImportMap().mergeWith(pdaFunctionFragment);
 

packages/renderers-js/src/getTypeManifestVisitor.ts

@@ -15,6 +15,7 @@ import {
 } from '@codama/nodes';
 import {
     extendVisitor,
+    findLastNodeFromPath,
     LinkableDictionary,
     NodeStack,
     pipe,
@@ -830,8 +831,9 @@ export function getTypeManifestVisitor(input: {
                     }
 
                     // Check if we are inside an instruction or account to use discriminator constants when available.
-                    const instructionNode = stack.find('instructionNode');
-                    const accountNode = stack.find('accountNode');
+                    const parentPath = stack.getPath();
+                    const instructionNode = findLastNodeFromPath(parentPath, 'instructionNode');
+                    const accountNode = findLastNodeFromPath(parentPath, 'accountNode');
                     const discriminatorPrefix = instructionNode ? instructionNode.name : accountNode?.name;
                     const discriminators =
                         (instructionNode ? instructionNode.discriminators : accountNode?.discriminators) ?? [];

packages/validators/README.md

@@ -37,7 +37,7 @@ type ValidationItem = {
     // The node that the validation item is related to.
     node: Node;
     // The path of nodes that led to the node above (including the node itself).
-    path: readonly Node[];
+    path: NodePath;
 };
 ```
 

packages/validators/src/ValidationItem.ts

@@ -1,5 +1,5 @@
 import { Node } from '@codama/nodes';
-import { NodeStack } from '@codama/visitors-core';
+import { NodePath, NodeStack } from '@codama/visitors-core';
 
 export const LOG_LEVELS = ['debug', 'trace', 'info', 'warn', 'error'] as const;
 export type LogLevel = (typeof LOG_LEVELS)[number];
@@ -8,20 +8,20 @@ export type ValidationItem = {
     level: LogLevel;
     message: string;
     node: Node;
-    path: readonly Node[];
+    path: NodePath;
 };
 
 export function validationItem(
     level: LogLevel,
     message: string,
     node: Node,
-    stack: Node[] | NodeStack,
+    path: NodePath | NodeStack,
 ): ValidationItem {
     return {
         level,
         message,
         node,
-        path: Array.isArray(stack) ? [...stack] : stack.all(),
+        path: Array.isArray(path) ? path : (path as NodeStack).getPath(),
     };
 }
 

packages/validators/src/getValidationItemsVisitor.ts

@@ -47,7 +47,7 @@ export function getValidationItemsVisitor(): Visitor<readonly ValidationItem[]>
                     const items = [] as ValidationItem[];
                     if (!node.name) {
                         items.push(validationItem('error', 'Pointing to a defined type with no name.', node, stack));
-                    } else if (!linkables.has(stack.getPath())) {
+                    } else if (!linkables.has(stack.getPath(node.kind))) {
                         items.push(
                             validationItem(
                                 'error',

packages/visitors-core/README.md

@@ -442,19 +442,11 @@ const lastNode = nodeStack.pop();
 // Peek at the last node in the stack.
 const lastNode = nodeStack.peek();
 // Get all the nodes in the stack as an array.
-const nodes = nodeStack.all();
-// Get the closest node in the stack matching one or several node kinds.
-const nodes = nodeStack.find('accountNode');
-// Get the closest program node in the stack.
-const nodes = nodeStack.getProgram();
-// Get the closest instruction node in the stack.
-const nodes = nodeStack.getInstruction();
+const path = nodeStack.getPath();
 // Check if the stack is empty.
 const isEmpty = nodeStack.isEmpty();
 // Clone the stack.
 const clonedStack = nodeStack.clone();
-// Get a string representation of the stack.
-const stackString = nodeStack.toString();
 ```
 
 ### `recordNodeStackVisitor`
@@ -470,7 +462,7 @@ const visitor = pipe(
     v => recordNodeStackVisitor(v, stack),
     v =>
         interceptVisitor(v, (node, next) => {
-            console.log(stack.clone().toString());
+            console.log(nodePathToString(stack.getPath()));
             return next(node);
         }),
 );

packages/visitors-core/src/NodePath.ts

@@ -1,6 +1,8 @@
 import { assertIsNode, GetNodeFromKind, InstructionNode, isNode, Node, NodeKind, ProgramNode } from '@codama/nodes';
 
-export type NodePath<TNode extends Node = Node> = readonly [...Node[], TNode];
+export type NodePath<TNode extends Node | undefined = undefined> = TNode extends undefined
+    ? readonly Node[]
+    : readonly [...Node[], TNode];
 
 export function getLastNodeFromPath<TNode extends Node>(path: NodePath<TNode>): TNode {
     return path[path.length - 1] as TNode;
@@ -47,16 +49,30 @@ export function getNodePathUntilLastNode<TKind extends NodeKind>(
     return path.slice(0, lastIndex + 1) as unknown as NodePath<GetNodeFromKind<TKind>>;
 }
 
+function isNotEmptyNodePath(path: NodePath | null | undefined): path is NodePath<Node> {
+    return !!path && path.length > 0;
+}
+
 export function isNodePath<TKind extends NodeKind>(
     path: NodePath | null | undefined,
     kind: TKind | TKind[],
 ): path is NodePath<GetNodeFromKind<TKind>> {
-    return isNode(path ? getLastNodeFromPath(path) : null, kind);
+    return isNode(isNotEmptyNodePath(path) ? getLastNodeFromPath<Node>(path) : null, kind);
 }
 
 export function assertIsNodePath<TKind extends NodeKind>(
     path: NodePath | null | undefined,
     kind: TKind | TKind[],
 ): asserts path is NodePath<GetNodeFromKind<TKind>> {
-    assertIsNode(path ? getLastNodeFromPath(path) : null, kind);
+    assertIsNode(isNotEmptyNodePath(path) ? getLastNodeFromPath<Node>(path) : null, kind);
+}
+
+export function nodePathToStringArray(path: NodePath): string[] {
+    return path.map((node): string => {
+        return 'name' in node ? `[${node.kind}]${node.name}` : `[${node.kind}]`;
+    });
+}
+
+export function nodePathToString(path: NodePath): string {
+    return nodePathToStringArray(path).join(' > ');
 }

packages/visitors-core/src/NodeSelector.ts

@@ -53,7 +53,8 @@ export const getNodeSelectorFunction = (selector: NodeSelector): NodeSelectorFun
     const nodeSelectors = selector.split('.');
     const lastNodeSelector = nodeSelectors.pop() as string;
 
-    return (node, stack) => checkNode(node, lastNodeSelector) && checkStack(stack.all() as Node[], [...nodeSelectors]);
+    return (node, stack) =>
+        checkNode(node, lastNodeSelector) && checkStack(stack.getPath() as Node[], [...nodeSelectors]);
 };
 
 export const getConjunctiveNodeSelectorFunction = (selector: NodeSelector | NodeSelector[]): NodeSelectorFunction => {

packages/visitors-core/src/NodeStack.ts

@@ -1,98 +1,72 @@
-import {
-    assertIsNode,
-    GetNodeFromKind,
-    InstructionNode,
-    Node,
-    NodeKind,
-    ProgramNode,
-    REGISTERED_NODE_KINDS,
-} from '@codama/nodes';
+import { GetNodeFromKind, Node, NodeKind } from '@codama/nodes';
 
-import { findLastNodeFromPath, NodePath } from './NodePath';
+import { assertIsNodePath, NodePath } from './NodePath';
+
+type MutableNodePath = Node[];
 
 export class NodeStack {
     /**
-     * Contains all the node stacks saved during the traversal.
+     * Contains all the node paths saved during the traversal.
      *
-     * - The very last stack is the current stack which is being
+     * - The very last path is the current path which is being
      *   used during the traversal.
-     * - The other stacks can be used to save and restore the
-     *   current stack when jumping to different parts of the tree.
+     * - The other paths can be used to save and restore the
+     *   current path when jumping to different parts of the tree.
      *
-     * There must at least be one stack in the heap at all times.
+     * There must at least be one path in the stack at all times.
      */
-    private readonly heap: [...Node[][], Node[]];
+    private readonly stack: [...MutableNodePath[], MutableNodePath];
 
-    constructor(...heap: readonly [...(readonly (readonly Node[])[]), readonly Node[]] | readonly []) {
-        this.heap = heap.length === 0 ? [[]] : ([...heap.map(nodes => [...nodes])] as [...Node[][], Node[]]);
+    constructor(...stack: readonly [...(readonly NodePath[]), NodePath] | readonly []) {
+        this.stack =
+            stack.length === 0
+                ? [[]]
+                : ([...stack.map(nodes => [...nodes])] as [...MutableNodePath[], MutableNodePath]);
     }
 
-    public get stack(): Node[] {
-        return this.heap[this.heap.length - 1];
+    private get currentPath(): MutableNodePath {
+        return this.stack[this.stack.length - 1];
     }
 
     public push(node: Node): void {
-        this.stack.push(node);
+        this.currentPath.push(node);
     }
 
     public pop(): Node | undefined {
-        return this.stack.pop();
+        return this.currentPath.pop();
     }
 
     public peek(): Node | undefined {
-        return this.isEmpty() ? undefined : this.stack[this.stack.length - 1];
+        return this.isEmpty() ? undefined : this.currentPath[this.currentPath.length - 1];
     }
 
-    public pushStack(newStack: readonly Node[] = []): void {
-        this.heap.push([...newStack]);
+    public pushPath(newPath: NodePath = []): void {
+        this.stack.push([...newPath]);
     }
 
-    public popStack(): readonly Node[] {
-        const oldStack = this.heap.pop() as Node[];
-        if (this.heap.length === 0) {
+    public popPath(): NodePath {
+        if (this.stack.length === 0) {
             // TODO: Coded error
-            throw new Error('The heap of stacks can never be empty.');
+            throw new Error('The stack of paths can never be empty.');
         }
-        return [...oldStack] as readonly Node[];
-    }
-
-    public find<TKind extends NodeKind>(kind: TKind | TKind[]): GetNodeFromKind<TKind> | undefined {
-        return findLastNodeFromPath([...this.stack] as unknown as NodePath<GetNodeFromKind<TKind>>, kind);
-    }
-
-    public getProgram(): ProgramNode | undefined {
-        return this.find('programNode');
-    }
-
-    public getInstruction(): InstructionNode | undefined {
-        return this.find('instructionNode');
+        return [...this.stack.pop()!];
     }
 
-    public all(): readonly Node[] {
-        return [...this.stack];
-    }
-
-    public getPath<TKind extends NodeKind>(kind?: TKind | TKind[]): NodePath<GetNodeFromKind<TKind>> {
-        const node = this.peek();
-        assertIsNode(node, kind ?? REGISTERED_NODE_KINDS);
-        return [...this.stack] as unknown as NodePath<GetNodeFromKind<TKind>>;
+    public getPath(): NodePath;
+    public getPath<TKind extends NodeKind>(kind: TKind | TKind[]): NodePath<GetNodeFromKind<TKind>>;
+    public getPath<TKind extends NodeKind>(kind?: TKind | TKind[]): NodePath {
+        const path = [...this.currentPath];
+        if (kind) {
+            assertIsNodePath(path, kind);
+        }
+        return path;
     }
 
     public isEmpty(): boolean {
-        return this.stack.length === 0;
+        return this.currentPath.length === 0;
     }
 
     public clone(): NodeStack {
-        return new NodeStack(...this.heap);
-    }
-
-    public toString(): string {
-        return this.toStringArray().join(' > ');
-    }
-
-    public toStringArray(): string[] {
-        return this.stack.map((node): string => {
-            return 'name' in node ? `[${node.kind}]${node.name}` : `[${node.kind}]`;
-        });
+        return new NodeStack(...this.stack);
     }
 }

packages/visitors-core/src/recordLinkablesVisitor.ts

@@ -18,7 +18,7 @@ export function getRecordLinkablesVisitor<TNodeKind extends NodeKind>(
         v =>
             interceptVisitor(v, (node, next) => {
                 if (isNode(node, LINKABLE_NODES)) {
-                    linkables.recordPath(stack.getPath());
+                    linkables.recordPath(stack.getPath(LINKABLE_NODES));
                 }
                 return next(node);
             }),

packages/visitors-core/test/bottomUpTransformerVisitor.test.ts

@@ -94,7 +94,7 @@ test('it can transform nodes using multiple node selectors', () => {
     // - the second one selects all nodes with more than one ancestor.
     const visitor = bottomUpTransformerVisitor([
         {
-            select: ['[numberTypeNode]', (_, nodeStack) => nodeStack.all().length > 1],
+            select: ['[numberTypeNode]', (_, nodeStack) => nodeStack.getPath().length > 1],
             transform: () => stringTypeNode('utf8'),
         },
     ]);

packages/visitors-core/test/recordNodeStackVisitor.test.ts

@@ -24,7 +24,7 @@ test('it records the current node stack of a visit', () => {
 
     // Then we expect the number stacks to have been recorded.
     expect(numberStacks.length).toBe(1);
-    expect(numberStacks[0].all()).toEqual([node, node.type]);
+    expect(numberStacks[0].getPath()).toEqual([node, node.type]);
 
     // And the current node stack to be empty.
     expect(stack.isEmpty()).toBe(true);
@@ -52,5 +52,5 @@ test('it includes the current node when applied last', () => {
     // Then we expect the number stacks to have been recorded
     // such that the number node themselves are included in the stack.
     expect(numberStacks.length).toBe(1);
-    expect(numberStacks[0].all()).toEqual([node, node.type, (node.type as TupleTypeNode).items[0]]);
+    expect(numberStacks[0].getPath()).toEqual([node, node.type, (node.type as TupleTypeNode).items[0]]);
 });

packages/visitors-core/test/topDownTransformerVisitor.test.ts

@@ -102,7 +102,7 @@ test('it can transform nodes using multiple node selectors', () => {
     // - the second one selects all nodes with more than one ancestor.
     const visitor = topDownTransformerVisitor([
         {
-            select: ['[numberTypeNode]', (_, nodeStack) => nodeStack.all().length > 1],
+            select: ['[numberTypeNode]', (_, nodeStack) => nodeStack.getPath().length > 1],
             transform: node => numberTypeNode('u64') as typeof node,
         },
     ]);

packages/visitors/src/updateAccountsVisitor.ts

@@ -12,7 +12,11 @@ import {
     programNode,
     transformNestedTypeNode,
 } from '@codama/nodes';
-import { BottomUpNodeTransformerWithSelector, bottomUpTransformerVisitor } from '@codama/visitors-core';
+import {
+    BottomUpNodeTransformerWithSelector,
+    bottomUpTransformerVisitor,
+    findProgramNodeFromPath,
+} from '@codama/visitors-core';
 
 import { renameStructNode } from './renameHelpers';
 
@@ -37,26 +41,27 @@ export function updateAccountsVisitor(map: Record<string, AccountUpdates>) {
                         assertIsNode(node, 'accountNode');
                         if ('delete' in updates) return null;
 
+                        const programNode = findProgramNodeFromPath(stack.getPath())!;
                         const { seeds, pda, ...assignableUpdates } = updates;
                         let newPda = node.pda;
                         if (pda && seeds !== undefined) {
                             newPda = pda;
                             pdasToUpsert.push({
                                 pda: pdaNode({ name: pda.name, seeds }),
-                                program: stack.getProgram()!.name,
+                                program: programNode.name,
                             });
                         } else if (pda) {
                             newPda = pda;
                         } else if (seeds !== undefined && node.pda) {
                             pdasToUpsert.push({
                                 pda: pdaNode({ name: node.pda.name, seeds }),
-                                program: stack.getProgram()!.name,
+                                program: programNode.name,
                             });
                         } else if (seeds !== undefined) {
                             newPda = pdaLinkNode(newName ?? node.name);
                             pdasToUpsert.push({
                                 pda: pdaNode({ name: newName ?? node.name, seeds }),
-                                program: stack.getProgram()!.name,
+                                program: programNode.name,
                             });
                         }