Add interceptFirstVisitVisitor helper (#172)

.changeset/tasty-laws-admire.md

@@ -0,0 +1,5 @@
+---
+'@kinobi-so/visitors-core': patch
+---
+
+Add `interceptFirstVisitVisitor` helper

packages/visitors-core/README.md

@@ -286,6 +286,27 @@ visit(tupleTypeNode([numberTypeNode('u32'), publicKeyTypeNode()]), visitor);
 // ]
 ```
 
+### `interceptFirstVisitVisitor`
+
+The `interceptFirstVisitVisitor` works the same way as the `interceptVisitor` but only intercepts the first visit of a node. This means that the provided function is called when visiting the specific node provided but not when visiting its children. The parameters are the same as for the `interceptVisitor`.
+
+For instance, the following visitor intercepts a `voidVisitor` and captures events only during the first visit.
+
+```ts
+const events: string[] = [];
+const visitor = interceptFirstVisitVisitor(voidVisitor(), (node, next) => {
+    events.push(`down:${node.kind}`);
+    next(node);
+    events.push(`up:${node.kind}`);
+});
+
+visit(tupleTypeNode([numberTypeNode('u32'), publicKeyTypeNode()]), visitor);
+// events === [
+//     'down:tupleTypeNode',
+//     'up:tupleTypeNode',
+// ]
+```
+
 ### `tapVisitor`
 
 The `tapVisitor` function allows us to tap into the visiting functions of a provided visitor without modifying its behaviour. This means the returned visitor will behave exactly like the base visitor except that the provided function will be called for the specified node kind.

packages/visitors-core/src/index.ts

@@ -11,6 +11,7 @@ export * from './getResolvedInstructionInputsVisitor';
 export * from './getUniqueHashStringVisitor';
 export * from './identityVisitor';
 export * from './interceptVisitor';
+export * from './interceptFirstVisitVisitor';
 export * from './mapVisitor';
 export * from './mergeVisitor';
 export * from './nonNullableIdentityVisitor';

packages/visitors-core/src/interceptFirstVisitVisitor.ts

@@ -0,0 +1,26 @@
+import type { NodeKind } from '@kinobi-so/nodes';
+
+import { interceptVisitor, VisitorInterceptor } from './interceptVisitor';
+import { Visitor } from './visitor';
+
+export function interceptFirstVisitVisitor<TReturn, TNodeKind extends NodeKind>(
+    visitor: Visitor<TReturn, TNodeKind>,
+    interceptor: VisitorInterceptor<TReturn>,
+): Visitor<TReturn, TNodeKind> {
+    let isFirstVisit = true;
+
+    return interceptVisitor(visitor, (node, next) => {
+        try {
+            if (isFirstVisit) {
+                isFirstVisit = false;
+                const result = interceptor(node, next);
+                isFirstVisit = true;
+                return result;
+            }
+            return next(node);
+        } catch (error) {
+            isFirstVisit = true;
+            throw error;
+        }
+    });
+}

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

@@ -0,0 +1,72 @@
+import { numberTypeNode, publicKeyTypeNode, tupleTypeNode } from '@kinobi-so/nodes';
+import { expect, test } from 'vitest';
+
+import { extendVisitor, interceptFirstVisitVisitor, visit, voidVisitor } from '../src';
+
+test('it returns a new visitor that only intercepts the first visit of a visitor', () => {
+    // Given the following 3-nodes tree.
+    const node = tupleTypeNode([numberTypeNode('u32'), publicKeyTypeNode()]);
+
+    // And an intercepted void visitor that records the events that happened during the first visit.
+    const events: string[] = [];
+    const baseVisitor = voidVisitor();
+    const visitor = interceptFirstVisitVisitor(baseVisitor, (node, next) => {
+        events.push(`down:${node.kind}`);
+        next(node);
+        events.push(`up:${node.kind}`);
+    });
+
+    // When we visit the tree using that visitor.
+    visit(node, visitor);
+
+    // Then we expect the following events to have happened.
+    expect(events).toEqual(['down:tupleTypeNode', 'up:tupleTypeNode']);
+
+    // And the intercepted visitor is a new instance.
+    expect(baseVisitor).not.toBe(visitor);
+});
+
+test('it still works on subsequent calls', () => {
+    // Given the following 3-nodes tree.
+    const node = tupleTypeNode([numberTypeNode('u32'), publicKeyTypeNode()]);
+
+    // And an intercepted void visitor that records the events that happened during the first visit.
+    const events: string[] = [];
+    const baseVisitor = voidVisitor();
+    const visitor = interceptFirstVisitVisitor(baseVisitor, (node, next) => {
+        events.push(`intercepting:${node.kind}`);
+        next(node);
+    });
+
+    // When we visit the tree twice using that visitor.
+    visit(node, visitor);
+    visit(node, visitor);
+
+    // Then we expect the following events to have happened.
+    expect(events).toEqual(['intercepting:tupleTypeNode', 'intercepting:tupleTypeNode']);
+});
+
+test('it resets the first visit boolean if an error is thrown', () => {
+    // Given the following 3-nodes tree.
+    const node = tupleTypeNode([numberTypeNode('u32'), publicKeyTypeNode()]);
+
+    // And an intercepted visitor that records the events that happened during the first visit
+    // but throws an error when it visits publicKeyTypeNodes.
+    const events: string[] = [];
+    const baseVisitor = extendVisitor(voidVisitor(), {
+        visitPublicKeyType: () => {
+            throw new Error('public key error');
+        },
+    });
+    const visitor = interceptFirstVisitVisitor(baseVisitor, (node, next) => {
+        events.push(`intercepting:${node.kind}`);
+        next(node);
+    });
+
+    // Then we expect errors to be thrown whenever we visit the three.
+    expect(() => visit(node, visitor)).toThrow('public key error');
+    expect(() => visit(node, visitor)).toThrow('public key error');
+
+    // But we still expect the following events to have happened.
+    expect(events).toEqual(['intercepting:tupleTypeNode', 'intercepting:tupleTypeNode']);
+});