Add program to link nodes and update LinkableDictionary (#180)

.changeset/tough-grapes-give.md

@@ -0,0 +1,9 @@
+---
+'@kinobi-so/visitors-core': minor
+'@kinobi-so/node-types': minor
+'@kinobi-so/nodes': minor
+'@kinobi-so/visitors': patch
+'@kinobi-so/errors': patch
+---
+
+Add optional `program` attribute to link nodes and namespace linkable nodes under their associated program.

packages/errors/src/context.ts

@@ -75,6 +75,7 @@ export type KinobiErrorContext = DefaultUnspecifiedErrorContextToUndefined<{
         kind: LinkNode['kind'];
         linkNode: LinkNode;
         name: CamelCaseString;
+        program?: CamelCaseString;
     };
     [KINOBI_ERROR__NODE_FILESYSTEM_FUNCTION_UNAVAILABLE]: {
         fsFunction: string;

packages/node-types/src/linkNodes/AccountLinkNode.ts

@@ -1,8 +1,12 @@
 import type { CamelCaseString } from '../shared';
+import type { ProgramLinkNode } from './ProgramLinkNode';
 
-export interface AccountLinkNode {
+export interface AccountLinkNode<TProgram extends ProgramLinkNode | undefined = ProgramLinkNode | undefined> {
     readonly kind: 'accountLinkNode';
 
+    // Children.
+    readonly program?: TProgram;
+
     // Data.
     readonly name: CamelCaseString;
 }

packages/node-types/src/linkNodes/DefinedTypeLinkNode.ts

@@ -1,8 +1,12 @@
 import type { CamelCaseString } from '../shared';
+import type { ProgramLinkNode } from './ProgramLinkNode';
 
-export interface DefinedTypeLinkNode {
+export interface DefinedTypeLinkNode<TProgram extends ProgramLinkNode | undefined = ProgramLinkNode | undefined> {
     readonly kind: 'definedTypeLinkNode';
 
+    // Children.
+    readonly program?: TProgram;
+
     // Data.
     readonly name: CamelCaseString;
 }

packages/node-types/src/linkNodes/PdaLinkNode.ts

@@ -1,8 +1,12 @@
 import type { CamelCaseString } from '../shared';
+import type { ProgramLinkNode } from './ProgramLinkNode';
 
-export interface PdaLinkNode {
+export interface PdaLinkNode<TProgram extends ProgramLinkNode | undefined = ProgramLinkNode | undefined> {
     readonly kind: 'pdaLinkNode';
 
+    // Children.
+    readonly program?: TProgram;
+
     // Data.
     readonly name: CamelCaseString;
 }

packages/nodes/docs/linkNodes/AccountLinkNode.md

@@ -13,14 +13,17 @@ This node represents a reference to an existing [`AccountNode`](../AccountNode.m
 
 ### Children
 
-_This node has no children._
+| Attribute | Type                                      | Description                                                                                                     |
+| --------- | ----------------------------------------- | --------------------------------------------------------------------------------------------------------------- |
+| `program` | [`ProgramLinkNode`](./ProgramLinkNode.md) | (Optional) The program associated with the linked account. Default to using the program we are currently under. |
 
 ## Functions
 
-### `accountLinkNode(name)`
+### `accountLinkNode(name, program?)`
 
-Helper function that creates a `AccountLinkNode` object from the name of the `AccountNode` we are referring to.
+Helper function that creates a `AccountLinkNode` object from the name of the `AccountNode` we are referring to. If the account is from another program, the `program` parameter must be provided as either a `string` or a `ProgramLinkNode`.
 
 ```ts
 const node = accountLinkNode('myAccount');
+const nodeFromAnotherProgram = accountLinkNode('myAccount', 'myOtherProgram');
 ```

packages/nodes/docs/linkNodes/DefinedTypeLinkNode.md

@@ -13,14 +13,17 @@ This node represents a reference to an existing [`DefinedTypeNode`](../DefinedTy
 
 ### Children
 
-_This node has no children._
+| Attribute | Type                                      | Description                                                                                                  |
+| --------- | ----------------------------------------- | ------------------------------------------------------------------------------------------------------------ |
+| `program` | [`ProgramLinkNode`](./ProgramLinkNode.md) | (Optional) The program associated with the linked type. Default to using the program we are currently under. |
 
 ## Functions
 
-### `definedTypeLinkNode(name)`
+### `definedTypeLinkNode(name, program?)`
 
-Helper function that creates a `DefinedTypeLinkNode` object from the name of the `DefinedTypeNode` we are referring to.
+Helper function that creates a `DefinedTypeLinkNode` object from the name of the `DefinedTypeNode` we are referring to. If the defined type is from another program, the `program` parameter must be provided as either a `string` or a `ProgramLinkNode`.
 
 ```ts
 const node = definedTypeLinkNode('myDefinedType');
+const nodeFromAnotherProgram = definedTypeLinkNode('myDefinedType', 'myOtherProgram');
 ```

packages/nodes/docs/linkNodes/PdaLinkNode.md

@@ -13,14 +13,17 @@ This node represents a reference to an existing [`PdaNode`](../PdaNode.md) in th
 
 ### Children
 
-_This node has no children._
+| Attribute | Type                                      | Description                                                                                                 |
+| --------- | ----------------------------------------- | ----------------------------------------------------------------------------------------------------------- |
+| `program` | [`ProgramLinkNode`](./ProgramLinkNode.md) | (Optional) The program associated with the linked PDA. Default to using the program we are currently under. |
 
 ## Functions
 
-### `pdaLinkNode(name)`
+### `pdaLinkNode(name, program?)`
 
-Helper function that creates a `PdaLinkNode` object from the name of the `PdaNode` we are referring to.
+Helper function that creates a `PdaLinkNode` object from the name of the `PdaNode` we are referring to. If the PDA is from another program, the `program` parameter must be provided as either a `string` or a `ProgramLinkNode`.
 
 ```ts
 const node = pdaLinkNode('myPda');
+const nodeFromAnotherProgram = pdaLinkNode('myPda', 'myOtherProgram');
 ```

packages/nodes/src/linkNodes/AccountLinkNode.ts

@@ -1,11 +1,15 @@
-import type { AccountLinkNode } from '@kinobi-so/node-types';
+import type { AccountLinkNode, ProgramLinkNode } from '@kinobi-so/node-types';
 
 import { camelCase } from '../shared';
+import { programLinkNode } from './ProgramLinkNode';
 
-export function accountLinkNode(name: string): AccountLinkNode {
+export function accountLinkNode(name: string, program?: ProgramLinkNode | string): AccountLinkNode {
     return Object.freeze({
         kind: 'accountLinkNode',
 
+        // Children.
+        ...(program === undefined ? {} : { program: typeof program === 'string' ? programLinkNode(program) : program }),
+
         // Data.
         name: camelCase(name),
     });

packages/nodes/src/linkNodes/DefinedTypeLinkNode.ts

@@ -1,11 +1,15 @@
-import type { DefinedTypeLinkNode } from '@kinobi-so/node-types';
+import type { DefinedTypeLinkNode, ProgramLinkNode } from '@kinobi-so/node-types';
 
 import { camelCase } from '../shared';
+import { programLinkNode } from './ProgramLinkNode';
 
-export function definedTypeLinkNode(name: string): DefinedTypeLinkNode {
+export function definedTypeLinkNode(name: string, program?: ProgramLinkNode | string): DefinedTypeLinkNode {
     return Object.freeze({
         kind: 'definedTypeLinkNode',
 
+        // Children.
+        ...(program === undefined ? {} : { program: typeof program === 'string' ? programLinkNode(program) : program }),
+
         // Data.
         name: camelCase(name),
     });

packages/nodes/src/linkNodes/PdaLinkNode.ts

@@ -1,11 +1,15 @@
-import type { PdaLinkNode } from '@kinobi-so/node-types';
+import type { PdaLinkNode, ProgramLinkNode } from '@kinobi-so/node-types';
 
 import { camelCase } from '../shared';
+import { programLinkNode } from './ProgramLinkNode';
 
-export function pdaLinkNode(name: string): PdaLinkNode {
+export function pdaLinkNode(name: string, program?: ProgramLinkNode | string): PdaLinkNode {
     return Object.freeze({
         kind: 'pdaLinkNode',
 
+        // Children.
+        ...(program === undefined ? {} : { program: typeof program === 'string' ? programLinkNode(program) : program }),
+
         // Data.
         name: camelCase(name),
     });

packages/visitors-core/README.md

@@ -651,11 +651,11 @@ It offers the following API:
 ```ts
 const linkables = new LinkableDictionary();
 
-// Record any linkable node — such as programs, PDAs, accounts and defined types.
+// Record program nodes.
 linkables.record(programNode);
 
-// Record multiple linkable nodes at once.
-linkables.recordAll([...accountNodes, ...pdaNodes]);
+// Record other linkable nodes with their associated program node.
+linkables.record(accountNode);
 
 // Get a linkable node using a link node, or throw an error if it is not found.
 const programNode = linkables.getOrThrow(programLinkNode);
@@ -664,6 +664,8 @@ const programNode = linkables.getOrThrow(programLinkNode);
 const accountNode = linkables.get(accountLinkNode);
 ```
 
+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.
+
 ### `recordLinkablesVisitor`
 
 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.
@@ -676,15 +678,17 @@ Here's an example that records a `LinkableDictionary` and uses it to log the amo
 const linkables = new LinkableDictionary();
 const visitor = pipe(
     baseVisitor,
-    v => recordLinkablesVisitor(v, linkables),
     v =>
         tapVisitor(v, 'pdaLinkNode', node => {
             const pdaNode = linkables.getOrThrow(node);
             console.log(`${pdaNode.seeds.length} seeds`);
         }),
+    v => recordLinkablesVisitor(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.
+
 ## Other useful visitors
 
 This package provides a few other visitors that may help build more complex visitors.

packages/visitors-core/src/LinkableDictionary.ts

@@ -2,6 +2,8 @@ import { KINOBI_ERROR__LINKED_NODE_NOT_FOUND, KinobiError } from '@kinobi-so/err
 import {
     AccountLinkNode,
     AccountNode,
+    camelCase,
+    CamelCaseString,
     DefinedTypeLinkNode,
     DefinedTypeNode,
     isNode,
@@ -12,52 +14,83 @@ import {
     ProgramNode,
 } from '@kinobi-so/nodes';
 
+import { NodeStack } from './NodeStack';
+
 export type LinkableNode = AccountNode | DefinedTypeNode | PdaNode | ProgramNode;
 
 export const LINKABLE_NODES: LinkableNode['kind'][] = ['accountNode', 'definedTypeNode', 'pdaNode', 'programNode'];
 
-export class LinkableDictionary {
-    private readonly programs: Map<string, ProgramNode> = new Map();
+type ProgramDictionary = {
+    accounts: Map<string, AccountNode>;
+    definedTypes: Map<string, DefinedTypeNode>;
+    pdas: Map<string, PdaNode>;
+    program: ProgramNode;
+};
 
-    private readonly pdas: Map<string, PdaNode> = new Map();
+type ProgramInput = ProgramLinkNode | ProgramNode | string;
 
-    private readonly accounts: Map<string, AccountNode> = new Map();
+function getProgramName(program: ProgramInput): CamelCaseString;
+function getProgramName(program: ProgramInput | undefined): CamelCaseString | undefined;
+function getProgramName(program: ProgramInput | undefined): CamelCaseString | undefined {
+    if (!program) return undefined;
+    return typeof program === 'string' ? camelCase(program) : program.name;
+}
 
-    private readonly definedTypes: Map<string, DefinedTypeNode> = new Map();
+export class LinkableDictionary {
+    readonly programs: Map<string, ProgramDictionary> = new Map();
+
+    readonly stack: NodeStack = new NodeStack();
+
+    private getOrCreateProgramDictionary(node: ProgramNode): ProgramDictionary {
+        let programDictionary = this.programs.get(node.name);
+        if (!programDictionary) {
+            programDictionary = {
+                accounts: new Map(),
+                definedTypes: new Map(),
+                pdas: new Map(),
+                program: node,
+            };
+            this.programs.set(node.name, programDictionary);
+        }
+        return programDictionary;
+    }
 
     record(node: LinkableNode): this {
         if (isNode(node, 'programNode')) {
-            this.programs.set(node.name, node);
+            this.getOrCreateProgramDictionary(node);
+            return this;
         }
+
+        // Do not record nodes that are outside of a program.
+        const program = this.stack.getProgram();
+        if (!program) return this;
+
+        const programDictionary = this.getOrCreateProgramDictionary(program);
         if (isNode(node, 'pdaNode')) {
-            this.pdas.set(node.name, node);
-        }
-        if (isNode(node, 'accountNode')) {
-            this.accounts.set(node.name, node);
-        }
-        if (isNode(node, 'definedTypeNode')) {
-            this.definedTypes.set(node.name, node);
+            programDictionary.pdas.set(node.name, node);
+        } else if (isNode(node, 'accountNode')) {
+            programDictionary.accounts.set(node.name, node);
+        } else if (isNode(node, 'definedTypeNode')) {
+            programDictionary.definedTypes.set(node.name, node);
         }
         return this;
     }
 
-    recordAll(nodes: LinkableNode[]): this {
-        nodes.forEach(node => this.record(node));
-        return this;
-    }
-
     getOrThrow(linkNode: ProgramLinkNode): ProgramNode;
     getOrThrow(linkNode: PdaLinkNode): PdaNode;
     getOrThrow(linkNode: AccountLinkNode): AccountNode;
     getOrThrow(linkNode: DefinedTypeLinkNode): DefinedTypeNode;
     getOrThrow(linkNode: LinkNode): LinkableNode {
-        const node = this.get(linkNode as ProgramLinkNode) as LinkableNode;
+        const node = this.get(linkNode as ProgramLinkNode) as LinkableNode | undefined;
 
         if (!node) {
             throw new KinobiError(KINOBI_ERROR__LINKED_NODE_NOT_FOUND, {
                 kind: linkNode.kind,
                 linkNode,
                 name: linkNode.name,
+                program: isNode(linkNode, 'pdaLinkNode')
+                    ? getProgramName(linkNode.program ?? this.stack.getProgram())
+                    : undefined,
             });
         }
 
@@ -70,17 +103,23 @@ export class LinkableDictionary {
     get(linkNode: DefinedTypeLinkNode): DefinedTypeNode | undefined;
     get(linkNode: LinkNode): LinkableNode | undefined {
         if (isNode(linkNode, 'programLinkNode')) {
-            return this.programs.get(linkNode.name);
+            return this.programs.get(linkNode.name)?.program;
         }
+
+        const programName = getProgramName(linkNode.program ?? this.stack.getProgram());
+        if (!programName) return undefined;
+
+        const programDictionary = this.programs.get(programName);
+        if (!programDictionary) return undefined;
+
         if (isNode(linkNode, 'pdaLinkNode')) {
-            return this.pdas.get(linkNode.name);
-        }
-        if (isNode(linkNode, 'accountLinkNode')) {
-            return this.accounts.get(linkNode.name);
-        }
-        if (isNode(linkNode, 'definedTypeLinkNode')) {
-            return this.definedTypes.get(linkNode.name);
+            return programDictionary.pdas.get(linkNode.name);
+        } else if (isNode(linkNode, 'accountLinkNode')) {
+            return programDictionary.accounts.get(linkNode.name);
+        } else if (isNode(linkNode, 'definedTypeLinkNode')) {
+            return programDictionary.definedTypes.get(linkNode.name);
         }
+
         return undefined;
     }
 
@@ -88,15 +127,21 @@ export class LinkableDictionary {
         if (isNode(linkNode, 'programLinkNode')) {
             return this.programs.has(linkNode.name);
         }
+
+        const programName = getProgramName(linkNode.program ?? this.stack.getProgram());
+        if (!programName) return false;
+
+        const programDictionary = this.programs.get(programName);
+        if (!programDictionary) return false;
+
         if (isNode(linkNode, 'pdaLinkNode')) {
-            return this.pdas.has(linkNode.name);
-        }
-        if (isNode(linkNode, 'accountLinkNode')) {
-            return this.accounts.has(linkNode.name);
-        }
-        if (isNode(linkNode, 'definedTypeLinkNode')) {
-            return this.definedTypes.has(linkNode.name);
+            return programDictionary.pdas.has(linkNode.name);
+        } else if (isNode(linkNode, 'accountLinkNode')) {
+            return programDictionary.accounts.has(linkNode.name);
+        } else if (isNode(linkNode, 'definedTypeLinkNode')) {
+            return programDictionary.definedTypes.has(linkNode.name);
         }
+
         return false;
     }
 }

packages/visitors-core/src/NodeStack.ts

@@ -1,4 +1,4 @@
-import { GetNodeFromKind, isNodeFilter, Node, NodeKind, ProgramNode } from '@kinobi-so/nodes';
+import { GetNodeFromKind, isNode, Node, NodeKind, ProgramNode } from '@kinobi-so/nodes';
 
 export class NodeStack {
     private readonly stack: Node[];
@@ -20,7 +20,11 @@ export class NodeStack {
     }
 
     public find<TKind extends NodeKind>(kind: TKind | TKind[]): GetNodeFromKind<TKind> | undefined {
-        return this.stack.find(isNodeFilter(kind));
+        for (let index = this.stack.length - 1; index >= 0; index--) {
+            const node = this.stack[index];
+            if (isNode(node, kind)) return node;
+        }
+        return undefined;
     }
 
     public getProgram(): ProgramNode | undefined {

packages/visitors-core/src/getDebugStringVisitor.ts

@@ -64,10 +64,11 @@ function getNodeDetails(node: Node): string[] {
         case 'errorNode':
             return [node.code.toString(), node.name];
         case 'programLinkNode':
+            return [node.name];
         case 'pdaLinkNode':
         case 'accountLinkNode':
         case 'definedTypeLinkNode':
-            return [node.name];
+            return [...(node.program ? [node.program.name] : []), node.name];
         case 'numberTypeNode':
             return [node.format, ...(node.endian === 'be' ? ['bigEndian'] : [])];
         case 'amountTypeNode':

packages/visitors-core/src/recordLinkablesVisitor.ts

@@ -3,6 +3,8 @@ import { isNode, type NodeKind } from '@kinobi-so/nodes';
 import { interceptFirstVisitVisitor } from './interceptFirstVisitVisitor';
 import { interceptVisitor } from './interceptVisitor';
 import { LINKABLE_NODES, LinkableDictionary } from './LinkableDictionary';
+import { pipe } from './pipe';
+import { recordNodeStackVisitor } from './recordNodeStackVisitor';
 import { visit, Visitor } from './visitor';
 import { voidVisitor } from './voidVisitor';
 
@@ -10,15 +12,25 @@ export function recordLinkablesVisitor<TReturn, TNodeKind extends NodeKind>(
     visitor: Visitor<TReturn, TNodeKind>,
     linkables: LinkableDictionary,
 ): Visitor<TReturn, TNodeKind> {
-    const recordingVisitor = interceptVisitor(voidVisitor(), (node, next) => {
-        if (isNode(node, LINKABLE_NODES)) {
-            linkables.record(node);
-        }
-        return next(node);
-    });
+    const recordingVisitor = pipe(
+        voidVisitor(),
+        v =>
+            interceptVisitor(v, (node, next) => {
+                if (isNode(node, LINKABLE_NODES)) {
+                    linkables.record(node);
+                }
+                return next(node);
+            }),
+        v => recordNodeStackVisitor(v, linkables.stack),
+    );
 
-    return interceptFirstVisitVisitor(visitor, (node, next) => {
-        visit(node, recordingVisitor);
-        return next(node);
-    });
+    return pipe(
+        visitor,
+        v =>
+            interceptFirstVisitVisitor(v, (node, next) => {
+                visit(node, recordingVisitor);
+                return next(node);
+            }),
+        v => recordNodeStackVisitor(v, linkables.stack),
+    );
 }

packages/visitors-core/test/nodes/linkNodes/AccountLinkNode.test.ts

@@ -8,7 +8,7 @@ import {
     expectMergeVisitorCount,
 } from '../_setup';
 
-const node = accountLinkNode('token');
+const node = accountLinkNode('token', 'splToken');
 
 test('mergeVisitor', () => {
     expectMergeVisitorCount(node, 1);
@@ -23,5 +23,5 @@ test('deleteNodesVisitor', () => {
 });
 
 test('debugStringVisitor', () => {
-    expectDebugStringVisitor(node, `accountLinkNode [token]`);
+    expectDebugStringVisitor(node, `accountLinkNode [splToken.token]`);
 });

packages/visitors-core/test/nodes/linkNodes/DefinedTypeLinkNode.test.ts

@@ -8,7 +8,7 @@ import {
     expectMergeVisitorCount,
 } from '../_setup';
 
-const node = definedTypeLinkNode('tokenState');
+const node = definedTypeLinkNode('tokenState', 'splToken');
 
 test('mergeVisitor', () => {
     expectMergeVisitorCount(node, 1);
@@ -23,5 +23,5 @@ test('deleteNodesVisitor', () => {
 });
 
 test('debugStringVisitor', () => {
-    expectDebugStringVisitor(node, `definedTypeLinkNode [tokenState]`);
+    expectDebugStringVisitor(node, `definedTypeLinkNode [splToken.tokenState]`);
 });

packages/visitors-core/test/nodes/linkNodes/PdaLinkNode.test.ts

@@ -8,7 +8,7 @@ import {
     expectMergeVisitorCount,
 } from '../_setup';
 
-const node = pdaLinkNode('associatedToken');
+const node = pdaLinkNode('associatedToken', 'splToken');
 
 test('mergeVisitor', () => {
     expectMergeVisitorCount(node, 1);
@@ -23,5 +23,5 @@ test('deleteNodesVisitor', () => {
 });
 
 test('debugStringVisitor', () => {
-    expectDebugStringVisitor(node, `pdaLinkNode [associatedToken]`);
+    expectDebugStringVisitor(node, `pdaLinkNode [splToken.associatedToken]`);
 });

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

@@ -1,8 +1,10 @@
 import {
     accountLinkNode,
+    AccountNode,
     accountNode,
     definedTypeLinkNode,
     definedTypeNode,
+    isNode,
     pdaLinkNode,
     pdaNode,
     programLinkNode,
@@ -12,7 +14,14 @@ import {
 } from '@kinobi-so/nodes';
 import { expect, test } from 'vitest';
 
-import { interceptFirstVisitVisitor, LinkableDictionary, recordLinkablesVisitor, visit, voidVisitor } from '../src';
+import {
+    interceptFirstVisitVisitor,
+    interceptVisitor,
+    LinkableDictionary,
+    recordLinkablesVisitor,
+    visit,
+    voidVisitor,
+} from '../src';
 
 test('it record all linkable nodes it finds when traversing the tree', () => {
     // Given the following root node containing multiple linkable nodes.
@@ -45,12 +54,12 @@ test('it record all linkable nodes it finds when traversing the tree', () => {
     // Then we expect all linkable nodes to be recorded.
     expect(linkables.get(programLinkNode('programA'))).toEqual(node.program);
     expect(linkables.get(programLinkNode('programB'))).toEqual(node.additionalPrograms[0]);
-    expect(linkables.get(pdaLinkNode('pdaA'))).toEqual(node.program.pdas[0]);
-    expect(linkables.get(pdaLinkNode('pdaB'))).toEqual(node.additionalPrograms[0].pdas[0]);
-    expect(linkables.get(accountLinkNode('accountA'))).toEqual(node.program.accounts[0]);
-    expect(linkables.get(accountLinkNode('accountB'))).toEqual(node.additionalPrograms[0].accounts[0]);
-    expect(linkables.get(definedTypeLinkNode('typeA'))).toEqual(node.program.definedTypes[0]);
-    expect(linkables.get(definedTypeLinkNode('typeB'))).toEqual(node.additionalPrograms[0].definedTypes[0]);
+    expect(linkables.get(pdaLinkNode('pdaA', 'programA'))).toEqual(node.program.pdas[0]);
+    expect(linkables.get(pdaLinkNode('pdaB', 'programB'))).toEqual(node.additionalPrograms[0].pdas[0]);
+    expect(linkables.get(accountLinkNode('accountA', 'programA'))).toEqual(node.program.accounts[0]);
+    expect(linkables.get(accountLinkNode('accountB', 'programB'))).toEqual(node.additionalPrograms[0].accounts[0]);
+    expect(linkables.get(definedTypeLinkNode('typeA', 'programA'))).toEqual(node.program.definedTypes[0]);
+    expect(linkables.get(definedTypeLinkNode('typeB', 'programB'))).toEqual(node.additionalPrograms[0].definedTypes[0]);
 });
 
 test('it records all linkable before the first visit of the base visitor', () => {
@@ -76,3 +85,52 @@ test('it records all linkable before the first visit of the base visitor', () =>
     // Then we expect all linkable nodes to be recorded.
     expect(events).toEqual(['programA:true', 'programB:true']);
 });
+
+test('it keeps track of the current program when extending a visitor', () => {
+    // Given the following root node containing two program containing an account with the same name.
+    const programA = programNode({
+        accounts: [accountNode({ name: 'someAccount' })],
+        name: 'programA',
+        publicKey: '1111',
+    });
+    const programB = programNode({
+        accounts: [accountNode({ name: 'someAccount' })],
+        name: 'programB',
+        publicKey: '2222',
+    });
+    const node = rootNode(programA, [programB]);
+
+    // And a recordLinkablesVisitor 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 dictionary: Record<string, AccountNode> = {};
+    const baseVisitor = interceptVisitor(voidVisitor(), (node, next) => {
+        if (isNode(node, 'programNode')) {
+            dictionary[node.name] = linkables.getOrThrow(accountLinkNode('someAccount'));
+        }
+        next(node);
+    });
+    const visitor = recordLinkablesVisitor(baseVisitor, linkables);
+
+    // When we visit the tree.
+    visit(node, visitor);
+
+    // Then we expect each program to have its own account.
+    expect(dictionary.programA).toBe(programA.accounts[0]);
+    expect(dictionary.programB).toBe(programB.accounts[0]);
+});
+
+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.
+    const linkables = new LinkableDictionary();
+    const visitor = recordLinkablesVisitor(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);
+});

packages/visitors/src/unwrapDefinedTypesVisitor.ts

@@ -16,7 +16,6 @@ export function unwrapDefinedTypesVisitor(typesToInline: string[] | '*' = '*') {
 
     return pipe(
         nonNullableIdentityVisitor(),
-        v => recordLinkablesVisitor(v, linkables),
         v =>
             extendVisitor(v, {
                 visitDefinedTypeLink(linkType, { self }) {
@@ -42,5 +41,6 @@ export function unwrapDefinedTypesVisitor(typesToInline: string[] | '*' = '*') {
                     });
                 },
             }),
+        v => recordLinkablesVisitor(v, linkables),
     );
 }

packages/visitors/test/fillDefaultPdaSeedValuesVisitor.test.ts

@@ -10,6 +10,7 @@ import {
     pdaNode,
     pdaSeedValueNode,
     pdaValueNode,
+    programNode,
     publicKeyTypeNode,
     variablePdaSeedNode,
 } from '@kinobi-so/nodes';
@@ -20,17 +21,25 @@ import { fillDefaultPdaSeedValuesVisitor } from '../src';
 
 test('it fills missing pda seed values with default values', () => {
     // Given a pdaNode with three variable seeds.
-    const pda = pdaNode({
-        name: 'myPda',
-        seeds: [
-            variablePdaSeedNode('seed1', numberTypeNode('u64')),
-            variablePdaSeedNode('seed2', numberTypeNode('u64')),
-            variablePdaSeedNode('seed3', publicKeyTypeNode()),
+    const program = programNode({
+        name: 'myProgram',
+        pdas: [
+            pdaNode({
+                name: 'myPda',
+                seeds: [
+                    variablePdaSeedNode('seed1', numberTypeNode('u64')),
+                    variablePdaSeedNode('seed2', numberTypeNode('u64')),
+                    variablePdaSeedNode('seed3', publicKeyTypeNode()),
+                ],
+            }),
         ],
+        publicKey: '1111',
     });
+    const pda = program.pdas[0];
 
     // And a linkable dictionary that recorded this PDA.
     const linkables = new LinkableDictionary();
+    linkables.stack.push(program);
     linkables.record(pda);
 
     // And a pdaValueNode with a single seed filled.
@@ -64,17 +73,25 @@ test('it fills missing pda seed values with default values', () => {
 
 test('it fills nested pda value nodes', () => {
     // Given a pdaNode with three variable seeds.
-    const pda = pdaNode({
-        name: 'myPda',
-        seeds: [
-            variablePdaSeedNode('seed1', numberTypeNode('u64')),
-            variablePdaSeedNode('seed2', numberTypeNode('u64')),
-            variablePdaSeedNode('seed3', publicKeyTypeNode()),
+    const program = programNode({
+        name: 'myProgram',
+        pdas: [
+            pdaNode({
+                name: 'myPda',
+                seeds: [
+                    variablePdaSeedNode('seed1', numberTypeNode('u64')),
+                    variablePdaSeedNode('seed2', numberTypeNode('u64')),
+                    variablePdaSeedNode('seed3', publicKeyTypeNode()),
+                ],
+            }),
         ],
+        publicKey: '1111',
     });
+    const pda = program.pdas[0];
 
     // And a linkable dictionary that recorded this PDA.
     const linkables = new LinkableDictionary();
+    linkables.stack.push(program);
     linkables.record(pda);
 
     // And a pdaValueNode nested inside a conditionalValueNode.
@@ -114,17 +131,25 @@ test('it fills nested pda value nodes', () => {
 
 test('it ignores default seeds missing from the instruction', () => {
     // Given a pdaNode with three variable seeds.
-    const pda = pdaNode({
-        name: 'myPda',
-        seeds: [
-            variablePdaSeedNode('seed1', numberTypeNode('u64')),
-            variablePdaSeedNode('seed2', numberTypeNode('u64')),
-            variablePdaSeedNode('seed3', publicKeyTypeNode()),
+    const program = programNode({
+        name: 'myProgram',
+        pdas: [
+            pdaNode({
+                name: 'myPda',
+                seeds: [
+                    variablePdaSeedNode('seed1', numberTypeNode('u64')),
+                    variablePdaSeedNode('seed2', numberTypeNode('u64')),
+                    variablePdaSeedNode('seed3', publicKeyTypeNode()),
+                ],
+            }),
         ],
+        publicKey: '1111',
     });
+    const pda = program.pdas[0];
 
     // And a linkable dictionary that recorded this PDA.
     const linkables = new LinkableDictionary();
+    linkables.stack.push(program);
     linkables.record(pda);
 
     // And a pdaValueNode with a single seed filled.