WIP new util function

2025-07-03 12:15:20 -04:00
parent 8955b5fcd3
commit a0fe33260e
3 changed files with 105 additions and 153 deletions
--- a/src/lang/modifyAst.test.ts
+++ b/src/lang/modifyAst.test.ts
@ -2,7 +2,9 @@ import type { Node } from '@rust/kcl-lib/bindings/Node'

 import {
  createArrayExpression,
+  createCallExpressionStdLibKw,
  createIdentifier,
+  createLabeledArg,
  createLiteral,
  createLiteralMaybeSuffix,
  createLocalName,
@ -19,6 +21,7 @@ import {
  createVariableExpressionsArray,
  deleteSegmentFromPipeExpression,
  moveValueIntoNewVariable,
+  setCallInAst,
  sketchOnExtrudedFace,
  splitPipedProfile,
 } from '@src/lang/modifyAst'
@ -1037,3 +1040,45 @@ extrude001 = extrude(profile001, length = 123)
    expect(node.node.end).toEqual(startOfKwargIndex + 3)
  })
 })
+
+describe('Testing setCallInAst', () => {
+  it('should push an extrude call with variable on variable profile', () => {
+    const code = `sketch001 = startSketchOn(XY)
+profile001 = circle(sketch001, center = [0, 0], radius = 1)
+`
+    const ast = assertParse(code)
+    const exprs = createVariableExpressionsArray([
+      createLocalName('profile001'),
+    ])
+    const call = createCallExpressionStdLibKw('extrude', exprs, [
+      createLabeledArg('length', createLiteral(5)),
+    ])
+    const pathToNode = setCallInAst(ast, call)
+    if (err(pathToNode)) {
+      throw pathToNode
+    }
+    const newCode = recast(ast)
+    expect(newCode).toContain(code)
+    expect(newCode).toContain(`extrude001 = extrude(profile001, length = 5)`)
+  })
+
+  it('should push an extrude call with variable on variable profile', () => {
+    const code = `startSketchOn(XY)
+  |> circle(sketch001, center = [0, 0], radius = 1)
+`
+    const ast = assertParse(code)
+    const exprs = createVariableExpressionsArray([
+      createLocalName('profile001'),
+    ])
+    const call = createCallExpressionStdLibKw('extrude', exprs, [
+      createLabeledArg('length', createLiteral(5)),
+    ])
+    const pathToNode = setCallInAst(ast, call)
+    if (err(pathToNode)) {
+      throw pathToNode
+    }
+    const newCode = recast(ast)
+    expect(newCode).toContain(code)
+    expect(newCode).toContain(`|> extrude(length = 5)`)
+  })
+})
--- a/src/lang/modifyAst.ts
+++ b/src/lang/modifyAst.ts
@ -1247,3 +1247,45 @@ export function createPathToNodeForLastVariable(

  return pathToCall
 }
