Adding path.stop()

README.md

@@ -355,6 +355,8 @@ node itself. The path object:
 
   // Don't visit any children of this node
   skip: () => void;
+  // Stop traversal entirely
+  stop: () => void;
   // Remove this node from the AST
   remove: () => void;
   // Replace this node with another AST node. See replaceWith() documentation.
@@ -401,6 +403,13 @@ const ast = parser.parse(`float a = 1.0;`);
 visitPreprocessedAst(ast, visitors);
 ```
 
+### Stopping traversal
+
+To skip all children of a node, call `path.skip()`.
+
+To stop traversal entirely, call `path.stop()` in either `enter()` or `exit()`.
+No future `enter()` nor `exit()` callbacks will fire.
+
 ### Visitor `.replaceWith()` Behavior
 
 When you visit a node and call `path.replaceWith(otherNode)` inside the visitor's `enter()` method:

package.json

@@ -3,7 +3,7 @@
   "engines": {
     "node": ">=16"
   },
-  "version": "4.0.0",
+  "version": "4.1.0",
   "type": "module",
   "description": "A GLSL ES 1.0 and 3.0 parser and preprocessor that can preserve whitespace and comments",
   "scripts": {

src/ast/ast.test.ts

@@ -4,7 +4,16 @@ import {
   IdentifierNode,
   LiteralNode,
 } from './ast-types.js';
-import { visit } from './visit.js';
+import { Path, visit } from './visit.js';
+
+const visitLogger = () => {
+  const visitLog: Array<['enter' | 'exit', AstNode['type']]> = [];
+  const track = (type: 'enter' | 'exit') => (path: Path<any>) =>
+    visitLog.push([type, path.node.type]);
+  const enter = track('enter');
+  const exit = track('exit');
+  return [visitLog, enter, exit, track] as const;
+};
 
 const literal = <T>(literal: T): LiteralNode<T> => ({
   type: 'literal',
@@ -69,7 +78,7 @@ test('visit()', () => {
 });
 
 test('visit with replace', () => {
-  const visitLog: Array<['enter' | 'exit', AstNode['type']]> = [];
+  const [visitLog, enter, exit] = visitLogger();
 
   const tree: BinaryNode = {
     type: 'binary',
@@ -94,42 +103,30 @@ test('visit with replace', () => {
   visit(tree, {
     group: {
       enter: (path) => {
-        visitLog.push(['enter', path.node.type]);
+        enter(path);
         path.replaceWith(identifier('baz'));
       },
-      exit: (path) => {
-        visitLog.push(['exit', path.node.type]);
-      },
+      exit,
     },
     binary: {
-      enter: (path) => {
-        visitLog.push(['enter', path.node.type]);
-      },
-      exit: (path) => {
-        visitLog.push(['exit', path.node.type]);
-      },
+      enter,
+      exit,
     },
     literal: {
-      enter: (path) => {
-        visitLog.push(['enter', path.node.type]);
-      },
-      exit: (path) => {
-        visitLog.push(['exit', path.node.type]);
-      },
+      enter,
+      exit,
     },
     identifier: {
       enter: (path) => {
-        visitLog.push(['enter', path.node.type]);
+        enter(path);
         if (path.node.identifier === 'baz') {
           sawBaz = true;
         }
         if (path.node.identifier === 'bar') {
           sawBar = true;
         }
       },
-      exit: (path) => {
-        visitLog.push(['exit', path.node.type]);
-      },
+      exit,
     },
   });
 
@@ -160,4 +157,65 @@ test('visit with replace', () => {
 
   // The children of the new replacement node should be visited
   expect(sawBaz).toBeTruthy();
-})
+});
+
+test('visit stop()', () => {
+  const [visitLog, enter, exit] = visitLogger();
+
+  const tree: BinaryNode = {
+    type: 'binary',
+    operator: literal('-'),
+    left: {
+      type: 'binary',
+      operator: literal('+'),
+      left: identifier('foo'),
+      right: identifier('bar'),
+    },
+    right: {
+      type: 'group',
+      lp: literal('('),
+      rp: literal(')'),
+      expression: identifier('baz'),
+    },
+  };
+
+  visit(tree, {
+    group: {
+      enter,
+      exit,
+    },
+    binary: {
+      enter,
+      exit,
+    },
+    literal: {
+      enter,
+      exit,
+    },
+    identifier: {
+      enter: (path) => {
+        enter(path);
+        if (path.node.identifier === 'foo') {
+          path.stop();
+        }
+      },
+      exit,
+    },
+  });
+
+  expect(visitLog).toEqual([
+    ['enter', 'binary'],
+
+    // tree.operator
+    ['enter', 'literal'],
+    ['exit', 'literal'],
+
+    // tree.left
+    ['enter', 'binary'],
+    ['enter', 'literal'],
+    ['exit', 'literal'],
+
+    // stop on first identifier!
+    ['enter', 'identifier'],
+  ]);
+});

src/ast/visit.ts

@@ -9,11 +9,13 @@ export type Path<NodeType> = {
   parentPath: Path<any> | undefined;
   key: string | undefined;
   index: number | undefined;
+  stop: () => void;
   skip: () => void;
   remove: () => void;
   replaceWith: (replacer: AstNode) => void;
   findParent: (test: (p: Path<any>) => boolean) => Path<any> | undefined;
 
+  stopped?: boolean;
   skipped?: boolean;
   removed?: boolean;
   replaced?: any;
@@ -31,6 +33,9 @@ const makePath = <NodeType>(
   parentPath,
   key,
   index,
+  stop: function () {
+    this.stopped = true;
+  },
   skip: function () {
     this.skipped = true;
   },
@@ -72,13 +77,20 @@ export type NodeVisitors = {
  * Apply the visitor pattern to an AST that conforms to this compiler's spec
  */
 export const visit = (ast: Program | AstNode, visitors: NodeVisitors) => {
+  let stopped = false;
+
   const visitNode = (
     node: AstNode | Program,
     parent?: AstNode | Program,
     parentPath?: Path<any>,
     key?: string,
     index?: number
   ) => {
+    // Handle case where stop happened at exit
+    if (stopped) {
+      return;
+    }
+
     const visitor = visitors[node.type];
     const path = makePath(node, parent, parentPath, key, index);
     const parentNode = parent as any;
@@ -115,6 +127,11 @@ export const visit = (ast: Program | AstNode, visitors: NodeVisitors) => {
       }
     }
 
+    if (path.stopped) {
+      stopped = true;
+      return;
+    }
+
     if (path.replaced) {
       const replacedNode = path.replaced as AstNode;
       visitNode(replacedNode, parent, parentPath, key, index);
@@ -123,19 +140,27 @@ export const visit = (ast: Program | AstNode, visitors: NodeVisitors) => {
         .filter(([_, nodeValue]) => isTraversable(nodeValue))
         .forEach(([nodeKey, nodeValue]) => {
           if (Array.isArray(nodeValue)) {
-            for (let i = 0, offset = 0; i - offset < nodeValue.length; i++) {
+            for (
+              let i = 0, offset = 0;
+              i - offset < nodeValue.length && !stopped;
+              i++
+            ) {
               const child = nodeValue[i - offset];
               const res = visitNode(child, node, path, nodeKey, i - offset);
               if (res?.removed) {
                 offset += 1;
               }
             }
           } else {
-            visitNode(nodeValue, node, path, nodeKey);
+            if (!stopped) {
+              visitNode(nodeValue, node, path, nodeKey);
+            }
           }
         });
 
-      visitor?.exit?.(path as any);
+      if (!stopped) {
+        visitor?.exit?.(path as any);
+      }
     }
   };