Remove LinkableDictionary's inner NodeStack (#266)

This PR removes the inner `NodeStack` inside the `LinkableDictionary` and requests stacks when resolving linkables in order to find the right program/instruction to fetch from.

This PR also tags places that need to be fixed and tested before the end of this refactoring stack.

.changeset/silver-foxes-hug.md

@@ -0,0 +1,5 @@
+---
+'@codama/visitors-core': minor
+---
+
+Remove `LinkableDictionary`'s inner `NodeStack`

packages/renderers-js-umi/src/getRenderMapVisitor.ts

@@ -25,8 +25,10 @@ import {
     getByteSizeVisitor,
     getResolvedInstructionInputsVisitor,
     LinkableDictionary,
+    NodeStack,
     pipe,
-    recordLinkablesVisitor,
+    recordLinkablesOnFirstVisitVisitor,
+    recordNodeStackVisitor,
     ResolvedInstructionAccount,
     ResolvedInstructionInput,
     staticVisitor,
@@ -60,7 +62,8 @@ export type GetRenderMapOptions = {
 
 export function getRenderMapVisitor(options: GetRenderMapOptions = {}): Visitor<RenderMap> {
     const linkables = new LinkableDictionary();
-    const byteSizeVisitor = getByteSizeVisitor(linkables);
+    const stack = new NodeStack();
+    const byteSizeVisitor = getByteSizeVisitor(linkables, stack);
     let program: ProgramNode | null = null;
 
     const renderParentInstructions = options.renderParentInstructions ?? false;
@@ -201,7 +204,7 @@ export function getRenderMapVisitor(options: GetRenderMapOptions = {}): Visitor<
                     }
 
                     // Seeds.
-                    const pda = node.pda ? linkables.get(node.pda) : undefined;
+                    const pda = node.pda ? linkables.get(node.pda, stack) : undefined;
                     const pdaSeeds = pda?.seeds ?? [];
                     const seeds = pdaSeeds.map(seed => {
                         if (isNode(seed, 'variablePdaSeedNode')) {
@@ -540,6 +543,7 @@ export function getRenderMapVisitor(options: GetRenderMapOptions = {}): Visitor<
                         .mergeWith(...getAllPrograms(node).map(p => visit(p, self)));
                 },
             }),
-        v => recordLinkablesVisitor(v, linkables),
+        v => recordNodeStackVisitor(v, stack),
+        v => recordLinkablesOnFirstVisitVisitor(v, linkables),
     );
 }

packages/renderers-js-umi/src/getTypeManifestVisitor.ts

@@ -17,7 +17,16 @@ import {
     structTypeNodeFromInstructionArgumentNodes,
     TypeNode,
 } from '@codama/nodes';
-import { extendVisitor, LinkableDictionary, pipe, staticVisitor, visit, Visitor } from '@codama/visitors-core';
+import {
+    extendVisitor,
+    LinkableDictionary,
+    NodeStack,
+    pipe,
+    recordNodeStackVisitor,
+    staticVisitor,
+    visit,
+    Visitor,
+} from '@codama/visitors-core';
 
 import { ImportMap } from './ImportMap';
 import { getBytesFromBytesValueNode, GetImportFromFunction, jsDocblock, ParsedCustomDataOptions } from './utils';
