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,33 +103,45 @@ 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;
     }
 
     has(linkNode: LinkNode): boolean {
         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;
     }
 }