[2.x] Add withPath and visitPath helpers to NodeStack

packages/dynamic-codecs/src/codecs.ts

@@ -154,12 +154,7 @@ export function getNodeValueCodecVisitor(
         node: TLinkNode,
     ) => {
         const path = linkables.getPathOrThrow(stack.getPath(node.kind) as NodePath<TLinkNode>);
-        stack.pushPath(path);
-        try {
-            return visit(getLastNodeFromPath(path), visitor);
-        } finally {
-            stack.popPath();
-        }
+        return stack.visitPath(path, visitor);
     };
 
     const getCollectionCodec = (item: Codec<unknown>, count: CountNode): Codec<unknown[]> => {

packages/visitors-core/README.md

@@ -510,6 +510,17 @@ stack.pushPath([rootNode, programNode, linkableNode]);
 const previousPath = stack.popPath();
 ```
 
+Prefer the `withPath` and `visitPath` helpers, which restore the previous path even if an error is thrown. `withPath` runs a callback with the given path as the current path, whereas `visitPath` visits the last node of the given path using the rest of that path as the current path. Since `visitPath` relies on the visitor to record the visited node, the visitor must record its nodes on the same stack (e.g. using `recordNodeStackVisitor`). Otherwise, use `withPath` instead.
+
+```ts
+// Run a callback with another path as the current path.
+const fields = stack.withPath([rootNode, programNode, definedType], () => getFields(definedType.type));
+
+// Visit a linked node with its own path as the current path.
+const linkedPath = linkables.getPathOrThrow(stack.getPath('definedTypeLinkNode'));
+const result = stack.visitPath(linkedPath, visitor);
+```
+
 ### `recordNodeStackVisitor`
 
 The `recordNodeStackVisitor` function gives us a convenient way to record the stack of each node currently being visited. It accepts a base visitor and an empty `NodeStack` instance that will automatically be pushed and popped as the visitor traverses the nodes. This means that we can inject the `NodeStack` instance into another extension of the visitor to access the stack whilst visiting the nodes.

packages/visitors-core/src/NodeStack.ts

@@ -1,7 +1,8 @@
 import { CODAMA_ERROR__VISITORS__CANNOT_REMOVE_LAST_PATH_IN_NODE_STACK, CodamaError } from '@codama/errors';
 import { GetNodeFromKind, Node, NodeKind } from '@codama/nodes';
 
-import { assertIsNodePath, NodePath, nodePathToString } from './NodePath';
+import { assertIsNodePath, getLastNodeFromPath, NodePath, nodePathToString } from './NodePath';
+import { visit, Visitor } from './visitor';
 
 type MutableNodePath = Node[];
 
@@ -52,6 +53,46 @@ export class NodeStack {
         return [...this.stack.pop()!];
     }
 
+    /**
+     * Run `callback` with `path` as the current path, then restore the
+     * previous one, even if `callback` throws. Use it to jump to another
+     * part of the tree, e.g. to the definition of a linked node.
+     *
+     * @example
+     * ```ts
+     * const fields = stack.withPath(definedTypePath, () => getFields(definedType.type));
+     * ```
+     */
+    public withPath<T>(path: NodePath, callback: () => T): T {
+        this.pushPath(path);
+        try {
+            return callback();
+        } finally {
+            this.popPath();
+        }
+    }
+
+    /**
+     * Visit the last node of `path` with the rest of `path` as the current
+     * path, then restore the previous one, even if the visit throws.
+     *
+     * The visitor must record its nodes on this stack (e.g. using
+     * `recordNodeStackVisitor`), which adds the visited node back to the
+     * path. Otherwise, use {@link NodeStack.withPath} instead.
+     *
+     * @example
+     * ```ts
+     * const linkedPath = linkables.getPathOrThrow(stack.getPath('definedTypeLinkNode'));
+     * return stack.visitPath(linkedPath, self);
+     * ```
+     */
+    public visitPath<TNode extends Node, TReturn>(
+        path: NodePath<TNode>,
+        visitor: Visitor<TReturn, TNode['kind']>,
+    ): TReturn {
+        return this.withPath(path.slice(0, -1), () => visit(getLastNodeFromPath(path), visitor));
+    }
+
     public getPath(): NodePath;
     public getPath<TKind extends NodeKind>(kind: TKind | TKind[]): NodePath<GetNodeFromKind<TKind>>;
     public getPath<TKind extends NodeKind>(kind?: TKind | TKind[]): NodePath {

packages/visitors-core/src/getByteSizeVisitor.ts

@@ -100,12 +100,7 @@ export function getByteSizeVisitor(
                     // This prevents infinite recursion by assuming cyclic types don't have a fixed size.
                     if (definedTypeStack.includes(linkedDefinedType.identifier)) return null;
 
-                    stack.pushPath(linkedDefinedPath);
-                    try {
-                        return visit(linkedDefinedType, self);
-                    } finally {
-                        stack.popPath();
-                    }
+                    return stack.visitPath(linkedDefinedPath, self);
                 },
 
                 visitEnumType(node, { self }) {

packages/visitors-core/src/getInstructionDataFields.ts

@@ -48,12 +48,7 @@ export function getInstructionDataFields(
             const definedType = getLastNodeFromPath(linkedPath);
             if (walkedDefinedTypes.has(definedType)) return;
             walkedDefinedTypes.add(definedType);
-            stack.pushPath(linkedPath);
-            try {
-                walk(definedType.type, prefix);
-            } finally {
-                stack.popPath();
-            }
+            stack.withPath(linkedPath, () => walk(definedType.type, prefix));
             return;
         }
         if (!isNode(type, 'structTypeNode')) return;

packages/visitors-core/src/getMaxByteSizeVisitor.ts

@@ -93,12 +93,7 @@ export function getMaxByteSizeVisitor(
                     // This prevents infinite recursion by assuming cyclic types don't have a fixed size.
                     if (definedTypeStack.includes(linkedDefinedType.identifier)) return null;
 
-                    stack.pushPath(linkedDefinedPath);
-                    try {
-                        return visit(linkedDefinedType, self);
-                    } finally {
-                        stack.popPath();
-                    }
+                    return stack.visitPath(linkedDefinedPath, self);
                 },
 
                 visitEnumType(node, { self }) {

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

@@ -0,0 +1,64 @@
+import { definedTypeNode, integerTypeNode, programNode, publicKeyTypeNode, rootNode } from '@codama/nodes';
+import { expect, test } from 'vitest';
+
+import { NodeStack, pipe, recordNodeStackVisitor, tapVisitor, voidVisitor } from '../src';
+
+const slot = definedTypeNode({ identifier: 'slot', type: integerTypeNode('u64') });
+const owner = definedTypeNode({ identifier: 'owner', type: publicKeyTypeNode() });
+const root = rootNode(programNode({ definedTypes: [slot, owner], identifier: 'myProgram', publicKey: '1111' }));
+
+test('withPath runs the callback with the given path and restores the previous one', () => {
+    // Given a stack within the owner type.
+    const stack = new NodeStack([root, root.program, owner]);
+
+    // When we run a callback with the path of the slot type.
+    const result = stack.withPath([root, root.program, slot], () => stack.getPath());
+
+    // Then the callback sees that path, and the previous path is restored.
+    expect(result).toStrictEqual([root, root.program, slot]);
+    expect(stack.getPath()).toStrictEqual([root, root.program, owner]);
+});
+
+test('withPath restores the previous path when the callback throws', () => {
+    const stack = new NodeStack([root, root.program, owner]);
+    expect(() =>
+        stack.withPath([root, root.program, slot], () => {
+            throw new Error('boom');
+        }),
+    ).toThrow('boom');
+    expect(stack.getPath()).toStrictEqual([root, root.program, owner]);
+});
+
+test('visitPath visits the last node of the path with the path recorded once', () => {
+    // Given a stack within the owner type and a visitor recording the paths of integer types.
+    const stack = new NodeStack([root, root.program, owner]);
+    const paths: unknown[] = [];
+    const visitor = pipe(
+        voidVisitor(),
+        v => tapVisitor(v, 'integerTypeNode', () => paths.push(stack.getPath())),
+        v => recordNodeStackVisitor(v, stack),
+    );
+
+    // When we visit the path of the slot type.
+    stack.visitPath([root, root.program, slot], visitor);
+
+    // Then its type is visited with the slot type recorded once, rather than twice.
+    expect(paths).toStrictEqual([[root, root.program, slot, slot.type]]);
+
+    // And the previous path is restored.
+    expect(stack.getPath()).toStrictEqual([root, root.program, owner]);
+});
+
+test('visitPath restores the previous path when the visit throws', () => {
+    const stack = new NodeStack([root, root.program, owner]);
+    const visitor = pipe(
+        voidVisitor(),
+        v => recordNodeStackVisitor(v, stack),
+        v =>
+            tapVisitor(v, 'integerTypeNode', () => {
+                throw new Error('boom');
+            }),
+    );
+    expect(() => stack.visitPath([root, root.program, slot], visitor)).toThrow('boom');
+    expect(stack.getPath()).toStrictEqual([root, root.program, owner]);
+});

packages/visitors/src/unwrapDefinedTypesVisitor.ts

@@ -1,4 +1,4 @@
-import { assertIsNode, IdentifierString, Node, programNode, TYPE_NODES } from '@codama/nodes';
+import { assertIsNode, IdentifierString, programNode, TYPE_NODES } from '@codama/nodes';
 import {
     extendVisitor,
     findProgramNodeFromPath,
@@ -48,13 +48,7 @@ export function unwrapDefinedTypesVisitor(typesToInline: string[] | '*' = '*') {
                     const definedTypePath = linkables.getPathOrThrow(stack.getPath('definedTypeLinkNode'));
                     const definedType = getLastNodeFromPath(definedTypePath);
 
-                    stack.pushPath(definedTypePath);
-                    let type: Node;
-                    try {
-                        type = visit(definedType.type, self);
-                    } finally {
-                        stack.popPath();
-                    }
+                    const type = stack.withPath(definedTypePath, () => visit(definedType.type, self));
                     assertIsNode(type, TYPE_NODES);
 
                     return inlineDefinedType(link, type, {