@@ -59,6 +68,7 @@ export function getTypeManifestVisitor(input: {
     const { linkables, nonScalarEnums, customAccountData, customInstructionData, getImportFrom } = input;
     let parentName = input.parentName ?? null;
     let parentSize: NumberTypeNode | number | null = null;
+    const stack = new NodeStack();
 
     return pipe(
         staticVisitor(
@@ -418,7 +428,8 @@ export function getTypeManifestVisitor(input: {
                     const variantName = pascalCase(node.variant);
                     const importFrom = getImportFrom(node.enum);
 
-                    const enumNode = linkables.get(node.enum)?.type;
+                    // FIXME(loris): No program node can ever be in this stack.
+                    const enumNode = linkables.get(node.enum, stack)?.type;
                     const isScalar =
                         enumNode && isNode(enumNode, 'enumTypeNode')
                             ? isScalarEnum(enumNode)
@@ -836,6 +847,7 @@ export function getTypeManifestVisitor(input: {
                     throw new CodamaError(CODAMA_ERROR__RENDERERS__UNSUPPORTED_NODE, { kind: node.kind, node });
                 },
             }),
+        v => recordNodeStackVisitor(v, stack),
     );
 }
 

packages/renderers-js/src/fragments/accountPdaHelpers.ts

@@ -1,4 +1,5 @@
-import { AccountNode, isNodeFilter, ProgramNode } from '@codama/nodes';
+import { AccountNode, isNodeFilter } from '@codama/nodes';
+import { NodeStack } from '@codama/visitors-core';
 
 import type { GlobalFragmentScope } from '../getRenderMapVisitor';
 import type { TypeManifest } from '../TypeManifest';
@@ -7,12 +8,12 @@ import { Fragment, fragment, fragmentFromTemplate } from './common';
 export function getAccountPdaHelpersFragment(
     scope: Pick<GlobalFragmentScope, 'customAccountData' | 'linkables' | 'nameApi'> & {
         accountNode: AccountNode;
-        programNode: ProgramNode;
+        accountStack: NodeStack;
         typeManifest: TypeManifest;
     },
 ): Fragment {
-    const { accountNode, programNode, nameApi, linkables, customAccountData, typeManifest } = scope;
-    const pdaNode = accountNode.pda ? linkables.get(accountNode.pda) : undefined;
+    const { accountNode, accountStack, nameApi, linkables, customAccountData, typeManifest } = scope;
+    const pdaNode = accountNode.pda ? linkables.get(accountNode.pda, accountStack) : undefined;
     if (!pdaNode) {
         return fragment('');
     }
@@ -38,7 +39,7 @@ export function getAccountPdaHelpersFragment(
         findPdaFunction,
         hasVariableSeeds,
         pdaSeedsType,
-        program: programNode,
+        program: accountStack.getProgram(),
     })
         .mergeImportsWith(accountTypeFragment)
         .addImports(importFrom, hasVariableSeeds ? [pdaSeedsType, findPdaFunction] : [findPdaFunction])

packages/renderers-js/src/fragments/instructionAccountTypeParam.ts

@@ -1,11 +1,5 @@
-import {
-    InstructionAccountNode,
-    InstructionInputValueNode,
-    InstructionNode,
-    pascalCase,
-    ProgramNode,
-} from '@codama/nodes';
-import { LinkableDictionary } from '@codama/visitors-core';
+import { InstructionAccountNode, InstructionInputValueNode, InstructionNode, pascalCase } from '@codama/nodes';
+import { LinkableDictionary, NodeStack } from '@codama/visitors-core';
 
 import type { GlobalFragmentScope } from '../getRenderMapVisitor';
 import { ImportMap } from '../ImportMap';
@@ -16,10 +10,10 @@ export function getInstructionAccountTypeParamFragment(
         allowAccountMeta: boolean;
         instructionAccountNode: InstructionAccountNode;
         instructionNode: InstructionNode;
-        programNode: ProgramNode;
+        instructionStack: NodeStack;
     },
 ): Fragment {
-    const { instructionNode, instructionAccountNode, programNode, allowAccountMeta, linkables } = scope;
+    const { instructionNode, instructionAccountNode, instructionStack, allowAccountMeta, linkables } = scope;
     const typeParam = `TAccount${pascalCase(instructionAccountNode.name)}`;
     const accountMeta = allowAccountMeta ? ' | IAccountMeta<string>' : '';
     const imports = new ImportMap();
@@ -31,7 +25,11 @@ export function getInstructionAccountTypeParamFragment(
         return fragment(`${typeParam} extends string${accountMeta} | undefined = undefined`, imports);
     }
 
-    const defaultAddress = getDefaultAddress(instructionAccountNode.defaultValue, programNode.publicKey, linkables);
+    const defaultAddress = getDefaultAddress(
+        instructionAccountNode.defaultValue,
+        instructionStack.getProgram()!.publicKey,
+        linkables,
+    );
 
     return fragment(`${typeParam} extends string${accountMeta} = ${defaultAddress}`, imports);
 }
@@ -45,8 +43,9 @@ function getDefaultAddress(
         case 'publicKeyValueNode':
             return `"${defaultValue.publicKey}"`;
         case 'programLinkNode':
+            // FIXME(loris): No need for a stack here.
             // eslint-disable-next-line no-case-declarations
-            const programNode = linkables.get(defaultValue);
+            const programNode = linkables.get(defaultValue, new NodeStack());
             return programNode ? `"${programNode.publicKey}"` : 'string';
         case 'programIdValueNode':
             return `"${programId}"`;

packages/renderers-js/src/fragments/instructionFunction.ts

@@ -1,13 +1,5 @@
-import {
-    camelCase,
-    InstructionArgumentNode,
-    InstructionNode,
-    isNode,
-    isNodeFilter,
-    pascalCase,
-    ProgramNode,
-} from '@codama/nodes';
-import { ResolvedInstructionInput } from '@codama/visitors-core';
+import { camelCase, InstructionArgumentNode, InstructionNode, isNode, isNodeFilter, pascalCase } from '@codama/nodes';
+import { NodeStack, ResolvedInstructionInput } from '@codama/visitors-core';
 
 import type { GlobalFragmentScope } from '../getRenderMapVisitor';
 import { NameApi } from '../nameTransformers';
@@ -27,7 +19,7 @@ export function getInstructionFunctionFragment(
         dataArgsManifest: TypeManifest;
         extraArgsManifest: TypeManifest;
         instructionNode: InstructionNode;
-        programNode: ProgramNode;
+        instructionStack: NodeStack;
         renamedArgs: Map<string, string>;
         resolvedInputs: ResolvedInstructionInput[];
         useAsync: boolean;
@@ -36,7 +28,7 @@ export function getInstructionFunctionFragment(
     const {
         useAsync,
         instructionNode,
-        programNode,
+        instructionStack,
         resolvedInputs,
         renamedArgs,
         dataArgsManifest,
@@ -74,7 +66,7 @@ export function getInstructionFunctionFragment(
     const hasAnyArgs = hasDataArgs || hasExtraArgs || hasRemainingAccountArgs;
     const hasInput = hasAccounts || hasAnyArgs;
     const instructionDataName = nameApi.instructionDataType(instructionNode.name);
-    const programAddressConstant = nameApi.programAddressConstant(programNode.name);
+    const programAddressConstant = nameApi.programAddressConstant(instructionStack.getProgram()!.name);
     const encoderFunction = customData
         ? dataArgsManifest.encoder.render
         : `${nameApi.encoderFunction(instructionDataName)}()`;

packages/renderers-js/src/fragments/instructionParseFunction.ts

@@ -1,4 +1,5 @@
-import { InstructionNode, ProgramNode } from '@codama/nodes';
+import { InstructionNode } from '@codama/nodes';
+import { NodeStack } from '@codama/visitors-core';
 
 import type { GlobalFragmentScope } from '../getRenderMapVisitor';
 import { TypeManifest } from '../TypeManifest';
@@ -8,10 +9,10 @@ export function getInstructionParseFunctionFragment(
     scope: Pick<GlobalFragmentScope, 'customInstructionData' | 'nameApi'> & {
         dataArgsManifest: TypeManifest;
         instructionNode: InstructionNode;
-        programNode: ProgramNode;
+        instructionStack: NodeStack;
     },
 ): Fragment {
-    const { instructionNode, programNode, dataArgsManifest, nameApi, customInstructionData } = scope;
+    const { instructionNode, instructionStack, dataArgsManifest, nameApi, customInstructionData } = scope;
     const customData = customInstructionData.get(instructionNode.name);
     const hasAccounts = instructionNode.accounts.length > 0;
     const hasOptionalAccounts = instructionNode.accounts.some(account => account.isOptional);
@@ -22,7 +23,7 @@ export function getInstructionParseFunctionFragment(
     const hasData = !!customData || instructionNode.arguments.length > 0;
 
     const instructionDataName = nameApi.instructionDataType(instructionNode.name);
-    const programAddressConstant = nameApi.programAddressConstant(programNode.name);
+    const programAddressConstant = nameApi.programAddressConstant(instructionStack.getProgram()!.name);
     const dataTypeFragment = fragment(
         customData ? dataArgsManifest.strictType.render : nameApi.dataType(instructionDataName),
     );

packages/renderers-js/src/fragments/instructionType.ts

@@ -1,4 +1,5 @@
-import { InstructionNode, pascalCase, ProgramNode } from '@codama/nodes';
+import { InstructionNode, pascalCase } from '@codama/nodes';
+import { NodeStack } from '@codama/visitors-core';
 
 import type { GlobalFragmentScope } from '../getRenderMapVisitor';
 import { Fragment, fragmentFromTemplate, mergeFragments } from './common';
@@ -8,10 +9,11 @@ import { getInstructionAccountTypeParamFragment } from './instructionAccountType
 export function getInstructionTypeFragment(
     scope: Pick<GlobalFragmentScope, 'customInstructionData' | 'linkables' | 'nameApi'> & {
         instructionNode: InstructionNode;
-        programNode: ProgramNode;
+        instructionStack: NodeStack;
     },
 ): Fragment {
-    const { instructionNode, programNode, nameApi, customInstructionData } = scope;
+    const { instructionNode, instructionStack, nameApi, customInstructionData } = scope;
+    const programNode = instructionStack.getProgram()!;
     const hasAccounts = instructionNode.accounts.length > 0;
     const customData = customInstructionData.get(instructionNode.name);
     const hasData = !!customData || instructionNode.arguments.length > 0;

packages/renderers-js/src/fragments/pdaFunction.ts

@@ -1,5 +1,5 @@
-import { isNode, isNodeFilter, PdaNode, ProgramNode } from '@codama/nodes';
-import { visit } from '@codama/visitors-core';
+import { isNode, isNodeFilter, PdaNode } from '@codama/nodes';
+import { NodeStack, visit } from '@codama/visitors-core';
 
 import type { GlobalFragmentScope } from '../getRenderMapVisitor';
 import { ImportMap } from '../ImportMap';
@@ -8,10 +8,10 @@ import { Fragment, fragmentFromTemplate } from './common';
 export function getPdaFunctionFragment(
     scope: Pick<GlobalFragmentScope, 'nameApi' | 'typeManifestVisitor'> & {
         pdaNode: PdaNode;
-        programNode: ProgramNode;
+        pdaStack: NodeStack;
     },
 ): Fragment {
-    const { pdaNode, programNode, typeManifestVisitor, nameApi } = scope;
+    const { pdaNode, pdaStack, typeManifestVisitor, nameApi } = scope;
 
     // Seeds.
     const imports = new ImportMap();
@@ -37,7 +37,7 @@ export function getPdaFunctionFragment(
         findPdaFunction: nameApi.pdaFindFunction(pdaNode.name),
         hasVariableSeeds,
         pdaSeedsType: nameApi.pdaSeedsType(pdaNode.name),
-        programAddress: pdaNode.programId ?? programNode.publicKey,
+        programAddress: pdaNode.programId ?? pdaStack.getProgram()!.publicKey,
         seeds,
     })
         .mergeImportsWith(imports)

packages/renderers-js/src/getRenderMapVisitor.ts

@@ -10,7 +10,6 @@ import {
     getAllPdas,
     getAllPrograms,
     InstructionNode,
-    ProgramNode,
     resolveNestedTypeNode,
     structTypeNodeFromInstructionArgumentNodes,
 } from '@codama/nodes';
@@ -19,8 +18,10 @@ import {
     extendVisitor,
     getResolvedInstructionInputsVisitor,
     LinkableDictionary,
+    NodeStack,
     pipe,
-    recordLinkablesVisitor,
+    recordLinkablesOnFirstVisitVisitor,
+    recordNodeStackVisitor,
     staticVisitor,
     visit,
 } from '@codama/visitors-core';
@@ -86,7 +87,7 @@ export type GlobalFragmentScope = {
 
 export function getRenderMapVisitor(options: GetRenderMapOptions = {}) {
     const linkables = new LinkableDictionary();
-    let program: ProgramNode | null = null;
+    const stack = new NodeStack();
 
     const nameTransformers = {
         ...DEFAULT_NAME_TRANSFORMERS,
@@ -140,14 +141,14 @@ export function getRenderMapVisitor(options: GetRenderMapOptions = {}) {
         v =>
             extendVisitor(v, {
                 visitAccount(node) {
-                    if (!program) {
+                    if (!stack.getProgram()) {
                         throw new Error('Account must be visited inside a program.');
                     }
 
                     const scope = {
                         ...globalScope,
                         accountNode: node,
-                        programNode: program,
+                        accountStack: stack.clone(),
                         typeManifest: visit(node, typeManifestVisitor),
                     };
 
@@ -221,7 +222,7 @@ export function getRenderMapVisitor(options: GetRenderMapOptions = {}) {
                 },
 
                 visitInstruction(node) {
-                    if (!program) {
+                    if (!stack.getProgram()) {
                         throw new Error('Instruction must be visited inside a program.');
                     }
 
@@ -237,7 +238,7 @@ export function getRenderMapVisitor(options: GetRenderMapOptions = {}) {
                             }),
                         ),
                         instructionNode: node,
-                        programNode: program,
+                        instructionStack: stack,
                         renamedArgs: getRenamedArgsMap(node),
                         resolvedInputs: visit(node, resolvedInstructionInputVisitor),
                     };
@@ -290,11 +291,11 @@ export function getRenderMapVisitor(options: GetRenderMapOptions = {}) {
                 },
 
                 visitPda(node) {
-                    if (!program) {
+                    if (!stack.getProgram()) {
                         throw new Error('Account must be visited inside a program.');
                     }
 
-                    const scope = { ...globalScope, pdaNode: node, programNode: program };
+                    const scope = { ...globalScope, pdaNode: node, pdaStack: stack };
                     const pdaFunctionFragment = getPdaFunctionFragment(scope);
                     const imports = new ImportMap().mergeWith(pdaFunctionFragment);
 
@@ -308,7 +309,6 @@ export function getRenderMapVisitor(options: GetRenderMapOptions = {}) {
                 },
 
                 visitProgram(node, { self }) {
-                    program = node;
                     const customDataDefinedType = [
                         ...getDefinedTypeNodesToExtract(node.accounts, customAccountData),
                         ...getDefinedTypeNodesToExtract(node.instructions, customInstructionData),
@@ -349,11 +349,10 @@ export function getRenderMapVisitor(options: GetRenderMapOptions = {}) {
                     );
 
                     renderMap.mergeWith(
-                        ...getAllInstructionsWithSubs(program, {
+                        ...getAllInstructionsWithSubs(node, {
                             leavesOnly: !renderParentInstructions,
                         }).map(ix => visit(ix, self)),
                     );
-                    program = null;
                     return renderMap;
                 },
 
@@ -435,7 +434,8 @@ export function getRenderMapVisitor(options: GetRenderMapOptions = {}) {
                         .mergeWith(...getAllPrograms(node).map(p => visit(p, self)));
                 },
             }),
-        v => recordLinkablesVisitor(v, linkables),
+        v => recordNodeStackVisitor(v, stack),
+        v => recordLinkablesOnFirstVisitVisitor(v, linkables),
     );
 }
 

packages/renderers-js/src/getTypeManifestVisitor.ts

@@ -354,7 +354,8 @@ export function getTypeManifestVisitor(input: {
                     const enumFunction = nameApi.discriminatedUnionFunction(node.enum.name);
                     const importFrom = getImportFrom(node.enum);
 
-                    const enumNode = linkables.get(node.enum)?.type;
+                    // FIXME(loris): No program node can ever be in this stack.
+                    const enumNode = linkables.get(node.enum, stack)?.type;
                     const isScalar =
                         enumNode && isNode(enumNode, 'enumTypeNode')
                             ? isScalarEnum(enumNode)

packages/renderers-rust/src/getRenderMapVisitor.ts

@@ -18,8 +18,10 @@ import { RenderMap } from '@codama/renderers-core';
 import {
     extendVisitor,
     LinkableDictionary,
+    NodeStack,
     pipe,
-    recordLinkablesVisitor,
+    recordLinkablesOnFirstVisitVisitor,
+    recordNodeStackVisitor,
     staticVisitor,
     visit,
 } from '@codama/visitors-core';
@@ -40,6 +42,7 @@ export type GetRenderMapOptions = {
 
 export function getRenderMapVisitor(options: GetRenderMapOptions = {}) {
     const linkables = new LinkableDictionary();
+    const stack = new NodeStack();
     let program: ProgramNode | null = null;
 
     const renderParentInstructions = options.renderParentInstructions ?? false;
@@ -61,7 +64,7 @@ export function getRenderMapVisitor(options: GetRenderMapOptions = {}) {
 
                     // Seeds.
                     const seedsImports = new ImportMap();
-                    const pda = node.pda ? linkables.get(node.pda) : undefined;
+                    const pda = node.pda ? linkables.get(node.pda, stack) : undefined;
                     const pdaSeeds = pda?.seeds ?? [];
                     const seeds = pdaSeeds.map(seed => {
                         if (isNode(seed, 'variablePdaSeedNode')) {
@@ -286,7 +289,8 @@ export function getRenderMapVisitor(options: GetRenderMapOptions = {}) {
                         .mergeWith(...getAllPrograms(node).map(p => visit(p, self)));
                 },
             }),
-        v => recordLinkablesVisitor(v, linkables),
+        v => recordNodeStackVisitor(v, stack),
+        v => recordLinkablesOnFirstVisitVisitor(v, linkables),
     );
 }
 

packages/validators/src/getValidationItemsVisitor.ts

@@ -6,7 +6,7 @@ import {
     mergeVisitor,
     NodeStack,
     pipe,
-    recordLinkablesVisitor,
+    recordLinkablesOnFirstVisitVisitor,
     recordNodeStackVisitor,
     visit,
     Visitor,
@@ -23,7 +23,7 @@ export function getValidationItemsVisitor(): Visitor<readonly ValidationItem[]>
             () => [] as readonly ValidationItem[],
             (_, items) => items.flat(),
         ),
-        v => recordLinkablesVisitor(v, linkables),
+        v => recordLinkablesOnFirstVisitVisitor(v, linkables),
         v => recordNodeStackVisitor(v, stack),
         v =>
             extendVisitor(v, {
@@ -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(node)) {
+                    } else if (!linkables.has(node, stack)) {
                         items.push(
                             validationItem(
                                 'error',

packages/visitors-core/README.md

@@ -654,23 +654,28 @@ It offers the following API:
 const linkables = new LinkableDictionary();
 
 // Record program nodes.
-linkables.record(programNode);
+linkables.record(programNode, stack);
 
 // Record other linkable nodes with their associated program node.
-linkables.record(accountNode);
+linkables.record(accountNode, stack);
 
 // Get a linkable node using a link node, or throw an error if it is not found.
-const programNode = linkables.getOrThrow(programLinkNode);
+const programNode = linkables.getOrThrow(programLinkNode, stack);
 
 // Get a linkable node using a link node, or return undefined if it is not found.
-const accountNode = linkables.get(accountLinkNode);
+const accountNode = linkables.get(accountLinkNode, stack);
 ```
 
-Note that this API must be used in conjunction with the `recordLinkablesVisitor` to record the linkable nodes and, later on, resolve the link nodes as we traverse the nodes. This is because the `LinkableDictionary` instance keeps track of its own internal `NodeStack` in order to understand which program node should be used for a given link node.
+Note that:
 
-### `recordLinkablesVisitor`
+-   The stack of the recorded node must be provided when recording a linkable node.
+-   The stack of the link node must be provided when getting a linkable node from it.
 
-Much like the `recordNodeStackVisitor`, the `recordLinkablesVisitor` allows us to record linkable nodes as we traverse the tree of nodes. It accepts a base visitor and `LinkableDictionary` instance; and records any linkable node it encounters.
+This API may be used with the `recordLinkablesOnFirstVisitVisitor` to record the linkable nodes before the first node visit; as well as the `recordNodeStackVisitor` to keep track of the current node stack when accessing the linkable nodes.
+
+### `recordLinkablesOnFirstVisitVisitor`
+
+Much like the `recordNodeStackVisitor`, the `recordLinkablesOnFirstVisitVisitor` allows us to record linkable nodes as we traverse the tree of nodes. It accepts a base visitor and `LinkableDictionary` instance; and records any linkable node it encounters.
 
 This means that we can inject the `LinkableDictionary` instance into another extension of the base visitor to resolve any link node we encounter.
 
@@ -678,6 +683,7 @@ Here's an example that records a `LinkableDictionary` and uses it to log the amo
 
 ```ts
 const linkables = new LinkableDictionary();
+const stack = new NodeStack();
 const visitor = pipe(
     baseVisitor,
     v =>
@@ -685,11 +691,12 @@ const visitor = pipe(
             const pdaNode = linkables.getOrThrow(node);
             console.log(`${pdaNode.seeds.length} seeds`);
         }),
-    v => recordLinkablesVisitor(v, linkables),
+    v => recordNodeStackVisitor(v, stack),
+    v => recordLinkablesOnFirstVisitVisitor(v, linkables),
 );
 ```
 
-Note that the `recordLinkablesVisitor` should always be the last visitor in the pipe to ensure that all linkable nodes are recorded before being used.
+Note that the `recordLinkablesOnFirstVisitVisitor` should be the last visitor in the pipe to ensure that all linkable nodes are recorded before being used.
 
 ## Other useful visitors
 

packages/visitors-core/src/LinkableDictionary.ts

@@ -57,12 +57,10 @@ type InstructionDictionary = {
 export class LinkableDictionary {
     readonly programs: Map<string, ProgramDictionary> = new Map();
 
-    readonly stack: NodeStack = new NodeStack();
-
-    record(node: LinkableNode): this {
-        const programDictionary = this.getOrCreateProgramDictionary(node);
+    record(node: LinkableNode, stack: NodeStack): this {
+        const programDictionary = this.getOrCreateProgramDictionary(node, stack);
         if (!programDictionary) return this; // Do not record nodes that are outside of a program.
-        const instructionDictionary = this.getOrCreateInstructionDictionary(programDictionary, node);
+        const instructionDictionary = this.getOrCreateInstructionDictionary(programDictionary, node, stack);
 
         if (isNode(node, 'accountNode')) {
             programDictionary.accounts.set(node.name, node);
@@ -79,39 +77,39 @@ export class LinkableDictionary {
         return this;
     }
 
-    getOrThrow(linkNode: AccountLinkNode): AccountNode;
-    getOrThrow(linkNode: DefinedTypeLinkNode): DefinedTypeNode;
-    getOrThrow(linkNode: InstructionAccountLinkNode): InstructionAccountNode;
-    getOrThrow(linkNode: InstructionArgumentLinkNode): InstructionArgumentNode;
-    getOrThrow(linkNode: InstructionLinkNode): InstructionNode;
-    getOrThrow(linkNode: PdaLinkNode): PdaNode;
-    getOrThrow(linkNode: ProgramLinkNode): ProgramNode;
-    getOrThrow(linkNode: LinkNode): LinkableNode {
-        const node = this.get(linkNode as ProgramLinkNode) as LinkableNode | undefined;
+    getOrThrow(linkNode: AccountLinkNode, stack: NodeStack): AccountNode;
+    getOrThrow(linkNode: DefinedTypeLinkNode, stack: NodeStack): DefinedTypeNode;
+    getOrThrow(linkNode: InstructionAccountLinkNode, stack: NodeStack): InstructionAccountNode;
+    getOrThrow(linkNode: InstructionArgumentLinkNode, stack: NodeStack): InstructionArgumentNode;
+    getOrThrow(linkNode: InstructionLinkNode, stack: NodeStack): InstructionNode;
+    getOrThrow(linkNode: PdaLinkNode, stack: NodeStack): PdaNode;
+    getOrThrow(linkNode: ProgramLinkNode, stack: NodeStack): ProgramNode;
+    getOrThrow(linkNode: LinkNode, stack: NodeStack): LinkableNode {
+        const node = this.get(linkNode as ProgramLinkNode, stack) as LinkableNode | undefined;
 
         if (!node) {
             throw new CodamaError(CODAMA_ERROR__LINKED_NODE_NOT_FOUND, {
                 kind: linkNode.kind,
                 linkNode,
                 name: linkNode.name,
-                stack: this.stack.all(),
+                stack: stack.all(),
             });
         }
 
         return node;
     }
 
-    get(linkNode: AccountLinkNode): AccountNode | undefined;
-    get(linkNode: DefinedTypeLinkNode): DefinedTypeNode | undefined;
-    get(linkNode: InstructionAccountLinkNode): InstructionAccountNode | undefined;
-    get(linkNode: InstructionArgumentLinkNode): InstructionArgumentNode | undefined;
-    get(linkNode: InstructionLinkNode): InstructionNode | undefined;
-    get(linkNode: PdaLinkNode): PdaNode | undefined;
-    get(linkNode: ProgramLinkNode): ProgramNode | undefined;
-    get(linkNode: LinkNode): LinkableNode | undefined {
-        const programDictionary = this.getProgramDictionary(linkNode);
+    get(linkNode: AccountLinkNode, stack: NodeStack): AccountNode | undefined;
+    get(linkNode: DefinedTypeLinkNode, stack: NodeStack): DefinedTypeNode | undefined;
+    get(linkNode: InstructionAccountLinkNode, stack: NodeStack): InstructionAccountNode | undefined;
+    get(linkNode: InstructionArgumentLinkNode, stack: NodeStack): InstructionArgumentNode | undefined;
+    get(linkNode: InstructionLinkNode, stack: NodeStack): InstructionNode | undefined;
+    get(linkNode: PdaLinkNode, stack: NodeStack): PdaNode | undefined;
+    get(linkNode: ProgramLinkNode, stack: NodeStack): ProgramNode | undefined;
+    get(linkNode: LinkNode, stack: NodeStack): LinkableNode | undefined {
+        const programDictionary = this.getProgramDictionary(linkNode, stack);
         if (!programDictionary) return undefined;
-        const instructionDictionary = this.getInstructionDictionary(programDictionary, linkNode);
+        const instructionDictionary = this.getInstructionDictionary(programDictionary, linkNode, stack);
 
         if (isNode(linkNode, 'accountLinkNode')) {
             return programDictionary.accounts.get(linkNode.name);
@@ -132,10 +130,10 @@ export class LinkableDictionary {
         return undefined;
     }
 
-    has(linkNode: LinkNode): boolean {
-        const programDictionary = this.getProgramDictionary(linkNode);
+    has(linkNode: LinkNode, stack: NodeStack): boolean {
+        const programDictionary = this.getProgramDictionary(linkNode, stack);
         if (!programDictionary) return false;
-        const instructionDictionary = this.getInstructionDictionary(programDictionary, linkNode);
+        const instructionDictionary = this.getInstructionDictionary(programDictionary, linkNode, stack);
 
         if (isNode(linkNode, 'accountLinkNode')) {
             return programDictionary.accounts.has(linkNode.name);
@@ -156,8 +154,8 @@ export class LinkableDictionary {
         return false;
     }
 
-    private getOrCreateProgramDictionary(node: LinkableNode): ProgramDictionary | undefined {
-        const programNode = isNode(node, 'programNode') ? node : this.stack.getProgram();
+    private getOrCreateProgramDictionary(node: LinkableNode, stack: NodeStack): ProgramDictionary | undefined {
+        const programNode = isNode(node, 'programNode') ? node : stack.getProgram();
         if (!programNode) return undefined;
 
         let programDictionary = this.programs.get(programNode.name);
@@ -178,8 +176,9 @@ export class LinkableDictionary {
     private getOrCreateInstructionDictionary(
         programDictionary: ProgramDictionary,
         node: LinkableNode,
+        stack: NodeStack,
     ): InstructionDictionary | undefined {
-        const instructionNode = isNode(node, 'instructionNode') ? node : this.stack.getInstruction();
+        const instructionNode = isNode(node, 'instructionNode') ? node : stack.getInstruction();
         if (!instructionNode) return undefined;
 
         let instructionDictionary = programDictionary.instructions.get(instructionNode.name);
@@ -195,7 +194,7 @@ export class LinkableDictionary {
         return instructionDictionary;
     }
 
-    private getProgramDictionary(linkNode: LinkNode): ProgramDictionary | undefined {
+    private getProgramDictionary(linkNode: LinkNode, stack: NodeStack): ProgramDictionary | undefined {
         let programName: CamelCaseString | undefined = undefined;
         if (isNode(linkNode, 'programLinkNode')) {
             programName = linkNode.name;
@@ -204,7 +203,7 @@ export class LinkableDictionary {
         } else if ('instruction' in linkNode) {
             programName = linkNode.instruction?.program?.name;
         }
-        programName = programName ?? this.stack.getProgram()?.name;
+        programName = programName ?? stack.getProgram()?.name;
 
         return programName ? this.programs.get(programName) : undefined;
     }
@@ -212,6 +211,7 @@ export class LinkableDictionary {
     private getInstructionDictionary(
         programDictionary: ProgramDictionary,
         linkNode: LinkNode,
+        stack: NodeStack,
     ): InstructionDictionary | undefined {
         let instructionName: CamelCaseString | undefined = undefined;
         if (isNode(linkNode, 'instructionLinkNode')) {
@@ -219,7 +219,7 @@ export class LinkableDictionary {
         } else if ('instruction' in linkNode) {
             instructionName = linkNode.instruction?.name;
         }
-        instructionName = instructionName ?? this.stack.getInstruction()?.name;
+        instructionName = instructionName ?? stack.getInstruction()?.name;
 
         return instructionName ? programDictionary.instructions.get(instructionName) : undefined;
     }

packages/visitors-core/src/getByteSizeVisitor.ts

@@ -1,7 +1,11 @@
 import { isNode, isScalarEnum, REGISTERED_TYPE_NODE_KINDS, RegisteredTypeNode } from '@codama/nodes';
 
+import { extendVisitor } from './extendVisitor';
 import { LinkableDictionary } from './LinkableDictionary';
 import { mergeVisitor } from './mergeVisitor';
+import { NodeStack } from './NodeStack';
+import { pipe } from './pipe';
+import { recordNodeStackVisitor } from './recordNodeStackVisitor';
 import { visit, Visitor } from './visitor';
 
 export type ByteSizeVisitorKeys =
@@ -12,14 +16,17 @@ export type ByteSizeVisitorKeys =
     | 'instructionArgumentNode'
     | 'instructionNode';
 
-export function getByteSizeVisitor(linkables: LinkableDictionary): Visitor<number | null, ByteSizeVisitorKeys> {
+export function getByteSizeVisitor(
+    linkables: LinkableDictionary,
+    stack: NodeStack,
+): Visitor<number | null, ByteSizeVisitorKeys> {
     const visitedDefinedTypes = new Map<string, number | null>();
     const definedTypeStack: string[] = [];
 
     const sumSizes = (values: (number | null)[]): number | null =>
         values.reduce((all, one) => (all === null || one === null ? null : all + one), 0 as number | null);
 
-    const visitor = mergeVisitor(
+    const baseVisitor = mergeVisitor(
         () => null as number | null,
         (_, values) => sumSizes(values),
         [
@@ -32,88 +39,92 @@ export function getByteSizeVisitor(linkables: LinkableDictionary): Visitor<numbe
         ],
     );
 
-    return {
-        ...visitor,
-
-        visitAccount(node) {
-            return visit(node.data, this);
-        },
-
-        visitArrayType(node) {
-            if (!isNode(node.count, 'fixedCountNode')) return null;
-            const fixedSize = node.count.value;
-            const itemSize = visit(node.item, this);
-            const arraySize = itemSize !== null ? itemSize * fixedSize : null;
-            return fixedSize === 0 ? 0 : arraySize;
-        },
-
-        visitDefinedType(node) {
-            if (visitedDefinedTypes.has(node.name)) {
-                return visitedDefinedTypes.get(node.name)!;
-            }
-            definedTypeStack.push(node.name);
-            const child = visit(node.type, this);
-            definedTypeStack.pop();
-            visitedDefinedTypes.set(node.name, child);
-            return child;
-        },
-
-        visitDefinedTypeLink(node) {
-            // Fetch the linked type and return null if not found.
-            // The validator visitor will throw a proper error later on.
-            const linkedDefinedType = linkables.get(node);
-            if (!linkedDefinedType) {
-                return null;
-            }
-
-            // This prevents infinite recursion by using assuming
-            // cyclic types don't have a fixed size.
-            if (definedTypeStack.includes(linkedDefinedType.name)) {
-                return null;
-            }
-
-            return visit(linkedDefinedType, this);
-        },
-
-        visitEnumEmptyVariantType() {
-            return 0;
-        },
-
-        visitEnumType(node) {
-            const prefix = visit(node.size, this) ?? 1;
-            if (isScalarEnum(node)) return prefix;
-            const variantSizes = node.variants.map(v => visit(v, this));
-            const allVariantHaveTheSameFixedSize = variantSizes.every((one, _, all) => one === all[0]);
-            return allVariantHaveTheSameFixedSize && variantSizes.length > 0 && variantSizes[0] !== null
-                ? variantSizes[0] + prefix
-                : null;
-        },
-
-        visitFixedSizeType(node) {
-            return node.size;
-        },
-
-        visitInstruction(node) {
-            return sumSizes(node.arguments.map(arg => visit(arg, this)));
-        },
-
-        visitInstructionArgument(node) {
-            return visit(node.type, this);
-        },
-
-        visitNumberType(node) {
-            return parseInt(node.format.slice(1), 10) / 8;
-        },
-
-        visitOptionType(node) {
-            if (!node.fixed) return null;
-            const prefixSize = visit(node.prefix, this) as number;
-            const itemSize = visit(node.item, this);
-            return itemSize !== null ? itemSize + prefixSize : null;
-        },
-
-        visitPublicKeyType() {
-            return 32;
-        },
-    };
+    return pipe(
+        baseVisitor,
+        v =>
+            extendVisitor(v, {
+                visitAccount(node, { self }) {
+                    return visit(node.data, self);
+                },
+
+                visitArrayType(node, { self }) {
+                    if (!isNode(node.count, 'fixedCountNode')) return null;
+                    const fixedSize = node.count.value;
+                    const itemSize = visit(node.item, self);
+                    const arraySize = itemSize !== null ? itemSize * fixedSize : null;
+                    return fixedSize === 0 ? 0 : arraySize;
+                },
+
+                visitDefinedType(node, { self }) {
+                    if (visitedDefinedTypes.has(node.name)) {
+                        return visitedDefinedTypes.get(node.name)!;
+                    }
+                    definedTypeStack.push(node.name);
+                    const child = visit(node.type, self);
+                    definedTypeStack.pop();
+                    visitedDefinedTypes.set(node.name, child);
+                    return child;
+                },
+
+                visitDefinedTypeLink(node, { self }) {
+                    // Fetch the linked type and return null if not found.
+                    // The validator visitor will throw a proper error later on.
+                    const linkedDefinedType = linkables.get(node, stack);
+                    if (!linkedDefinedType) {
+                        return null;
+                    }
+
+                    // This prevents infinite recursion by using assuming
+                    // cyclic types don't have a fixed size.
+                    if (definedTypeStack.includes(linkedDefinedType.name)) {
+                        return null;
+                    }
+
+                    return visit(linkedDefinedType, self);
+                },
+
+                visitEnumEmptyVariantType() {
+                    return 0;
+                },
+
+                visitEnumType(node, { self }) {
+                    const prefix = visit(node.size, self) ?? 1;
+                    if (isScalarEnum(node)) return prefix;
+                    const variantSizes = node.variants.map(v => visit(v, self));
+                    const allVariantHaveTheSameFixedSize = variantSizes.every((one, _, all) => one === all[0]);
+                    return allVariantHaveTheSameFixedSize && variantSizes.length > 0 && variantSizes[0] !== null
+                        ? variantSizes[0] + prefix
+                        : null;
+                },
+
+                visitFixedSizeType(node) {
+                    return node.size;
+                },
+
+                visitInstruction(node, { self }) {
+                    return sumSizes(node.arguments.map(arg => visit(arg, self)));
+                },
+
+                visitInstructionArgument(node, { self }) {
+                    return visit(node.type, self);
+                },
+
+                visitNumberType(node) {
+                    if (node.format === 'shortU16') return null;
+                    return parseInt(node.format.slice(1), 10) / 8;
+                },
+
+                visitOptionType(node, { self }) {
+                    if (!node.fixed) return null;
+                    const prefixSize = visit(node.prefix, self) as number;
+                    const itemSize = visit(node.item, self);
+                    return itemSize !== null ? itemSize + prefixSize : null;
+                },
+
+                visitPublicKeyType() {
+                    return 32;
+                },
+            }),
+        v => recordNodeStackVisitor(v, stack),
+    );
 }

packages/visitors-core/src/index.ts

@@ -1,6 +1,3 @@
-export * from './LinkableDictionary';
-export * from './NodeSelector';
-export * from './NodeStack';
 export * from './bottomUpTransformerVisitor';
 export * from './consoleLogVisitor';
 export * from './deleteNodesVisitor';
@@ -10,10 +7,13 @@ export * from './getDebugStringVisitor';
 export * from './getResolvedInstructionInputsVisitor';
 export * from './getUniqueHashStringVisitor';
 export * from './identityVisitor';
-export * from './interceptVisitor';
 export * from './interceptFirstVisitVisitor';
+export * from './interceptVisitor';
+export * from './LinkableDictionary';
 export * from './mapVisitor';
 export * from './mergeVisitor';
+export * from './NodeSelector';
+export * from './NodeStack';
 export * from './nonNullableIdentityVisitor';
 export * from './pipe';
 export * from './recordLinkablesVisitor';

packages/visitors-core/src/recordLinkablesVisitor.ts

@@ -3,34 +3,39 @@ import { isNode, type NodeKind } from '@codama/nodes';
 import { interceptFirstVisitVisitor } from './interceptFirstVisitVisitor';
 import { interceptVisitor } from './interceptVisitor';
 import { LINKABLE_NODES, LinkableDictionary } from './LinkableDictionary';
+import { NodeStack } from './NodeStack';
 import { pipe } from './pipe';
 import { recordNodeStackVisitor } from './recordNodeStackVisitor';
 import { visit, Visitor } from './visitor';
 import { voidVisitor } from './voidVisitor';
 
-export function recordLinkablesVisitor<TReturn, TNodeKind extends NodeKind>(
-    visitor: Visitor<TReturn, TNodeKind>,
+export function getRecordLinkablesVisitor<TNodeKind extends NodeKind>(
     linkables: LinkableDictionary,
-): Visitor<TReturn, TNodeKind> {
-    const recordingVisitor = pipe(
+): Visitor<void, TNodeKind> {
+    const stack = new NodeStack();
+    return pipe(
         voidVisitor(),
         v =>
             interceptVisitor(v, (node, next) => {
                 if (isNode(node, LINKABLE_NODES)) {
-                    linkables.record(node);
+                    linkables.record(node, stack);
                 }
                 return next(node);
             }),
-        v => recordNodeStackVisitor(v, linkables.stack),
+        v => recordNodeStackVisitor(v, stack),
     );
+}
 
-    return pipe(
-        visitor,
-        v =>
-            interceptFirstVisitVisitor(v, (node, next) => {
-                visit(node, recordingVisitor);
-                return next(node);
-            }),
-        v => recordNodeStackVisitor(v, linkables.stack),
+export function recordLinkablesOnFirstVisitVisitor<TReturn, TNodeKind extends NodeKind>(
+    visitor: Visitor<TReturn, TNodeKind>,
+    linkables: LinkableDictionary,
+): Visitor<TReturn, TNodeKind> {
+    const recordingVisitor = getRecordLinkablesVisitor(linkables);
+
+    return pipe(visitor, v =>
+        interceptFirstVisitVisitor(v, (node, next) => {
+            visit(node, recordingVisitor);
+            return next(node);
+        }),
     );
 }

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

@@ -15,10 +15,12 @@ import {
 } from '@codama/nodes';
 import { expect, test } from 'vitest';
 
-import { getByteSizeVisitor, LinkableDictionary, visit, Visitor } from '../src';
+import { getByteSizeVisitor, LinkableDictionary, NodeStack, visit, Visitor } from '../src';
 
 const expectSize = (node: Node, expectedSize: number | null) => {
-    expect(visit(node, getByteSizeVisitor(new LinkableDictionary()) as Visitor<number | null>)).toBe(expectedSize);
+    expect(visit(node, getByteSizeVisitor(new LinkableDictionary(), new NodeStack()) as Visitor<number | null>)).toBe(
+        expectedSize,
+    );
 };
 
 test.each([

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

@@ -26,7 +26,8 @@ import {
     interceptFirstVisitVisitor,
     interceptVisitor,
     LinkableDictionary,
-    recordLinkablesVisitor,
+    NodeStack,
+    recordLinkablesOnFirstVisitVisitor,
     visit,
     voidVisitor,
 } from '../src';
@@ -37,16 +38,17 @@ test('it records program nodes', () => {
         programNode({ name: 'programB', publicKey: '2222' }),
     ]);
 
-    // And a recordLinkablesVisitor extending any visitor.
+    // And a recordLinkablesOnFirstVisitVisitor extending any visitor.
     const linkables = new LinkableDictionary();
-    const visitor = recordLinkablesVisitor(voidVisitor(), linkables);
+    const visitor = recordLinkablesOnFirstVisitVisitor(voidVisitor(), linkables);
 
     // When we visit the tree.
     visit(node, visitor);
 
     // Then we expect program nodes to be recorded and retrievable.
-    expect(linkables.get(programLinkNode('programA'))).toEqual(node.program);
-    expect(linkables.get(programLinkNode('programB'))).toEqual(node.additionalPrograms[0]);
+    const emptyStack = new NodeStack();
+    expect(linkables.get(programLinkNode('programA'), emptyStack)).toEqual(node.program);
+    expect(linkables.get(programLinkNode('programB'), emptyStack)).toEqual(node.additionalPrograms[0]);
 });
 
 test('it records account nodes', () => {
@@ -57,16 +59,17 @@ test('it records account nodes', () => {
         publicKey: '1111',
     });
 
-    // And a recordLinkablesVisitor extending any visitor.
+    // And a recordLinkablesOnFirstVisitVisitor extending any visitor.
     const linkables = new LinkableDictionary();
-    const visitor = recordLinkablesVisitor(voidVisitor(), linkables);
+    const visitor = recordLinkablesOnFirstVisitVisitor(voidVisitor(), linkables);
 
     // When we visit the tree.
     visit(node, visitor);
 
     // Then we expect account nodes to be recorded and retrievable.
-    expect(linkables.get(accountLinkNode('accountA', 'myProgram'))).toEqual(node.accounts[0]);
-    expect(linkables.get(accountLinkNode('accountB', 'myProgram'))).toEqual(node.accounts[1]);
+    const emptyStack = new NodeStack();
+    expect(linkables.get(accountLinkNode('accountA', 'myProgram'), emptyStack)).toEqual(node.accounts[0]);
+    expect(linkables.get(accountLinkNode('accountB', 'myProgram'), emptyStack)).toEqual(node.accounts[1]);
 });
 
 test('it records defined type nodes', () => {
@@ -80,16 +83,17 @@ test('it records defined type nodes', () => {
         publicKey: '1111',
     });
 
-    // And a recordLinkablesVisitor extending any visitor.
+    // And a recordLinkablesOnFirstVisitVisitor extending any visitor.
     const linkables = new LinkableDictionary();
-    const visitor = recordLinkablesVisitor(voidVisitor(), linkables);
+    const visitor = recordLinkablesOnFirstVisitVisitor(voidVisitor(), linkables);
 
     // When we visit the tree.
     visit(node, visitor);
 
     // Then we expect defined type nodes to be recorded and retrievable.
-    expect(linkables.get(definedTypeLinkNode('typeA', 'myProgram'))).toEqual(node.definedTypes[0]);
-    expect(linkables.get(definedTypeLinkNode('typeB', 'myProgram'))).toEqual(node.definedTypes[1]);
+    const emptyStack = new NodeStack();
+    expect(linkables.get(definedTypeLinkNode('typeA', 'myProgram'), emptyStack)).toEqual(node.definedTypes[0]);
+    expect(linkables.get(definedTypeLinkNode('typeB', 'myProgram'), emptyStack)).toEqual(node.definedTypes[1]);
 });
 
 test('it records pda nodes', () => {
@@ -100,16 +104,17 @@ test('it records pda nodes', () => {
         publicKey: '1111',
     });
 
-    // And a recordLinkablesVisitor extending any visitor.
+    // And a recordLinkablesOnFirstVisitVisitor extending any visitor.
     const linkables = new LinkableDictionary();
-    const visitor = recordLinkablesVisitor(voidVisitor(), linkables);
+    const visitor = recordLinkablesOnFirstVisitVisitor(voidVisitor(), linkables);
 
     // When we visit the tree.
     visit(node, visitor);
 
     // Then we expect pda nodes to be recorded and retrievable.
-    expect(linkables.get(pdaLinkNode('pdaA', 'myProgram'))).toEqual(node.pdas[0]);
-    expect(linkables.get(pdaLinkNode('pdaB', 'myProgram'))).toEqual(node.pdas[1]);
+    const emptyStack = new NodeStack();
+    expect(linkables.get(pdaLinkNode('pdaA', 'myProgram'), emptyStack)).toEqual(node.pdas[0]);
+    expect(linkables.get(pdaLinkNode('pdaB', 'myProgram'), emptyStack)).toEqual(node.pdas[1]);
 });
 
 test('it records instruction nodes', () => {
@@ -120,81 +125,76 @@ test('it records instruction nodes', () => {
         publicKey: '1111',
     });
 
-    // And a recordLinkablesVisitor extending any visitor.
+    // And a recordLinkablesOnFirstVisitVisitor extending any visitor.
     const linkables = new LinkableDictionary();
-    const visitor = recordLinkablesVisitor(voidVisitor(), linkables);
+    const visitor = recordLinkablesOnFirstVisitVisitor(voidVisitor(), linkables);
 
     // When we visit the tree.
     visit(node, visitor);
 
     // Then we expect instruction nodes to be recorded and retrievable.
-    expect(linkables.get(instructionLinkNode('instructionA', 'myProgram'))).toEqual(node.instructions[0]);
-    expect(linkables.get(instructionLinkNode('instructionB', 'myProgram'))).toEqual(node.instructions[1]);
+    const emptyStack = new NodeStack();
+    expect(linkables.get(instructionLinkNode('instructionA', 'myProgram'), emptyStack)).toEqual(node.instructions[0]);
+    expect(linkables.get(instructionLinkNode('instructionB', 'myProgram'), emptyStack)).toEqual(node.instructions[1]);
 });
 
 test('it records instruction account nodes', () => {
     // Given the following instruction node containing multiple accounts.
+    const instructionAccounts = [
+        instructionAccountNode({ isSigner: true, isWritable: false, name: 'accountA' }),
+        instructionAccountNode({ isSigner: false, isWritable: true, name: 'accountB' }),
+    ];
     const node = programNode({
-        instructions: [
-            instructionNode({
-                accounts: [
-                    instructionAccountNode({ isSigner: true, isWritable: false, name: 'accountA' }),
-                    instructionAccountNode({ isSigner: false, isWritable: true, name: 'accountB' }),
-                ],
-                name: 'myInstruction',
-            }),
-        ],
+        instructions: [instructionNode({ accounts: instructionAccounts, name: 'myInstruction' })],
         name: 'myProgram',
         publicKey: '1111',
     });
 
-    // And a recordLinkablesVisitor extending any visitor.
+    // And a recordLinkablesOnFirstVisitVisitor extending any visitor.
     const linkables = new LinkableDictionary();
-    const visitor = recordLinkablesVisitor(voidVisitor(), linkables);
+    const visitor = recordLinkablesOnFirstVisitVisitor(voidVisitor(), linkables);
 
     // When we visit the tree.
     visit(node, visitor);
 
     // Then we expect instruction account nodes to be recorded and retrievable.
+    const emptyStack = new NodeStack();
     const instruction = instructionLinkNode('myInstruction', 'myProgram');
-    expect(linkables.get(instructionAccountLinkNode('accountA', instruction))).toEqual(
-        node.instructions[0].accounts[0],
+    expect(linkables.get(instructionAccountLinkNode('accountA', instruction), emptyStack)).toEqual(
+        instructionAccounts[0],
     );
-    expect(linkables.get(instructionAccountLinkNode('accountB', instruction))).toEqual(
-        node.instructions[0].accounts[1],
+    expect(linkables.get(instructionAccountLinkNode('accountB', instruction), emptyStack)).toEqual(
+        instructionAccounts[1],
     );
 });
 
 test('it records instruction argument nodes', () => {
     // Given the following instruction node containing multiple arguments.
+    const instructionArguments = [
+        instructionArgumentNode({ name: 'argumentA', type: numberTypeNode('u32') }),
+        instructionArgumentNode({ name: 'argumentB', type: numberTypeNode('u32') }),
+    ];
     const node = programNode({
-        instructions: [
-            instructionNode({
-                arguments: [
-                    instructionArgumentNode({ name: 'argumentA', type: numberTypeNode('u32') }),
-                    instructionArgumentNode({ name: 'argumentB', type: numberTypeNode('u32') }),
-                ],
-                name: 'myInstruction',
-            }),
-        ],
+        instructions: [instructionNode({ arguments: instructionArguments, name: 'myInstruction' })],
         name: 'myProgram',
         publicKey: '1111',
     });
 
-    // And a recordLinkablesVisitor extending any visitor.
+    // And a recordLinkablesOnFirstVisitVisitor extending any visitor.
     const linkables = new LinkableDictionary();
-    const visitor = recordLinkablesVisitor(voidVisitor(), linkables);
+    const visitor = recordLinkablesOnFirstVisitVisitor(voidVisitor(), linkables);
 
     // When we visit the tree.
     visit(node, visitor);
 
     // Then we expect instruction argument nodes to be recorded and retrievable.
+    const emptyStack = new NodeStack();
     const instruction = instructionLinkNode('myInstruction', 'myProgram');
-    expect(linkables.get(instructionArgumentLinkNode('argumentA', instruction))).toEqual(
-        node.instructions[0].arguments[0],
+    expect(linkables.get(instructionArgumentLinkNode('argumentA', instruction), emptyStack)).toEqual(
+        instructionArguments[0],
     );
-    expect(linkables.get(instructionArgumentLinkNode('argumentB', instruction))).toEqual(
-        node.instructions[0].arguments[1],
+    expect(linkables.get(instructionArgumentLinkNode('argumentB', instruction), emptyStack)).toEqual(
+        instructionArguments[1],
     );
 });
 
@@ -204,16 +204,17 @@ test('it records all linkable before the first visit of the base visitor', () =>
         programNode({ name: 'programB', publicKey: '2222' }),
     ]);
 
-    // And a recordLinkablesVisitor extending a base visitor that
+    // And a recordLinkablesOnFirstVisitVisitor extending a base visitor that
     // stores the linkable programs available at every visit.
     const linkables = new LinkableDictionary();
+    const emptyStack = new NodeStack();
     const events: string[] = [];
     const baseVisitor = interceptFirstVisitVisitor(voidVisitor(), (node, next) => {
-        events.push(`programA:${linkables.has(programLinkNode('programA'))}`);
-        events.push(`programB:${linkables.has(programLinkNode('programB'))}`);
+        events.push(`programA:${linkables.has(programLinkNode('programA'), emptyStack)}`);
+        events.push(`programB:${linkables.has(programLinkNode('programB'), emptyStack)}`);
         next(node);
     });
-    const visitor = recordLinkablesVisitor(baseVisitor, linkables);
+    const visitor = recordLinkablesOnFirstVisitVisitor(baseVisitor, linkables);
 
     // When we visit the tree.
     visit(node, visitor);
@@ -236,17 +237,20 @@ test('it keeps track of the current program when extending a visitor', () => {
     });
     const node = rootNode(programA, [programB]);
 
-    // And a recordLinkablesVisitor extending a base visitor that checks
+    // And a recordLinkablesOnFirstVisitVisitor extending a base visitor that checks
     // the result of getting the linkable node with the same name for each program.
     const linkables = new LinkableDictionary();
+    const stack = new NodeStack();
     const dictionary: Record<string, AccountNode> = {};
     const baseVisitor = interceptVisitor(voidVisitor(), (node, next) => {
+        stack.push(node);
         if (isNode(node, 'programNode')) {
-            dictionary[node.name] = linkables.getOrThrow(accountLinkNode('someAccount'));
+            dictionary[node.name] = linkables.getOrThrow(accountLinkNode('someAccount'), stack);
         }
         next(node);
+        stack.pop();
     });
-    const visitor = recordLinkablesVisitor(baseVisitor, linkables);
+    const visitor = recordLinkablesOnFirstVisitVisitor(baseVisitor, linkables);
 
     // When we visit the tree.
     visit(node, visitor);
@@ -273,17 +277,20 @@ test('it keeps track of the current instruction when extending a visitor', () =>
         publicKey: '1111',
     });
 
-    // And a recordLinkablesVisitor extending a base visitor that checks
+    // And a recordLinkablesOnFirstVisitVisitor extending a base visitor that checks
     // the result of getting the linkable node with the same name for each instruction.
     const linkables = new LinkableDictionary();
+    const stack = new NodeStack();
     const dictionary: Record<string, InstructionAccountNode> = {};
     const baseVisitor = interceptVisitor(voidVisitor(), (node, next) => {
+        stack.push(node);
         if (isNode(node, 'instructionNode')) {
-            dictionary[node.name] = linkables.getOrThrow(instructionAccountLinkNode('someAccount'));
+            dictionary[node.name] = linkables.getOrThrow(instructionAccountLinkNode('someAccount'), stack);
         }
         next(node);
+        stack.pop();
     });
-    const visitor = recordLinkablesVisitor(baseVisitor, linkables);
+    const visitor = recordLinkablesOnFirstVisitVisitor(baseVisitor, linkables);
 
     // When we visit the tree.
     visit(node, visitor);
@@ -297,15 +304,16 @@ test('it does not record linkable types that are not under a program node', () =
     // Given the following account node that is not under a program node.
     const node = accountNode({ name: 'someAccount' });
 
-    // And a recordLinkablesVisitor extending a void visitor.
+    // And a recordLinkablesOnFirstVisitVisitor extending a void visitor.
     const linkables = new LinkableDictionary();
-    const visitor = recordLinkablesVisitor(voidVisitor(), linkables);
+    const visitor = recordLinkablesOnFirstVisitVisitor(voidVisitor(), linkables);
 
     // When we visit the node.
     visit(node, visitor);
 
     // Then we expect the account node to not be recorded.
-    expect(linkables.has(accountLinkNode('someAccount'))).toBe(false);
+    const emptyStack = new NodeStack();
+    expect(linkables.has(accountLinkNode('someAccount'), emptyStack)).toBe(false);
 });
 
 test('it can throw an exception when trying to retrieve a missing linked node', () => {
@@ -318,11 +326,12 @@ test('it can throw an exception when trying to retrieve a missing linked node',
 
     // And a recorded LinkableDictionary.
     const linkables = new LinkableDictionary();
-    const visitor = recordLinkablesVisitor(voidVisitor(), linkables);
+    const visitor = recordLinkablesOnFirstVisitVisitor(voidVisitor(), linkables);
     visit(node, visitor);
 
     // When we try to retrieve a missing account node.
-    const getMissingAccount = () => linkables.getOrThrow(accountLinkNode('missingAccount', 'myProgram'));
+    const emptyStack = new NodeStack();
+    const getMissingAccount = () => linkables.getOrThrow(accountLinkNode('missingAccount', 'myProgram'), emptyStack);
 
     // Then we expect an exception to be thrown.
     expect(getMissingAccount).toThrow(

packages/visitors/src/createSubInstructionsFromEnumArgsVisitor.ts

@@ -14,13 +14,17 @@ import {
     BottomUpNodeTransformerWithSelector,
     bottomUpTransformerVisitor,
     LinkableDictionary,
-    recordLinkablesVisitor,
+    NodeStack,
+    pipe,
+    recordLinkablesOnFirstVisitVisitor,
+    recordNodeStackVisitor,
 } from '@codama/visitors-core';
 
 import { flattenInstructionArguments } from './flattenInstructionDataArgumentsVisitor';
 
 export function createSubInstructionsFromEnumArgsVisitor(map: Record<string, string>) {
     const linkables = new LinkableDictionary();
+    const stack = new NodeStack();
 
     const visitor = bottomUpTransformerVisitor(
         Object.entries(map).map(
@@ -44,8 +48,8 @@ export function createSubInstructionsFromEnumArgsVisitor(map: Record<string, str
                     let argType: EnumTypeNode;
                     if (isNode(argField.type, 'enumTypeNode')) {
                         argType = argField.type;
-                    } else if (isNode(argField.type, 'definedTypeLinkNode') && linkables.has(argField.type)) {
-                        const linkedType = linkables.get(argField.type)?.type ?? null;
+                    } else if (isNode(argField.type, 'definedTypeLinkNode') && linkables.has(argField.type, stack)) {
+                        const linkedType = linkables.get(argField.type, stack)?.type ?? null;
                         assertIsNode(linkedType, 'enumTypeNode');
                         argType = linkedType;
                     } else {
@@ -100,5 +104,9 @@ export function createSubInstructionsFromEnumArgsVisitor(map: Record<string, str
         ),
     );
 
-    return recordLinkablesVisitor(visitor, linkables);
+    return pipe(
+        visitor,
+        v => recordNodeStackVisitor(v, stack),
+        v => recordLinkablesOnFirstVisitVisitor(v, linkables),
+    );
 }

packages/visitors/src/fillDefaultPdaSeedValuesVisitor.ts

@@ -14,7 +14,7 @@ import {
     pdaSeedValueNode,
     pdaValueNode,
 } from '@codama/nodes';
-import { extendVisitor, identityVisitor, LinkableDictionary, pipe, Visitor } from '@codama/visitors-core';
+import { extendVisitor, identityVisitor, LinkableDictionary, NodeStack, pipe, Visitor } from '@codama/visitors-core';
 
 /**
  * Fills in default values for variable PDA seeds that are not explicitly provided.
@@ -30,6 +30,7 @@ import { extendVisitor, identityVisitor, LinkableDictionary, pipe, Visitor } fro
  */
 export function fillDefaultPdaSeedValuesVisitor(
     instruction: InstructionNode,
+    stack: NodeStack,
     linkables: LinkableDictionary,
     strictMode: boolean = false,
 ) {
@@ -38,7 +39,9 @@ export function fillDefaultPdaSeedValuesVisitor(
             visitPdaValue(node, { next }) {
                 const visitedNode = next(node);
                 assertIsNode(visitedNode, 'pdaValueNode');
-                const foundPda = isNode(visitedNode.pda, 'pdaNode') ? visitedNode.pda : linkables.get(visitedNode.pda);
+                const foundPda = isNode(visitedNode.pda, 'pdaNode')
+                    ? visitedNode.pda
+                    : linkables.get(visitedNode.pda, stack);
                 if (!foundPda) return visitedNode;
                 const seeds = addDefaultSeedValuesFromPdaWhenMissing(instruction, foundPda, visitedNode.seeds);
                 if (strictMode && !allSeedsAreValid(instruction, foundPda, seeds)) {

packages/visitors/src/setFixedAccountSizesVisitor.ts

@@ -2,14 +2,18 @@ import { accountNode, assertIsNode, isNode } from '@codama/nodes';
 import {
     getByteSizeVisitor,
     LinkableDictionary,
-    recordLinkablesVisitor,
+    NodeStack,
+    pipe,
+    recordLinkablesOnFirstVisitVisitor,
+    recordNodeStackVisitor,
     topDownTransformerVisitor,
     visit,
 } from '@codama/visitors-core';
 
 export function setFixedAccountSizesVisitor() {
     const linkables = new LinkableDictionary();
-    const byteSizeVisitor = getByteSizeVisitor(linkables);
+    const stack = new NodeStack();
+    const byteSizeVisitor = getByteSizeVisitor(linkables, stack);
 
     const visitor = topDownTransformerVisitor(
         [
@@ -26,5 +30,9 @@ export function setFixedAccountSizesVisitor() {
         ['rootNode', 'programNode', 'accountNode'],
     );
 
-    return recordLinkablesVisitor(visitor, linkables);
+    return pipe(
+        visitor,
+        v => recordNodeStackVisitor(v, stack),
+        v => recordLinkablesOnFirstVisitVisitor(v, linkables),
+    );
 }

packages/visitors/src/setInstructionAccountDefaultValuesVisitor.ts

@@ -12,9 +12,11 @@ import {
 import {
     extendVisitor,
     LinkableDictionary,
+    NodeStack,
     nonNullableIdentityVisitor,
     pipe,
-    recordLinkablesVisitor,
+    recordLinkablesOnFirstVisitVisitor,
+    recordNodeStackVisitor,
     visit,
 } from '@codama/visitors-core';
 
@@ -136,6 +138,7 @@ export const getCommonInstructionAccountDefaultRules = (): InstructionAccountDef
 
 export function setInstructionAccountDefaultValuesVisitor(rules: InstructionAccountDefaultRule[]) {
     const linkables = new LinkableDictionary();
+    const stack = new NodeStack();
 
     // Place the rules with instructions first.
     const sortedRules = rules.sort((a, b) => {
@@ -161,7 +164,6 @@ export function setInstructionAccountDefaultValuesVisitor(rules: InstructionAcco
 
     return pipe(
         nonNullableIdentityVisitor(['rootNode', 'programNode', 'instructionNode']),
-        v => recordLinkablesVisitor(v, linkables),
         v =>
             extendVisitor(v, {
                 visitInstruction(node) {
@@ -178,7 +180,7 @@ export function setInstructionAccountDefaultValuesVisitor(rules: InstructionAcco
                                 ...account,
                                 defaultValue: visit(
                                     rule.defaultValue,
-                                    fillDefaultPdaSeedValuesVisitor(node, linkables, true),
+                                    fillDefaultPdaSeedValuesVisitor(node, stack, linkables, true),
                                 ),
                             };
                         } catch (error) {
@@ -192,5 +194,7 @@ export function setInstructionAccountDefaultValuesVisitor(rules: InstructionAcco
                     });
                 },
             }),
+        v => recordNodeStackVisitor(v, stack),
+        v => recordLinkablesOnFirstVisitVisitor(v, linkables),
     );
 }

packages/visitors/src/unwrapDefinedTypesVisitor.ts

@@ -2,14 +2,17 @@ import { assertIsNodeFilter, camelCase, CamelCaseString, programNode } from '@co
 import {
     extendVisitor,
     LinkableDictionary,
+    NodeStack,
     nonNullableIdentityVisitor,
     pipe,
-    recordLinkablesVisitor,
+    recordLinkablesOnFirstVisitVisitor,
+    recordNodeStackVisitor,
     visit,
 } from '@codama/visitors-core';
 
 export function unwrapDefinedTypesVisitor(typesToInline: string[] | '*' = '*') {
     const linkables = new LinkableDictionary();
+    const stack = new NodeStack();
     const typesToInlineMainCased = typesToInline === '*' ? '*' : typesToInline.map(camelCase);
     const shouldInline = (definedType: CamelCaseString): boolean =>
         typesToInlineMainCased === '*' || typesToInlineMainCased.includes(definedType);
@@ -22,7 +25,7 @@ export function unwrapDefinedTypesVisitor(typesToInline: string[] | '*' = '*') {
                     if (!shouldInline(linkType.name)) {
                         return linkType;
                     }
-                    return visit(linkables.getOrThrow(linkType).type, self);
+                    return visit(linkables.getOrThrow(linkType, stack).type, self);
                 },
 
                 visitProgram(program, { self }) {
@@ -41,6 +44,7 @@ export function unwrapDefinedTypesVisitor(typesToInline: string[] | '*' = '*') {
                     });
                 },
             }),
-        v => recordLinkablesVisitor(v, linkables),
+        v => recordNodeStackVisitor(v, stack),
+        v => recordLinkablesOnFirstVisitVisitor(v, linkables),
     );
 }

packages/visitors/src/unwrapTypeDefinedLinksVisitor.ts

@@ -3,20 +3,27 @@ import {
     BottomUpNodeTransformerWithSelector,
     bottomUpTransformerVisitor,
     LinkableDictionary,
+    NodeStack,
     pipe,
-    recordLinkablesVisitor,
+    recordLinkablesOnFirstVisitVisitor,
+    recordNodeStackVisitor,
 } from '@codama/visitors-core';
 
 export function unwrapTypeDefinedLinksVisitor(definedLinksType: string[]) {
     const linkables = new LinkableDictionary();
+    const stack = new NodeStack();
 
     const transformers: BottomUpNodeTransformerWithSelector[] = definedLinksType.map(selector => ({
         select: ['[definedTypeLinkNode]', selector],
         transform: node => {
             assertIsNode(node, 'definedTypeLinkNode');
-            return linkables.getOrThrow(node).type;
+            return linkables.getOrThrow(node, stack).type;
         },
     }));
 
-    return pipe(bottomUpTransformerVisitor(transformers), v => recordLinkablesVisitor(v, linkables));
+    return pipe(
+        bottomUpTransformerVisitor(transformers),
+        v => recordNodeStackVisitor(v, stack),
+        v => recordLinkablesOnFirstVisitVisitor(v, linkables),
+    );
 }

packages/visitors/src/updateInstructionsVisitor.ts

@@ -16,8 +16,10 @@ import {
     BottomUpNodeTransformerWithSelector,
     bottomUpTransformerVisitor,
     LinkableDictionary,
+    NodeStack,
     pipe,
-    recordLinkablesVisitor,
+    recordLinkablesOnFirstVisitVisitor,
+    recordNodeStackVisitor,
     visit,
 } from '@codama/visitors-core';
 
@@ -59,6 +61,7 @@ export type InstructionArgumentUpdates = Record<
 
 export function updateInstructionsVisitor(map: Record<string, InstructionUpdates>) {
     const linkables = new LinkableDictionary();
+    const stack = new NodeStack();
 
     const transformers = Object.entries(map).map(
         ([selector, updates]): BottomUpNodeTransformerWithSelector => ({
@@ -72,7 +75,7 @@ export function updateInstructionsVisitor(map: Record<string, InstructionUpdates
                 const { accounts: accountUpdates, arguments: argumentUpdates, ...metadataUpdates } = updates;
                 const { newArguments, newExtraArguments } = handleInstructionArguments(node, argumentUpdates ?? {});
                 const newAccounts = node.accounts.map(account =>
-                    handleInstructionAccount(node, account, accountUpdates ?? {}, linkables),
+                    handleInstructionAccount(node, stack, account, accountUpdates ?? {}, linkables),
                 );
                 return instructionNode({
                     ...node,
@@ -85,11 +88,16 @@ export function updateInstructionsVisitor(map: Record<string, InstructionUpdates
         }),
     );
 
-    return pipe(bottomUpTransformerVisitor(transformers), v => recordLinkablesVisitor(v, linkables));
+    return pipe(
+        bottomUpTransformerVisitor(transformers),
+        v => recordNodeStackVisitor(v, stack),
+        v => recordLinkablesOnFirstVisitVisitor(v, linkables),
+    );
 }
 
 function handleInstructionAccount(
     instruction: InstructionNode,
+    stack: NodeStack,
     account: InstructionAccountNode,
     accountUpdates: InstructionAccountUpdates,
     linkables: LinkableDictionary,
@@ -107,7 +115,7 @@ function handleInstructionAccount(
 
     return instructionAccountNode({
         ...acountWithoutDefault,
-        defaultValue: visit(defaultValue, fillDefaultPdaSeedValuesVisitor(instruction, linkables)),
+        defaultValue: visit(defaultValue, fillDefaultPdaSeedValuesVisitor(instruction, stack, linkables)),
     });
 }
 

packages/visitors/test/fillDefaultPdaSeedValuesVisitor.test.ts

@@ -14,7 +14,7 @@ import {
     publicKeyTypeNode,
     variablePdaSeedNode,
 } from '@codama/nodes';
-import { LinkableDictionary, visit } from '@codama/visitors-core';
+import { LinkableDictionary, NodeStack, visit } from '@codama/visitors-core';
 import { expect, test } from 'vitest';
 
 import { fillDefaultPdaSeedValuesVisitor } from '../src';
@@ -39,8 +39,7 @@ test('it fills missing pda seed values with default values', () => {
 
     // And a linkable dictionary that recorded this PDA.
     const linkables = new LinkableDictionary();
-    linkables.stack.push(program);
-    linkables.record(pda);
+    linkables.record(pda, new NodeStack([program, pda]));
 
     // And a pdaValueNode with a single seed filled.
     const node = pdaValueNode('myPda', [pdaSeedValueNode('seed1', numberValueNode(42))]);
@@ -57,9 +56,10 @@ test('it fills missing pda seed values with default values', () => {
         arguments: [instructionArgumentNode({ name: 'seed2', type: numberTypeNode('u64') })],
         name: 'myInstruction',
     });
+    const instructionStack = new NodeStack([program, instruction]);
 
     // When we fill the PDA seeds with default values.
-    const result = visit(node, fillDefaultPdaSeedValuesVisitor(instruction, linkables));
+    const result = visit(node, fillDefaultPdaSeedValuesVisitor(instruction, instructionStack, linkables));
 
     // Then we expect the following pdaValueNode to be returned.
     expect(result).toEqual(
@@ -91,8 +91,7 @@ test('it fills nested pda value nodes', () => {
 
     // And a linkable dictionary that recorded this PDA.
     const linkables = new LinkableDictionary();
-    linkables.stack.push(program);
-    linkables.record(pda);
+    linkables.record(pda, new NodeStack([program, pda]));
 
     // And a pdaValueNode nested inside a conditionalValueNode.
     const node = conditionalValueNode({
@@ -112,9 +111,10 @@ test('it fills nested pda value nodes', () => {
         arguments: [instructionArgumentNode({ name: 'seed2', type: numberTypeNode('u64') })],
         name: 'myInstruction',
     });
+    const instructionStack = new NodeStack([program, instruction]);
 
     // When we fill the PDA seeds with default values.
-    const result = visit(node, fillDefaultPdaSeedValuesVisitor(instruction, linkables));
+    const result = visit(node, fillDefaultPdaSeedValuesVisitor(instruction, instructionStack, linkables));
 
     // Then we expect the following conditionalValueNode to be returned.
     expect(result).toEqual(
@@ -149,8 +149,7 @@ test('it ignores default seeds missing from the instruction', () => {
 
     // And a linkable dictionary that recorded this PDA.
     const linkables = new LinkableDictionary();
-    linkables.stack.push(program);
-    linkables.record(pda);
+    linkables.record(pda, new NodeStack([program, pda]));
 
     // And a pdaValueNode with a single seed filled.
     const node = pdaValueNode('myPda', [pdaSeedValueNode('seed1', numberValueNode(42))]);
@@ -160,9 +159,10 @@ test('it ignores default seeds missing from the instruction', () => {
         arguments: [instructionArgumentNode({ name: 'seed2', type: numberTypeNode('u64') })],
         name: 'myInstruction',
     });
+    const instructionStack = new NodeStack([program, instruction]);
 
     // When we fill the PDA seeds with default values.
-    const result = visit(node, fillDefaultPdaSeedValuesVisitor(instruction, linkables));
+    const result = visit(node, fillDefaultPdaSeedValuesVisitor(instruction, instructionStack, linkables));
 
     // Then we expect the following pdaValueNode to be returned.
     expect(result).toEqual(