Fix finding path to node for import and if-else (#4483)

src/lang/queryAst.test.ts

@@ -1,4 +1,4 @@
-import { parse, recast, initPromise, PathToNode } from './wasm'
+import { parse, recast, initPromise, PathToNode, Identifier } from './wasm'
 import {
   findAllPreviousVariables,
   isNodeSafeToReplace,
@@ -10,6 +10,7 @@ import {
   hasSketchPipeBeenExtruded,
   doesSceneHaveSweepableSketch,
   traverse,
+  getNodeFromPath,
 } from './queryAst'
 import { enginelessExecutor } from '../lib/testHelpers'
 import {
@@ -266,6 +267,86 @@ describe('testing getNodePathFromSourceRange', () => {
     ])
     expect(selectWholeThing).toEqual(expected)
   })
+
+  it('finds the node in if-else condition', () => {
+    const code = `y = 0
+    x = if x > y {
+      x + 1
+    } else {
+      y
+    }`
+    const searchLn = `x > y`
+    const sourceIndex = code.indexOf(searchLn)
+    const ast = parse(code)
+    if (err(ast)) throw ast
+
+    const result = getNodePathFromSourceRange(ast, [sourceIndex, sourceIndex])
+    expect(result).toEqual([
+      ['body', ''],
+      [1, 'index'],
+      ['declarations', 'VariableDeclaration'],
+      [0, 'index'],
+      ['init', ''],
+      ['cond', 'IfExpression'],
+      ['left', 'BinaryExpression'],
+    ])
+    const _node = getNodeFromPath<Identifier>(ast, result)
+    if (err(_node)) throw _node
+    expect(_node.node.type).toEqual('Identifier')
+    expect(_node.node.name).toEqual('x')
+  })
+
+  it('finds the node in if-else then', () => {
+    const code = `y = 0
+    x = if x > y {
+      x + 1
+    } else {
+      y
+    }`
+    const searchLn = `x + 1`
+    const sourceIndex = code.indexOf(searchLn)
+    const ast = parse(code)
+    if (err(ast)) throw ast
+
+    const result = getNodePathFromSourceRange(ast, [sourceIndex, sourceIndex])
+    expect(result).toEqual([
+      ['body', ''],
+      [1, 'index'],
+      ['declarations', 'VariableDeclaration'],
+      [0, 'index'],
+      ['init', ''],
+      ['then_val', 'IfExpression'],
+      ['body', 'IfExpression'],
+      [0, 'index'],
+      ['expression', 'ExpressionStatement'],
+      ['left', 'BinaryExpression'],
+    ])
+    const _node = getNodeFromPath<Identifier>(ast, result)
+    if (err(_node)) throw _node
+    expect(_node.node.type).toEqual('Identifier')
+    expect(_node.node.name).toEqual('x')
+  })
+
+  it('finds the node in import statement item', () => {
+    const code = `import foo, bar as baz from 'thing.kcl'`
+    const searchLn = `bar`
+    const sourceIndex = code.indexOf(searchLn)
+    const ast = parse(code)
+    if (err(ast)) throw ast
+
+    const result = getNodePathFromSourceRange(ast, [sourceIndex, sourceIndex])
+    expect(result).toEqual([
+      ['body', ''],
+      [0, 'index'],
+      ['items', 'ImportStatement'],
+      [1, 'index'],
+      ['name', 'ImportItem'],
+    ])
+    const _node = getNodeFromPath<Identifier>(ast, result)
+    if (err(_node)) throw _node
+    expect(_node.node.type).toEqual('Identifier')
+    expect(_node.node.name).toEqual('bar')
+  })
 })
 
 describe('testing doesPipeHave', () => {

src/lang/queryAst.ts

@@ -317,6 +317,62 @@ function moreNodePathFromSourceRange(
   }
 
   if (_node.type === 'PipeSubstitution' && isInRange) return path
+
+  if (_node.type === 'IfExpression' && isInRange) {
+    const { cond, then_val, else_ifs, final_else } = _node
+    if (cond.start <= start && cond.end >= end) {
+      path.push(['cond', 'IfExpression'])
+      return moreNodePathFromSourceRange(cond, sourceRange, path)
+    }
+    if (then_val.start <= start && then_val.end >= end) {
+      path.push(['then_val', 'IfExpression'])
+      path.push(['body', 'IfExpression'])
+      return getNodePathFromSourceRange(then_val, sourceRange, path)
+    }
+    for (let i = 0; i < else_ifs.length; i++) {
+      const else_if = else_ifs[i]
+      if (else_if.start <= start && else_if.end >= end) {
+        path.push(['else_ifs', 'IfExpression'])
+        path.push([i, 'index'])
+        const { cond, then_val } = else_if
+        if (cond.start <= start && cond.end >= end) {
+          path.push(['cond', 'IfExpression'])
+          return moreNodePathFromSourceRange(cond, sourceRange, path)
+        }
+        path.push(['then_val', 'IfExpression'])
+        path.push(['body', 'IfExpression'])
+        return getNodePathFromSourceRange(then_val, sourceRange, path)
+      }
+    }
+    if (final_else.start <= start && final_else.end >= end) {
+      path.push(['final_else', 'IfExpression'])
+      path.push(['body', 'IfExpression'])
+      return getNodePathFromSourceRange(final_else, sourceRange, path)
+    }
+    return path
+  }
+
+  if (_node.type === 'ImportStatement' && isInRange) {
+    const { items } = _node
+    for (let i = 0; i < items.length; i++) {
+      const item = items[i]
+      if (item.start <= start && item.end >= end) {
+        path.push(['items', 'ImportStatement'])
+        path.push([i, 'index'])
+        if (item.name.start <= start && item.name.end >= end) {
+          path.push(['name', 'ImportItem'])
+          return path
+        }
+        if (item.alias && item.alias.start <= start && item.alias.end >= end) {
+          path.push(['alias', 'ImportItem'])
+          return path
+        }
+        return path
+      }
+    }
+    return path
+  }
+
   console.error('not implemented: ' + node.type)
 
   return path