From b4009526e5efbb2586627309893811141c7da2bf Mon Sep 17 00:00:00 2001
From: Dave Pagurek <davepagurek@gmail.com>
Date: Thu, 30 May 2024 15:23:33 -0400
Subject: [PATCH] Visit replaced nodes

---
 src/ast/ast.test.ts | 94 +++++++++++++++++++++++++++++++++++++++++++++
 src/ast/visit.ts    | 37 +++++++++---------
 2 files changed, 113 insertions(+), 18 deletions(-)

diff --git a/src/ast/ast.test.ts b/src/ast/ast.test.ts
index c99e82e..dca1011 100644
--- a/src/ast/ast.test.ts
+++ b/src/ast/ast.test.ts
@@ -67,3 +67,97 @@ test('visit()', () => {
   expect(parent?.type).toBe('binary');
   expect(unfound).not.toBeDefined();
 });
+
+test('visit with replace', () => {
+  const visitLog: Array<['enter' | 'exit', AstNode['type']]> = [];
+
+  const tree: BinaryNode = {
+    type: 'binary',
+    operator: literal('-'),
+    // mock location data
+    location: {
+      start: { line: 0, column: 0, offset: 0 },
+      end: { line: 0, column: 0, offset: 0 },
+    },
+    left: identifier('foo'),
+    right: {
+      type: 'group',
+      lp: literal('('),
+      rp: literal(')'),
+      expression: identifier('bar'),
+    },
+  };
+
+  let sawBar = false;
+  let sawBaz = false;
+
+  visit(tree, {
+    group: {
+      enter: (path) => {
+        visitLog.push(['enter', path.node.type]);
+        path.replaceWith(identifier('baz'));
+      },
+      exit: (path) => {
+        visitLog.push(['exit', path.node.type]);
+      },
+    },
+    binary: {
+      enter: (path) => {
+        visitLog.push(['enter', path.node.type]);
+      },
+      exit: (path) => {
+        visitLog.push(['exit', path.node.type]);
+      },
+    },
+    literal: {
+      enter: (path) => {
+        visitLog.push(['enter', path.node.type]);
+      },
+      exit: (path) => {
+        visitLog.push(['exit', path.node.type]);
+      },
+    },
+    identifier: {
+      enter: (path) => {
+        visitLog.push(['enter', path.node.type]);
+        if (path.node.identifier === 'baz') {
+          sawBaz = true;
+        }
+        if (path.node.identifier === 'bar') {
+          sawBar = true;
+        }
+      },
+      exit: (path) => {
+        visitLog.push(['exit', path.node.type]);
+      },
+    },
+  });
+
+  expect(visitLog).toEqual([
+    ['enter', 'binary'],
+
+    // tree.operator
+    ['enter', 'literal'],
+    ['exit', 'literal'],
+
+    // tree.left
+    ['enter', 'identifier'],
+    ['exit', 'identifier'],
+
+    // tree.right
+    ['enter', 'group'],
+    // No exit because it got replaced
+
+    // Replaced tree.right
+    ['enter', 'identifier'],
+    ['exit', 'identifier'],
+
+    ['exit', 'binary'],
+  ]);
+
+  // The children of the node that got replaced should not be visited
+  expect(sawBar).toBeFalsy();
+
+  // The children of the new replacement node should be visited
+  expect(sawBaz).toBeTruthy();
+})
diff --git a/src/ast/visit.ts b/src/ast/visit.ts
index e799550..ac3721b 100644
--- a/src/ast/visit.ts
+++ b/src/ast/visit.ts
@@ -115,27 +115,28 @@ export const visit = (ast: Program | AstNode, visitors: NodeVisitors) => {
       }
     }
 
-    const toTraverse = (path.replaced as AstNode | undefined) ?? node;
-    const newPath = path.replaced
-      ? makePath(toTraverse, parent, parentPath, key, index)
-      : path;
-    Object.entries(toTraverse)
-      .filter(([_, nodeValue]) => isTraversable(nodeValue))
-      .forEach(([nodeKey, nodeValue]) => {
-        if (Array.isArray(nodeValue)) {
-          for (let i = 0, offset = 0; i - offset < nodeValue.length; i++) {
-            const child = nodeValue[i - offset];
-            const res = visitNode(child, toTraverse, newPath, nodeKey, i - offset);
-            if (res?.removed) {
-              offset += 1;
+    if (path.replaced) {
+      const replacedNode = path.replaced as AstNode;
+      visitNode(replacedNode, parent, parentPath, key, index);
+    } else {
+      Object.entries(node)
+        .filter(([_, nodeValue]) => isTraversable(nodeValue))
+        .forEach(([nodeKey, nodeValue]) => {
+          if (Array.isArray(nodeValue)) {
+            for (let i = 0, offset = 0; i - offset < nodeValue.length; 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);
           }
-        } else {
-          visitNode(nodeValue, toTraverse, newPath, nodeKey);
-        }
-      });
+        });
 
-    visitor?.exit?.(path as any);
+      visitor?.exit?.(path as any);
+    }
   };
 
   visitNode(ast);