Use NodePaths in NodeSelectors

Author: lorisleiva

packages/visitors-core/src/NodePath.ts

@@ -2,7 +2,7 @@ import { assertIsNode, GetNodeFromKind, InstructionNode, isNode, Node, NodeKind,
 
 export type NodePath<TNode extends Node | undefined = undefined> = TNode extends undefined
     ? readonly Node[]
-    : readonly [...Node[], TNode];
+    : readonly [...(readonly Node[]), TNode];
 
 export function getLastNodeFromPath<TNode extends Node>(path: NodePath<TNode>): TNode {
     return path[path.length - 1] as TNode;
@@ -49,22 +49,22 @@ export function getNodePathUntilLastNode<TKind extends NodeKind>(
     return path.slice(0, lastIndex + 1) as unknown as NodePath<GetNodeFromKind<TKind>>;
 }
 
-function isNotEmptyNodePath(path: NodePath | null | undefined): path is NodePath<Node> {
+export function isFilledNodePath(path: NodePath | null | undefined): path is NodePath<Node> {
     return !!path && path.length > 0;
 }
 
 export function isNodePath<TKind extends NodeKind>(
     path: NodePath | null | undefined,
     kind: TKind | TKind[],
 ): path is NodePath<GetNodeFromKind<TKind>> {
-    return isNode(isNotEmptyNodePath(path) ? getLastNodeFromPath<Node>(path) : null, kind);
+    return isNode(isFilledNodePath(path) ? getLastNodeFromPath<Node>(path) : null, kind);
 }
 
 export function assertIsNodePath<TKind extends NodeKind>(
     path: NodePath | null | undefined,
     kind: TKind | TKind[],
 ): asserts path is NodePath<GetNodeFromKind<TKind>> {
-    assertIsNode(isNotEmptyNodePath(path) ? getLastNodeFromPath<Node>(path) : null, kind);
+    assertIsNode(isFilledNodePath(path) ? getLastNodeFromPath<Node>(path) : null, kind);
 }
 
 export function nodePathToStringArray(path: NodePath): string[] {

packages/visitors-core/src/NodeSelector.ts

@@ -1,6 +1,6 @@
 import { camelCase, CamelCaseString, Node } from '@codama/nodes';
 
-import type { NodeStack } from './NodeStack';
+import { NodePath } from './NodePath';
 
 export type NodeSelector = NodeSelectorFunction | NodeSelectorPath;
 
@@ -11,11 +11,11 @@ export type NodeSelector = NodeSelectorFunction | NodeSelectorPath;
  * - `[someNode]` matches a node of the given kind.
  * - `[someNode|someOtherNode]` matches a node with any of the given kind.
  * - `[someNode]someText` matches both the kind and the name of a node.
- * - `a.b.c` matches a node `c` such that its parent stack contains `a` and `b` in order (but not necessarily subsequent).
+ * - `a.b.c` matches a node `c` such that its ancestors contains `a` and `b` in order (but not necessarily subsequent).
  */
 export type NodeSelectorPath = string;
 
-export type NodeSelectorFunction = (node: Node, stack: NodeStack) => boolean;
+export type NodeSelectorFunction = (path: NodePath) => boolean;
 
 export const getNodeSelectorFunction = (selector: NodeSelector): NodeSelectorFunction => {
     if (typeof selector === 'function') return selector;
@@ -40,25 +40,29 @@ export const getNodeSelectorFunction = (selector: NodeSelector): NodeSelectorFun
         return true;
     };
 
-    const checkStack = (nodeStack: Node[], nodeSelectors: string[]): boolean => {
+    const checkPath = (path: Node[], nodeSelectors: string[]): boolean => {
         if (nodeSelectors.length === 0) return true;
-        if (nodeStack.length === 0) return false;
-        const lastNode = nodeStack.pop() as Node;
+        if (path.length === 0) return false;
+        const lastNode = path.pop() as Node;
         const lastNodeSelector = nodeSelectors.pop() as string;
         return checkNode(lastNode, lastNodeSelector)
-            ? checkStack(nodeStack, nodeSelectors)
-            : checkStack(nodeStack, [...nodeSelectors, lastNodeSelector]);
+            ? checkPath(path, nodeSelectors)
+            : checkPath(path, [...nodeSelectors, lastNodeSelector]);
     };
 
-    const nodeSelectors = selector.split('.');
-    const lastNodeSelector = nodeSelectors.pop() as string;
+    const checkInitialPath = (path: Node[], nodeSelectors: string[]): boolean => {
+        if (nodeSelectors.length === 0 || path.length === 0) return false;
+        const lastNode = path.pop() as Node;
+        const lastNodeSelector = nodeSelectors.pop() as string;
+        return checkNode(lastNode, lastNodeSelector) && checkPath(path, nodeSelectors);
+    };
 
-    return (node, stack) =>
-        checkNode(node, lastNodeSelector) && checkStack(stack.getPath() as Node[], [...nodeSelectors]);
+    const nodeSelectors = selector.split('.');
+    return path => checkInitialPath([...path], [...nodeSelectors]);
 };
 
 export const getConjunctiveNodeSelectorFunction = (selector: NodeSelector | NodeSelector[]): NodeSelectorFunction => {
     const selectors = Array.isArray(selector) ? selector : [selector];
     const selectorFunctions = selectors.map(getNodeSelectorFunction);
-    return (node, stack) => selectorFunctions.every(fn => fn(node, stack));
+    return path => selectorFunctions.every(fn => fn(path));
 };

packages/visitors-core/src/bottomUpTransformerVisitor.ts

@@ -8,11 +8,11 @@ import { pipe } from './pipe';
 import { recordNodeStackVisitor } from './recordNodeStackVisitor';
 import { Visitor } from './visitor';
 
-export type BottomUpNodeTransformer<TNode extends Node = Node> = (node: TNode, stack: NodeStack) => Node | null;
+export type BottomUpNodeTransformer = (node: Node, stack: NodeStack) => Node | null;
 
-export type BottomUpNodeTransformerWithSelector<TNode extends Node = Node> = {
+export type BottomUpNodeTransformerWithSelector = {
     select: NodeSelector | NodeSelector[];
-    transform: BottomUpNodeTransformer<TNode>;
+    transform: BottomUpNodeTransformer;
 };
 
 export function bottomUpTransformerVisitor<TNodeKind extends NodeKind = NodeKind>(
@@ -22,21 +22,21 @@ export function bottomUpTransformerVisitor<TNodeKind extends NodeKind = NodeKind
     const transformerFunctions = transformers.map((transformer): BottomUpNodeTransformer => {
         if (typeof transformer === 'function') return transformer;
         return (node, stack) =>
-            getConjunctiveNodeSelectorFunction(transformer.select)(node, stack)
+            getConjunctiveNodeSelectorFunction(transformer.select)(stack.getPath())
                 ? transformer.transform(node, stack)
                 : node;
     });
 
     const stack = new NodeStack();
     return pipe(
         identityVisitor(nodeKeys),
-        v => recordNodeStackVisitor(v, stack),
         v =>
-            interceptVisitor(v, (node, next) =>
-                transformerFunctions.reduce(
-                    (acc, transformer) => (acc === null ? null : transformer(acc, stack.clone())),
+            interceptVisitor(v, (node, next) => {
+                return transformerFunctions.reduce(
+                    (acc, transformer) => (acc === null ? null : transformer(acc, stack)),
                     next(node),
-                ),
-            ),
+                );
+            }),
+        v => recordNodeStackVisitor(v, stack),
     );
 }

packages/visitors-core/src/topDownTransformerVisitor.ts

@@ -8,14 +8,11 @@ import { pipe } from './pipe';
 import { recordNodeStackVisitor } from './recordNodeStackVisitor';
 import { Visitor } from './visitor';
 
-export type TopDownNodeTransformer<TNode extends Node = Node> = <T extends TNode = TNode>(
-    node: T,
-    stack: NodeStack,
-) => T | null;
+export type TopDownNodeTransformer = <TNode extends Node>(node: TNode, stack: NodeStack) => TNode | null;
 
-export type TopDownNodeTransformerWithSelector<TNode extends Node = Node> = {
+export type TopDownNodeTransformerWithSelector = {
     select: NodeSelector | NodeSelector[];
-    transform: TopDownNodeTransformer<TNode>;
+    transform: TopDownNodeTransformer;
 };
 
 export function topDownTransformerVisitor<TNodeKind extends NodeKind = NodeKind>(
@@ -25,23 +22,23 @@ export function topDownTransformerVisitor<TNodeKind extends NodeKind = NodeKind>
     const transformerFunctions = transformers.map((transformer): TopDownNodeTransformer => {
         if (typeof transformer === 'function') return transformer;
         return (node, stack) =>
-            getConjunctiveNodeSelectorFunction(transformer.select)(node, stack)
+            getConjunctiveNodeSelectorFunction(transformer.select)(stack.getPath())
                 ? transformer.transform(node, stack)
                 : node;
     });
 
     const stack = new NodeStack();
     return pipe(
         identityVisitor(nodeKeys),
-        v => recordNodeStackVisitor(v, stack),
         v =>
             interceptVisitor(v, (node, next) => {
                 const appliedNode = transformerFunctions.reduce(
-                    (acc, transformer) => (acc === null ? null : transformer(acc, stack.clone())),
+                    (acc, transformer) => (acc === null ? null : transformer(acc, stack)),
                     node as Parameters<typeof next>[0] | null,
                 );
                 if (appliedNode === null) return null;
                 return next(appliedNode);
             }),
+        v => recordNodeStackVisitor(v, stack),
     );
 }

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

@@ -10,7 +10,6 @@ import {
     instructionAccountNode,
     instructionArgumentNode,
     instructionNode,
-    isNode,
     Node,
     numberTypeNode,
     optionTypeNode,
@@ -23,9 +22,11 @@ import {
 import { expect, test } from 'vitest';
 
 import {
+    getLastNodeFromPath,
     getNodeSelectorFunction,
     identityVisitor,
     interceptVisitor,
+    isNodePath,
     NodeSelector,
     NodeStack,
     pipe,
@@ -196,12 +197,12 @@ const macro = (selector: NodeSelector, expectedSelected: Node[]) => {
         const selected = [] as Node[];
         const visitor = pipe(
             identityVisitor(),
-            v => recordNodeStackVisitor(v, stack),
             v =>
                 interceptVisitor(v, (node, next) => {
-                    if (selectorFunction(node, stack.clone())) selected.push(node);
+                    if (selectorFunction(stack.getPath())) selected.push(node);
                     return next(node);
                 }),
+            v => recordNodeStackVisitor(v, stack),
         );
 
         // When we visit the tree.
@@ -329,4 +330,7 @@ macro('[accountNode]gift.[publicKeyTypeNode|booleanTypeNode]', [
 ]);
 
 // Select using functions.
-macro(node => isNode(node, 'numberTypeNode') && node.format === 'u32', [tokenDelegatedAmountOption.prefix]);
+macro(
+    path => isNodePath(path, 'numberTypeNode') && getLastNodeFromPath(path).format === 'u32',
+    [tokenDelegatedAmountOption.prefix],
+);

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

@@ -94,7 +94,7 @@ test('it can transform nodes using multiple node selectors', () => {
     // - the second one selects all nodes with more than one ancestor.
     const visitor = bottomUpTransformerVisitor([
         {
-            select: ['[numberTypeNode]', (_, nodeStack) => nodeStack.getPath().length > 1],
+            select: ['[numberTypeNode]', path => path.length > 2],
             transform: () => stringTypeNode('utf8'),
         },
     ]);

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

@@ -102,7 +102,7 @@ test('it can transform nodes using multiple node selectors', () => {
     // - the second one selects all nodes with more than one ancestor.
     const visitor = topDownTransformerVisitor([
         {
-            select: ['[numberTypeNode]', (_, nodeStack) => nodeStack.getPath().length > 1],
+            select: ['[numberTypeNode]', path => path.length > 2],
             transform: node => numberTypeNode('u64') as typeof node,
         },
     ]);

packages/visitors/src/setFixedAccountSizesVisitor.ts

@@ -1,6 +1,8 @@
-import { accountNode, assertIsNode, isNode } from '@codama/nodes';
+import { accountNode, assertIsNode } from '@codama/nodes';
 import {
     getByteSizeVisitor,
+    getLastNodeFromPath,
+    isNodePath,
     LinkableDictionary,
     NodeStack,
     pipe,
@@ -18,7 +20,7 @@ export function setFixedAccountSizesVisitor() {
     const visitor = topDownTransformerVisitor(
         [
             {
-                select: node => isNode(node, 'accountNode') && node.size === undefined,
+                select: path => isNodePath(path, 'accountNode') && getLastNodeFromPath(path).size === undefined,
                 transform: node => {
                     assertIsNode(node, 'accountNode');
                     const size = visit(node.data, byteSizeVisitor);

packages/visitors/src/unwrapTupleEnumWithSingleStructVisitor.ts

@@ -3,7 +3,6 @@ import {
     CamelCaseString,
     DefinedTypeNode,
     enumStructVariantTypeNode,
-    EnumTupleVariantTypeNode,
     getAllDefinedTypes,
     isNode,
     resolveNestedTypeNode,
@@ -28,8 +27,7 @@ export function unwrapTupleEnumWithSingleStructVisitor(enumsOrVariantsToUnwrap:
             ? [() => true]
             : enumsOrVariantsToUnwrap.map(selector => getNodeSelectorFunction(selector));
 
-    const shouldUnwrap = (node: EnumTupleVariantTypeNode, stack: NodeStack): boolean =>
-        selectorFunctions.some(selector => selector(node, stack));
+    const shouldUnwrap = (stack: NodeStack): boolean => selectorFunctions.some(selector => selector(stack.getPath()));
 
     return rootNodeVisitor(root => {
         const typesToPotentiallyUnwrap: string[] = [];
@@ -44,7 +42,7 @@ export function unwrapTupleEnumWithSingleStructVisitor(enumsOrVariantsToUnwrap:
                     select: '[enumTupleVariantTypeNode]',
                     transform: (node, stack) => {
                         assertIsNode(node, 'enumTupleVariantTypeNode');
-                        if (!shouldUnwrap(node, stack)) return node;
+                        if (!shouldUnwrap(stack)) return node;
                         const tupleNode = resolveNestedTypeNode(node.tuple);
                         if (tupleNode.items.length !== 1) return node;
                         let item = tupleNode.items[0];