Improve traverse (#2782)

improve travers

src/lang/queryAst.test.ts

@@ -1,4 +1,4 @@
-import { parse, recast, initPromise } from './wasm'
+import { parse, recast, initPromise, PathToNode } from './wasm'
 import {
   findAllPreviousVariables,
   isNodeSafeToReplace,
@@ -9,6 +9,7 @@ import {
   findUsesOfTagInPipe,
   hasSketchPipeBeenExtruded,
   hasExtrudableGeometry,
+  traverse,
 } from './queryAst'
 import { enginelessExecutor } from '../lib/testHelpers'
 import {
@@ -538,3 +539,53 @@ const extrude001 = extrude(10, sketch001)
     expect(extrudable).toBeFalsy()
   })
 })
+
+describe.only('Testing traverse and pathToNode', () => {
+  it.each([
+    ['basic', '2.73'],
+    [
+      'very nested, array, object, callExpression, array, memberExpression',
+      '.yo',
+    ],
+  ])('testing %s', async (testName, literalOfInterest) => {
+    const code = `const myVar = 5
+const sketch001 = startSketchOn('XZ')
+  |> startProfileAt([3.29, 7.86], %)
+  |> line([2.48, 2.44], %)
+  |> line([-3.86, -2.73], %)
+  |> line([-17.67, 0.85], %)
+  |> close(%)
+const bing = { yo: 55 }
+const myNestedVar = [
+  {
+  prop:   line([bing.yo, 21], sketch001)
+}
+]
+  `
+    const ast = parse(code)
+    if (err(ast)) throw ast
+    let pathToNode: PathToNode = []
+    traverse(ast, {
+      enter: (node, path) => {
+        if (
+          node.type === 'Literal' &&
+          String(node.value) === literalOfInterest
+        ) {
+          pathToNode = path
+        } else if (
+          node.type === 'Identifier' &&
+          literalOfInterest.includes(node.name)
+        ) {
+          pathToNode = path
+        }
+      },
+    })
+
+    const literalIndex = code.indexOf(literalOfInterest)
+    const pathToNode2 = getNodePathFromSourceRange(ast, [
+      literalIndex + 2,
+      literalIndex + 2,
+    ])
+    expect(pathToNode).toEqual(pathToNode2)
+  })
+})

src/lang/queryAst.ts

@@ -270,6 +270,18 @@ function moreNodePathFromSourceRange(
       }
     }
   }
+  if (_node.type === 'MemberExpression' && isInRange) {
+    const { object, property } = _node
+    if (object.start <= start && object.end >= end) {
+      path.push(['object', 'MemberExpression'])
+      return moreNodePathFromSourceRange(object, sourceRange, path)
+    }
+    if (property.start <= start && property.end >= end) {
+      path.push(['property', 'MemberExpression'])
+      return moreNodePathFromSourceRange(property, sourceRange, path)
+    }
+    return path
+  }
   if (_node.type === 'PipeSubstitution' && isInRange) return path
   console.error('not implemented: ' + node.type)
   return path
