Record linkables before the first visit (#173)

.changeset/clever-news-worry.md

@@ -0,0 +1,5 @@
+---
+'@kinobi-so/visitors-core': patch
+---
+
+Record linkables before the first visit

packages/visitors-core/src/LinkableDictionary.ts

@@ -14,6 +14,8 @@ import {
 
 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();
 

packages/visitors-core/src/recordLinkablesVisitor.ts

@@ -1,51 +1,24 @@
-import { getAllAccounts, getAllDefinedTypes, getAllPdas, getAllPrograms, NodeKind } from '@kinobi-so/nodes';
+import { isNode, type NodeKind } from '@kinobi-so/nodes';
 
-import { extendVisitor, VisitorOverrides } from './extendVisitor';
-import { LinkableDictionary } from './LinkableDictionary';
-import { Visitor } from './visitor';
+import { interceptFirstVisitVisitor } from './interceptFirstVisitVisitor';
+import { interceptVisitor } from './interceptVisitor';
+import { LINKABLE_NODES, LinkableDictionary } from './LinkableDictionary';
+import { visit, Visitor } from './visitor';
+import { voidVisitor } from './voidVisitor';
 
 export function recordLinkablesVisitor<TReturn, TNodeKind extends NodeKind>(
     visitor: Visitor<TReturn, TNodeKind>,
     linkables: LinkableDictionary,
 ): Visitor<TReturn, TNodeKind> {
-    const overriddenFunctions: VisitorOverrides<
-        TReturn,
-        'accountNode' | 'definedTypeNode' | 'pdaNode' | 'programNode' | 'rootNode'
-    > = {};
-    if ('visitRoot' in visitor) {
-        overriddenFunctions.visitRoot = function visitRoot(node, { next }) {
-            linkables.recordAll([
-                ...getAllPrograms(node),
-                ...getAllPdas(node),
-                ...getAllAccounts(node),
-                ...getAllDefinedTypes(node),
-            ]);
-            return next(node);
-        };
-    }
-    if ('visitProgram' in visitor) {
-        overriddenFunctions.visitProgram = function visitProgram(node, { next }) {
-            linkables.recordAll([node, ...node.pdas, ...node.accounts, ...node.definedTypes]);
-            return next(node);
-        };
-    }
-    if ('visitPda' in visitor) {
-        overriddenFunctions.visitPda = function visitPda(node, { next }) {
+    const recordingVisitor = interceptVisitor(voidVisitor(), (node, next) => {
+        if (isNode(node, LINKABLE_NODES)) {
             linkables.record(node);
-            return next(node);
-        };
-    }
-    if ('visitAccount' in visitor) {
-        overriddenFunctions.visitAccount = function visitAccount(node, { next }) {
-            linkables.record(node);
-            return next(node);
-        };
-    }
-    if ('visitDefinedType' in visitor) {
-        overriddenFunctions.visitDefinedType = function visitDefinedType(node, { next }) {
-            linkables.record(node);
-            return next(node);
-        };
-    }
-    return extendVisitor(visitor, overriddenFunctions as VisitorOverrides<TReturn, TNodeKind>);
+        }
+        return next(node);
+    });
+
+    return interceptFirstVisitVisitor(visitor, (node, next) => {
+        visit(node, recordingVisitor);
+        return next(node);
+    });
 }

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

@@ -12,7 +12,7 @@ import {
 } from '@kinobi-so/nodes';
 import { expect, test } from 'vitest';
 
-import { LinkableDictionary, recordLinkablesVisitor, visit, voidVisitor } from '../src';
+import { interceptFirstVisitVisitor, 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.
@@ -52,3 +52,27 @@ test('it record all linkable nodes it finds when traversing the tree', () => {
     expect(linkables.get(definedTypeLinkNode('typeA'))).toEqual(node.program.definedTypes[0]);
     expect(linkables.get(definedTypeLinkNode('typeB'))).toEqual(node.additionalPrograms[0].definedTypes[0]);
 });
+
+test('it records all linkable before the first visit of the base visitor', () => {
+    // Given the following root node with two programs.
+    const node = rootNode(programNode({ name: 'programA', publicKey: '1111' }), [
+        programNode({ name: 'programB', publicKey: '2222' }),
+    ]);
+
+    // And a recordLinkablesVisitor extending a base visitor that
+    // stores the linkable programs available at every visit.
+    const linkables = new LinkableDictionary();
+    const events: string[] = [];
+    const baseVisitor = interceptFirstVisitVisitor(voidVisitor(), (node, next) => {
+        events.push(`programA:${linkables.has(programLinkNode('programA'))}`);
+        events.push(`programB:${linkables.has(programLinkNode('programB'))}`);
+        next(node);
+    });
+    const visitor = recordLinkablesVisitor(baseVisitor, linkables);
+
+    // When we visit the tree.
+    visit(node, visitor);
+
+    // Then we expect all linkable nodes to be recorded.
+    expect(events).toEqual(['programA:true', 'programB:true']);
+});