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.

Author: lorisleiva

.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)