@@ -307,48 +319,87 @@ type KCLNode =
   | ReturnStatement
 
 export function traverse(
-  node: KCLNode,
+  node: KCLNode | Program,
   option: {
-    enter?: (node: KCLNode) => void
+    enter?: (node: KCLNode, pathToNode: PathToNode) => void
     leave?: (node: KCLNode) => void
-  }
+  },
+  pathToNode: PathToNode = []
 ) {
-  option?.enter?.(node)
-  const _traverse = (node: KCLNode) => traverse(node, option)
-
-  if (node.type === 'VariableDeclaration') {
-    node.declarations.forEach(_traverse)
-  } else if (node.type === 'VariableDeclarator') {
-    _traverse(node.init)
-  } else if (node.type === 'PipeExpression') {
-    node.body.forEach(_traverse)
-  } else if (node.type === 'CallExpression') {
-    _traverse(node.callee)
-    node.arguments.forEach(_traverse)
-  } else if (node.type === 'BinaryExpression') {
-    _traverse(node.left)
-    _traverse(node.right)
-  } else if (node.type === 'Identifier') {
+  const _node = node as KCLNode
+  option?.enter?.(_node, pathToNode)
+  const _traverse = (node: KCLNode, pathToNode: PathToNode) =>
+    traverse(node, option, pathToNode)
+
+  if (_node.type === 'VariableDeclaration') {
+    _node.declarations.forEach((declaration, index) =>
+      _traverse(declaration, [
+        ...pathToNode,
+        ['declarations', 'VariableDeclaration'],
+        [index, 'index'],
+      ])
+    )
+  } else if (_node.type === 'VariableDeclarator') {
+    _traverse(_node.init, [...pathToNode, ['init', '']])
+  } else if (_node.type === 'PipeExpression') {
+    _node.body.forEach((expression, index) =>
+      _traverse(expression, [
+        ...pathToNode,
+        ['body', 'PipeExpression'],
+        [index, 'index'],
+      ])
+    )
+  } else if (_node.type === 'CallExpression') {
+    _traverse(_node.callee, [...pathToNode, ['callee', 'CallExpression']])
+    _node.arguments.forEach((arg, index) =>
+      _traverse(arg, [
+        ...pathToNode,
+        ['arguments', 'CallExpression'],
+        [index, 'index'],
+      ])
+    )
+  } else if (_node.type === 'BinaryExpression') {
+    _traverse(_node.left, [...pathToNode, ['left', 'BinaryExpression']])
+    _traverse(_node.right, [...pathToNode, ['right', 'BinaryExpression']])
+  } else if (_node.type === 'Identifier') {
     // do nothing
-  } else if (node.type === 'Literal') {
+  } else if (_node.type === 'Literal') {
     // do nothing
-  } else if (node.type === 'ArrayExpression') {
-    node.elements.forEach(_traverse)
-  } else if (node.type === 'ObjectExpression') {
-    node.properties.forEach(({ key, value }) => {
-      _traverse(key)
-      _traverse(value)
+  } else if (_node.type === 'ArrayExpression') {
+    _node.elements.forEach((el, index) =>
+      _traverse(el, [
+        ...pathToNode,
+        ['elements', 'ArrayExpression'],
+        [index, 'index'],
+      ])
+    )
+  } else if (_node.type === 'ObjectExpression') {
+    _node.properties.forEach(({ key, value }, index) => {
+      _traverse(key, [
+        ...pathToNode,
+        ['properties', 'ObjectExpression'],
+        [index, 'index'],
+        ['key', 'Property'],
+      ])
+      _traverse(value, [
+        ...pathToNode,
+        ['properties', 'ObjectExpression'],
+        [index, 'index'],
+        ['value', 'Property'],
+      ])
     })
-  } else if (node.type === 'UnaryExpression') {
-    _traverse(node.argument)
-  } else if (node.type === 'MemberExpression') {
+  } else if (_node.type === 'UnaryExpression') {
+    _traverse(_node.argument, [...pathToNode, ['argument', 'UnaryExpression']])
+  } else if (_node.type === 'MemberExpression') {
     // hmm this smell
-    _traverse(node.object)
-    _traverse(node.property)
-  } else if ('body' in node && Array.isArray(node.body)) {
-    node.body.forEach(_traverse)
+    _traverse(_node.object, [...pathToNode, ['object', 'MemberExpression']])
+    _traverse(_node.property, [...pathToNode, ['property', 'MemberExpression']])
+  } else if ('body' in _node && Array.isArray(_node.body)) {
+    _node.body.forEach((expression, index) =>
+      _traverse(expression, [...pathToNode, ['body', ''], [index, 'index']])
+    )
   }
-  option?.leave?.(node)
+  option?.leave?.(_node)
 }
 
 export interface PrevVariable<T> {