WIP new util function

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)`)
+  })
+})

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
+}

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 {
+import type { PathToNode, Program, VariableDeclaration } from '@src/lang/wasm'
-  CallExpressionKw,
-  PathToNode,
-  PipeExpression,
-  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
+  const lastPath = variableExpressions.paths.pop() // TODO: check if this is correct
-  if (nodeToEdit) {
+  const pathToNode = setCallInAst(modifiedAst, call, nodeToEdit, lastPath)
-    const result = getNodeFromPath<CallExpressionKw>(
+  if (err(pathToNode)) {
-      modifiedAst,
+    return pathToNode
-      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)
-    }
   }
 
   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
+  const lastPath = variableExprs.paths.pop() // TODO: check if this is correct
-  if (nodeToEdit) {
+  const pathToNode = setCallInAst(modifiedAst, call, nodeToEdit, lastPath)
-    const result = getNodeFromPath<CallExpressionKw>(
+  if (err(pathToNode)) {
-      modifiedAst,
+    return pathToNode
-      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)
-    }
   }
 
   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
+  const lastPath = variableExprs.paths.pop() // TODO: check if this is correct
-  if (nodeToEdit) {
+  const pathToNode = setCallInAst(modifiedAst, call, nodeToEdit, lastPath)
-    const result = getNodeFromPath<CallExpressionKw>(
+  if (err(pathToNode)) {
-      modifiedAst,
+    return pathToNode
-      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)
-    }
   }
 
   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
+  const lastPath = variableExprs.paths.pop() // TODO: check if this is correct
-  if (nodeToEdit) {
+  const pathToNode = setCallInAst(modifiedAst, call, nodeToEdit, lastPath)
-    const result = getNodeFromPath<CallExpressionKw>(
+  if (err(pathToNode)) {
-      modifiedAst,
+    return pathToNode
-      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)
-    }
   }
 
   return {