From 655bc944388bf86d3d634c5f0a632db73d04f34d Mon Sep 17 00:00:00 2001 From: Loris Leiva Date: Mon, 28 Sep 2026 15:49:54 +0100 Subject: [PATCH] Add withPath and visitPath helpers to NodeStack --- packages/dynamic-codecs/src/codecs.ts | 7 +- packages/visitors-core/README.md | 11 ++++ packages/visitors-core/src/NodeStack.ts | 43 ++++++++++++- .../visitors-core/src/getByteSizeVisitor.ts | 7 +- .../src/getInstructionDataFields.ts | 7 +- .../src/getMaxByteSizeVisitor.ts | 7 +- packages/visitors-core/test/NodeStack.test.ts | 64 +++++++++++++++++++ .../visitors/src/unwrapDefinedTypesVisitor.ts | 10 +-- 8 files changed, 123 insertions(+), 33 deletions(-) create mode 100644 packages/visitors-core/test/NodeStack.test.ts diff --git a/packages/dynamic-codecs/src/codecs.ts b/packages/dynamic-codecs/src/codecs.ts index 5094cc435..b63bdb6c7 100644 --- a/packages/dynamic-codecs/src/codecs.ts +++ b/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); - stack.pushPath(path); - try { - return visit(getLastNodeFromPath(path), visitor); - } finally { - stack.popPath(); - } + return stack.visitPath(path, visitor); }; const getCollectionCodec = (item: Codec, count: CountNode): Codec => { diff --git a/packages/visitors-core/README.md b/packages/visitors-core/README.md index 1b8519e95..8848e70b2 100644 --- a/packages/visitors-core/README.md +++ b/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. diff --git a/packages/visitors-core/src/NodeStack.ts b/packages/visitors-core/src/NodeStack.ts index 531124121..af66ce60c 100644 --- a/packages/visitors-core/src/NodeStack.ts +++ b/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(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( + path: NodePath, + visitor: Visitor, + ): TReturn { + return this.withPath(path.slice(0, -1), () => visit(getLastNodeFromPath(path), visitor)); + } + public getPath(): NodePath; public getPath(kind: TKind | TKind[]): NodePath>; public getPath(kind?: TKind | TKind[]): NodePath { diff --git a/packages/visitors-core/src/getByteSizeVisitor.ts b/packages/visitors-core/src/getByteSizeVisitor.ts index 2afc28e5c..76fe788e6 100644 --- a/packages/visitors-core/src/getByteSizeVisitor.ts +++ b/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 }) { diff --git a/packages/visitors-core/src/getInstructionDataFields.ts b/packages/visitors-core/src/getInstructionDataFields.ts index 7dd7daaf5..c323b8e0c 100644 --- a/packages/visitors-core/src/getInstructionDataFields.ts +++ b/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; diff --git a/packages/visitors-core/src/getMaxByteSizeVisitor.ts b/packages/visitors-core/src/getMaxByteSizeVisitor.ts index 9365d5e87..2c1734b1c 100644 --- a/packages/visitors-core/src/getMaxByteSizeVisitor.ts +++ b/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 }) { diff --git a/packages/visitors-core/test/NodeStack.test.ts b/packages/visitors-core/test/NodeStack.test.ts new file mode 100644 index 000000000..25c74e408 --- /dev/null +++ b/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]); +}); diff --git a/packages/visitors/src/unwrapDefinedTypesVisitor.ts b/packages/visitors/src/unwrapDefinedTypesVisitor.ts index 800c7fd94..e800f79ad 100644 --- a/packages/visitors/src/unwrapDefinedTypesVisitor.ts +++ b/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, {