+
+export function setCallInAst(
+  ast: Node<Program>,
+  call: Node<CallExpressionKw>,
+  nodeToEdit?: PathToNode,
+  lastPathToNode?: PathToNode
+): Error | PathToNode {
+  let pathToNode: PathToNode | undefined
+  if (nodeToEdit) {
+    const result = getNodeFromPath<CallExpressionKw>(
+      ast,
+      nodeToEdit,
+      'CallExpressionKw'
+    )
+    if (err(result)) {
+      return result
+    }
+
+    Object.assign(result.node, call)
+    pathToNode = nodeToEdit
+  } else {
+    if (!call.unlabeled && lastPathToNode) {
+      const pipe = getNodeFromPath<PipeExpression>(
+        ast,
+        lastPathToNode,
+        'PipeExpression'
+      )
+      if (err(pipe)) {
+        return pipe
+      }
+      pipe.node.body.push(call)
+      pathToNode = lastPathToNode
+    } else {
+      const name = findUniqueName(ast, call.callee.name.name)
+      const declaration = createVariableDeclaration(name, call)
+      ast.body.push(declaration)
+      pathToNode = createPathToNodeForLastVariable(ast)
+    }
+  }
+
+  return pathToNode
+}
--- a/src/lang/modifyAst/addSweep.ts
+++ b/src/lang/modifyAst/addSweep.ts
@ -5,13 +5,11 @@ import {
  createLabeledArg,
  createLiteral,
  createLocalName,
-  createVariableDeclaration,
-  findUniqueName,
 } from '@src/lang/create'
 import {
-  createPathToNodeForLastVariable,
  createVariableExpressionsArray,
  insertVariableAndOffsetPathToNode,
+  setCallInAst,
 } from '@src/lang/modifyAst'
 import {
  getEdgeTagCall,
@ -23,15 +21,8 @@ import {
  valueOrVariable,
 } from '@src/lang/queryAst'
 import { getNodePathFromSourceRange } from '@src/lang/queryAstNodePathUtils'
-import type {
-  CallExpressionKw,
-  PathToNode,
-  PipeExpression,
-  Program,
-  VariableDeclaration,
-} from '@src/lang/wasm'
+import type { PathToNode, Program, VariableDeclaration } from '@src/lang/wasm'
 import type { KclCommandValue } from '@src/lib/commandTypes'
-import { KCL_DEFAULT_CONSTANT_PREFIXES } from '@src/lib/constants'
 import type { Selections } from '@src/lib/selections'
 import { err } from '@src/lib/trap'

@ -116,42 +107,10 @@ export function addExtrude({

  // 3. If edit, we assign the new function call declaration to the existing node,
  // otherwise just push to the end
-  let pathToNode: PathToNode | undefined
-  if (nodeToEdit) {
-    const result = getNodeFromPath<CallExpressionKw>(
-      modifiedAst,
-      nodeToEdit,
-      'CallExpressionKw'
-    )
-    if (err(result)) {
-      return result
-    }
-
-    Object.assign(result.node, call)
-    pathToNode = nodeToEdit
-  } else {
-    const lastPathToNode: PathToNode | undefined =
-      variableExpressions.paths.pop()
-    if (sketchesExpr === null && lastPathToNode) {
-      const pipe = getNodeFromPath<PipeExpression>(
-        modifiedAst,
-        lastPathToNode,
-        'PipeExpression'
-      )
-      if (err(pipe)) {
-        return pipe
-      }
-      pipe.node.body.push(call)
-      pathToNode = lastPathToNode
-    } else {
-      const name = findUniqueName(
-        modifiedAst,
-        KCL_DEFAULT_CONSTANT_PREFIXES.EXTRUDE
-      )
-      const declaration = createVariableDeclaration(name, call)
-      modifiedAst.body.push(declaration)
-      pathToNode = createPathToNodeForLastVariable(modifiedAst)
-    }
+  const lastPath = variableExpressions.paths.pop() // TODO: check if this is correct
+  const pathToNode = setCallInAst(modifiedAst, call, nodeToEdit, lastPath)
+  if (err(pathToNode)) {
+    return pathToNode
  }

  return {
@ -222,41 +181,10 @@ export function addSweep({

  // 3. If edit, we assign the new function call declaration to the existing node,
  // otherwise just push to the end
-  let pathToNode: PathToNode | undefined
-  if (nodeToEdit) {
-    const result = getNodeFromPath<CallExpressionKw>(
-      modifiedAst,
-      nodeToEdit,
-      'CallExpressionKw'
-    )
-    if (err(result)) {
-      return result
-    }
-
-    Object.assign(result.node, call)
-    pathToNode = nodeToEdit
-  } else {
-    const lastPathToNode: PathToNode | undefined = variableExprs.paths.pop()
-    if (sketchesExpr === null && lastPathToNode) {
-      const pipe = getNodeFromPath<PipeExpression>(
-        modifiedAst,
-        lastPathToNode,
-        'PipeExpression'
-      )
-      if (err(pipe)) {
-        return pipe
-      }
-      pipe.node.body.push(call)
-      pathToNode = lastPathToNode
-    } else {
-      const name = findUniqueName(
-        modifiedAst,
-        KCL_DEFAULT_CONSTANT_PREFIXES.SWEEP
-      )
-      const declaration = createVariableDeclaration(name, call)
-      modifiedAst.body.push(declaration)
-      pathToNode = createPathToNodeForLastVariable(modifiedAst)
-    }
+  const lastPath = variableExprs.paths.pop() // TODO: check if this is correct
+  const pathToNode = setCallInAst(modifiedAst, call, nodeToEdit, lastPath)
+  if (err(pathToNode)) {
+    return pathToNode
  }

  return {
@ -312,42 +240,10 @@ export function addLoft({

  // 3. If edit, we assign the new function call declaration to the existing node,
  // otherwise just push to the end
-  let pathToNode: PathToNode | undefined
-  if (nodeToEdit) {
-    const result = getNodeFromPath<CallExpressionKw>(
-      modifiedAst,
-      nodeToEdit,
-      'CallExpressionKw'
-    )
-    if (err(result)) {
-      return result
-    }
-
-    Object.assign(result.node, call)
-    pathToNode = nodeToEdit
-  } else {
-    const lastPathToNode: PathToNode | undefined = variableExprs.paths.pop()
-    if (sketchesExpr === null && lastPathToNode) {
-      const pipe = getNodeFromPath<PipeExpression>(
-        modifiedAst,
-        lastPathToNode,
-        'PipeExpression'
-      )
-      if (err(pipe)) {
-        return pipe
-      }
-      pipe.node.body.push(call)
-      pathToNode = lastPathToNode
-    } else {
-      const name = findUniqueName(
-        modifiedAst,
-        KCL_DEFAULT_CONSTANT_PREFIXES.LOFT
-      )
-      const declaration = createVariableDeclaration(name, call)
-      modifiedAst.body.push(declaration)
-      const toFirstKwarg = !!vDegree
-      pathToNode = createPathToNodeForLastVariable(modifiedAst, toFirstKwarg)
-    }
+  const lastPath = variableExprs.paths.pop() // TODO: check if this is correct
+  const pathToNode = setCallInAst(modifiedAst, call, nodeToEdit, lastPath)
+  if (err(pathToNode)) {
+    return pathToNode
  }

  return {
@ -447,41 +343,10 @@ export function addRevolve({

  // 3. If edit, we assign the new function call declaration to the existing node,
  // otherwise just push to the end
-  let pathToNode: PathToNode | undefined
-  if (nodeToEdit) {
-    const result = getNodeFromPath<CallExpressionKw>(
-      modifiedAst,
-      nodeToEdit,
-      'CallExpressionKw'
-    )
-    if (err(result)) {
-      return result
-    }
-
-    Object.assign(result.node, call)
-    pathToNode = nodeToEdit
-  } else {
-    const lastPathToNode: PathToNode | undefined = variableExprs.paths.pop()
-    if (sketchesExpr === null && lastPathToNode) {
-      const pipe = getNodeFromPath<PipeExpression>(
-        modifiedAst,
-        lastPathToNode,
-        'PipeExpression'
-      )
-      if (err(pipe)) {
-        return pipe
-      }
-      pipe.node.body.push(call)
-      pathToNode = lastPathToNode
-    } else {
-      const name = findUniqueName(
-        modifiedAst,
-        KCL_DEFAULT_CONSTANT_PREFIXES.REVOLVE
-      )
-      const declaration = createVariableDeclaration(name, call)
-      modifiedAst.body.push(declaration)
-      pathToNode = createPathToNodeForLastVariable(modifiedAst)
-    }
+  const lastPath = variableExprs.paths.pop() // TODO: check if this is correct
+  const pathToNode = setCallInAst(modifiedAst, call, nodeToEdit, lastPath)
+  if (err(pathToNode)) {
+    return pathToNode
  }

  return {