Merge pull request #25 from davepagurek/fix/traverse-replaced

Only traverse replaced node after calling .replace()

src/ast/visit.ts

@@ -115,19 +115,23 @@ export const visit = (ast: Program | AstNode, visitors: NodeVisitors) => {
       }
     }
 
-    Object.entries(node)
+    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, node, path, nodeKey, i - offset);
+            const res = visitNode(child, toTraverse, newPath, nodeKey, i - offset);
             if (res?.removed) {
               offset += 1;
             }
           }
         } else {
-          visitNode(nodeValue, node, path, nodeKey);
+          visitNode(nodeValue, toTraverse, newPath, nodeKey);
         }
       });
 

src/preprocessor/preprocessor.test.ts

@@ -117,6 +117,39 @@ true
 `);
 });
 
+test('define inside if/else is properly expanded when the if branch is chosen', () => {
+  const program = `
+#define MACRO
+#ifdef MACRO
+#define BRANCH a
+#else
+#define BRANCH b
+#endif
+BRANCH
+`;
+  const ast = parse(program);
+  preprocessAst(ast);
+  expect(generate(ast)).toBe(`
+a
+`);
+});
+
+test('define inside if/else is properly expanded when the else branch is chosen', () => {
+  const program = `
+#ifdef MACRO
+#define BRANCH a
+#else
+#define BRANCH b
+#endif
+BRANCH
+`;
+  const ast = parse(program);
+  preprocessAst(ast);
+  expect(generate(ast)).toBe(`
+b
+`);
+});
+
 test('ifdef inside else is properly expanded', () => {
   // Regression: Make sure #ifdef MACRO inside #else isn't expanded
   const program = `