Merge branch 'main' into pierremtb/issue7657-Allow-all-sweeps-to-work-on-variable-less-profiles

Fixes #7657

src/lang/modifyAst.test.ts

@@ -2,9 +2,12 @@ import type { Node } from '@rust/kcl-lib/bindings/Node'
 
 import {
   createArrayExpression,
+  createCallExpressionStdLibKw,
   createIdentifier,
+  createLabeledArg,
   createLiteral,
   createLiteralMaybeSuffix,
+  createLocalName,
   createObjectExpression,
   createPipeExpression,
   createPipeSubstitution,
@@ -14,12 +17,19 @@ import {
 } from '@src/lang/create'
 import {
   addSketchTo,
+  createPathToNodeForLastVariable,
+  createVariableExpressionsArray,
   deleteSegmentFromPipeExpression,
   moveValueIntoNewVariable,
+  setCallInAst,
   sketchOnExtrudedFace,
   splitPipedProfile,
 } from '@src/lang/modifyAst'
-import { findUsesOfTagInPipe } from '@src/lang/queryAst'
+import {
+  findUsesOfTagInPipe,
+  getNodeFromPath,
+  getVariableExprsFromSelection,
+} from '@src/lang/queryAst'
 import { getNodePathFromSourceRange } from '@src/lang/queryAstNodePathUtils'
 import type { Artifact } from '@src/lang/std/artifactGraph'
 import { codeRefFromRange } from '@src/lang/std/artifactGraph'
@@ -31,6 +41,7 @@ import { enginelessExecutor } from '@src/lib/testHelpers'
 import { err } from '@src/lib/trap'
 import { deleteFromSelection } from '@src/lang/modifyAst/deleteFromSelection'
 import { assertNotErr } from '@src/unitTestUtils'
+import type { Selections } from '@src/lib/selections'
 
 beforeAll(async () => {
   await initPromise
@@ -917,3 +928,212 @@ extrude001 = extrude(part001, length = 5)
     expect(result instanceof Error).toBe(true)
   })
 })
+
+describe('Testing createVariableExpressionsArray', () => {
+  it('should return null for any number of pipe substitutions', () => {
+    const onePipe = [createPipeSubstitution()]
+    const twoPipes = [createPipeSubstitution(), createPipeSubstitution()]
+    const threePipes = [
+      createPipeSubstitution(),
+      createPipeSubstitution(),
+      createPipeSubstitution(),
+    ]
+    expect(createVariableExpressionsArray(onePipe)).toBeNull()
+    expect(createVariableExpressionsArray(twoPipes)).toBeNull()
+    expect(createVariableExpressionsArray(threePipes)).toBeNull()
+  })
+
+  it('should create a variable expressions for one variable', () => {
+    const oneVariableName = [createLocalName('var1')]
+    const expr = createVariableExpressionsArray(oneVariableName)
+    if (expr?.type !== 'Name') {
+      throw new Error(`Expected Literal type, got ${expr?.type}`)
+    }
+
+    expect(expr.name.name).toBe('var1')
+  })
+
+  it('should create an array of variable expressions for two variables', () => {
+    const twoVariableNames = [createLocalName('var1'), createLocalName('var2')]
+    const exprs = createVariableExpressionsArray(twoVariableNames)
+    if (exprs?.type !== 'ArrayExpression') {
+      throw new Error('Expected ArrayExpression type')
+    }
+
+    expect(exprs.elements).toHaveLength(2)
+    if (
+      exprs.elements[0].type !== 'Name' ||
+      exprs.elements[1].type !== 'Name'
+    ) {
+      throw new Error(
+        `Expected elements to be of type Name, got ${exprs.elements[0].type} and ${exprs.elements[1].type}`
+      )
+    }
+    expect(exprs.elements[0].name.name).toBe('var1')
+    expect(exprs.elements[1].name.name).toBe('var2')
+  })
+
+  // This would catch the issue at https://github.com/KittyCAD/modeling-app/issues/7669
+  // TODO: fix function to get this test to pass
+  // it('should create one expr if the array of variable names are the same', () => {
+  //   const twoVariableNames = [createLocalName('var1'), createLocalName('var1')]
+  //   const expr = createVariableExpressionsArray(twoVariableNames)
+  //   if (expr?.type !== 'Name') {
+  //     throw new Error(`Expected Literal type, got ${expr?.type}`)
+  //   }
+
+  //   expect(expr.name.name).toBe('var1')
+  // })
+
+  it('should create an array of variable expressions for one variable and a pipe', () => {
+    const oneVarOnePipe = [createPipeSubstitution(), createLocalName('var1')]
+    const exprs = createVariableExpressionsArray(oneVarOnePipe)
+    if (exprs?.type !== 'ArrayExpression') {
+      throw new Error('Expected ArrayExpression type')
+    }
+
+    expect(exprs.elements).toHaveLength(2)
+    expect(exprs.elements[0].type).toBe('PipeSubstitution')
+    if (exprs.elements[1].type !== 'Name') {
+      throw new Error(
+        `Expected elements[1] to be of type Name, got ${exprs.elements[1].type}`
+      )
+    }
+
+    expect(exprs.elements[1].name.name).toBe('var1')
+  })
+})
+
+describe('Testing createPathToNodeForLastVariable', () => {
+  it('should create a path to the last variable in the array', () => {
+    const circleProfileInVar = `sketch001 = startSketchOn(XY)
+profile001 = circle(sketch001, center = [0, 0], radius = 1)
+extrude001 = extrude(profile001, length = 5)
+`
+    const ast = assertParse(circleProfileInVar)
+    const path = createPathToNodeForLastVariable(ast, false)
+    expect(path.length).toEqual(4)
+
+    // Verify we can get the right node
+    const node = getNodeFromPath<any>(ast, path)
+    if (err(node)) {
+      throw node
+    }
+    // With the expected range
+    const startOfExtrudeIndex = circleProfileInVar.indexOf('extrude(')
+    expect(node.node.start).toEqual(startOfExtrudeIndex)
+    expect(node.node.end).toEqual(circleProfileInVar.length - 1)
+  })
+
+  it('should create a path to the first kwarg in the last expression', () => {
+    const circleProfileInVar = `sketch001 = startSketchOn(XY)
+profile001 = circle(sketch001, center = [0, 0], radius = 1)
+extrude001 = extrude(profile001, length = 123)
+`
+    const ast = assertParse(circleProfileInVar)
+    const path = createPathToNodeForLastVariable(ast, true)
+    expect(path.length).toEqual(7)
+
+    // Verify we can get the right node
+    const node = getNodeFromPath<any>(ast, path)
+    if (err(node)) {
+      throw node
+    }
+    // With the expected range
+    const startOfKwargIndex = circleProfileInVar.indexOf('123')
+    expect(node.node.start).toEqual(startOfKwargIndex)
+    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 in pipe is selection was in variable-less pipe', async () => {
+    const code = `startSketchOn(XY)
+  |> circle(center = [0, 0], radius = 1)
+`
+    const ast = assertParse(code)
+    const { artifactGraph } = await enginelessExecutor(ast)
+    const artifact = artifactGraph.values().find((a) => a.type === 'path')
+    if (!artifact) {
+      throw new Error('Artifact not found in the graph')
+    }
+    const selections: Selections = {
+      graphSelections: [
+        {
+          codeRef: artifact.codeRef,
+          artifact,
+        },
+      ],
+      otherSelections: [],
+    }
+    const variableExprs = getVariableExprsFromSelection(selections, ast)
+    if (err(variableExprs)) throw variableExprs
+    const exprs = createVariableExpressionsArray(variableExprs.exprs)
+    const call = createCallExpressionStdLibKw('extrude', exprs, [
+      createLabeledArg('length', createLiteral(5)),
+    ])
+    const lastPathToNode = variableExprs.paths.pop()
+    const pathToNode = setCallInAst(ast, call, undefined, lastPathToNode)
+    if (err(pathToNode)) {
+      throw pathToNode
+    }
+    const newCode = recast(ast)
+    expect(newCode).toContain(code)
+    expect(newCode).toContain(`|> extrude(length = 5)`)
+  })
+
+  it('should push an extrude call with variable if selection was in variable pipe', async () => {
+    const code = `profile001 = startSketchOn(XY)
+  |> circle(center = [0, 0], radius = 1)
+`
+    const ast = assertParse(code)
+    const { artifactGraph } = await enginelessExecutor(ast)
+    const artifact = artifactGraph.values().find((a) => a.type === 'path')
+    if (!artifact) {
+      throw new Error('Artifact not found in the graph')
+    }
+    const selections: Selections = {
+      graphSelections: [
+        {
+          codeRef: artifact.codeRef,
+          artifact,
+        },
+      ],
+      otherSelections: [],
+    }
+    const variableExprs = getVariableExprsFromSelection(selections, ast)
+    if (err(variableExprs)) throw variableExprs
+    const exprs = createVariableExpressionsArray(variableExprs.exprs)
+    const call = createCallExpressionStdLibKw('extrude', exprs, [
+      createLabeledArg('length', createLiteral(5)),
+    ])
+    const lastPathToNode = variableExprs.paths.pop()
+    const pathToNode = setCallInAst(ast, call, undefined, lastPathToNode)
+    if (err(pathToNode)) {
+      throw pathToNode
+    }
+    const newCode = recast(ast)
+    expect(newCode).toContain(code)
+    expect(newCode).toContain(`extrude001 = extrude(profile001, length = 5)`)
+  })
+})

src/lang/modifyAst.ts

@@ -1209,3 +1209,83 @@ export function insertVariableAndOffsetPathToNode(
     }
   }
 }
+
+// Create an array expression for variables,
+// or keep it null if all are PipeSubstitutions
+export function createVariableExpressionsArray(sketches: Expr[]): Expr | null {
+  let exprs: Expr | null = null
+  if (sketches.every((s) => s.type === 'PipeSubstitution')) {
+    // Keeping null so we don't even put it the % sign
+  } else if (sketches.length === 1) {
+    exprs = sketches[0]
+  } else {
+    exprs = createArrayExpression(sketches)
+  }
+  return exprs
+}
+
+// Create a path to node to the last variable declaroator of an ast
+// Optionally, can point to the first kwarg of the CallExpressionKw
+export function createPathToNodeForLastVariable(
+  ast: Node<Program>,
+  toFirstKwarg = true
+): PathToNode {
+  const argIndex = 0 // first kwarg for all sweeps here
+  const pathToCall: PathToNode = [
+    ['body', ''],
+    [ast.body.length - 1, 'index'],
+    ['declaration', 'VariableDeclaration'],
+    ['init', 'VariableDeclarator'],
+  ]
+  if (toFirstKwarg) {
+    pathToCall.push(
+      ['arguments', 'CallExpressionKw'],
+      [argIndex, ARG_INDEX_FIELD],
+      ['arg', LABELED_ARG_FIELD]
+    )
+  }
+
+  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/sweeps.test.ts

@@ -0,0 +1,252 @@
+import {
+  type Artifact,
+  assertParse,
+  type CodeRef,
+  type Program,
+  recast,
+} from '@src/lang/wasm'
+import type { Selection, Selections } from '@src/lib/selections'
+import { enginelessExecutor } from '@src/lib/testHelpers'
+import { err } from '@src/lib/trap'
+import {
+  addExtrude,
+  addLoft,
+  addRevolve,
+  addSweep,
+} from '@src/lang/modifyAst/sweeps'
+import { stringToKclExpression } from '@src/lib/kclHelpers'
+import type { Node } from '@rust/kcl-lib/bindings/Node'
+
+async function getAstAndArtifactGraph(code: string) {
+  const ast = assertParse(code)
+  if (err(ast)) throw ast
+
+  const { artifactGraph } = await enginelessExecutor(ast)
+  return { ast, artifactGraph }
+}
+
+function createSelectionFromPathArtifact(
+  artifacts: (Artifact & { codeRef: CodeRef })[]
+): Selections {
+  const graphSelections = artifacts.map(
+    (artifact) =>
+      ({
+        codeRef: artifact.codeRef,
+        artifact,
+      }) as Selection
+  )
+  return {
+    graphSelections,
+    otherSelections: [],
+  }
+}
+
+async function getAstAndSketchSelections(code: string) {
+  const { ast, artifactGraph } = await getAstAndArtifactGraph(code)
+  const artifacts = [...artifactGraph.values()].filter((a) => a.type === 'path')
+  if (artifacts.length === 0) {
+    throw new Error('Artifact not found in the graph')
+  }
+  const sketches = createSelectionFromPathArtifact(artifacts)
+  return { ast, sketches }
+}
+
+async function getKclCommandValue(value: string) {
+  const result = await stringToKclExpression(value)
+  if (err(result) || 'errors' in result) {
+    throw new Error(`Couldn't create kcl expression`)
+  }
+
+  return result
+}
+
+async function runNewAstAndCheckForSweep(ast: Node<Program>) {
+  const { artifactGraph } = await enginelessExecutor(ast)
+  console.log('artifactGraph', artifactGraph)
+  const sweepArtifact = artifactGraph.values().find((a) => a.type === 'sweep')
+  expect(sweepArtifact).toBeDefined()
+}
+
+describe('Testing addExtrude', () => {
+  it('should push a call in pipe if selection was in variable-less pipe', async () => {
+    const code = `startSketchOn(XY)
+  |> circle(center = [0, 0], radius = 1)
+`
+    const { ast, sketches } = await getAstAndSketchSelections(code)
+    const length = await getKclCommandValue('1')
+    const result = addExtrude({ ast, sketches, length })
+    if (err(result)) throw result
+    const newCode = recast(result.modifiedAst)
+    expect(newCode).toContain(code)
+    expect(newCode).toContain(`|> extrude(length = 1)`)
+    await runNewAstAndCheckForSweep(result.modifiedAst)
+  })
+
+  it('should push a call with variable if selection was in variable profile', async () => {
+    const code = `sketch001 = startSketchOn(XY)
+profile001 = circle(sketch001, center = [0, 0], radius = 1)
+`
+    const { ast, sketches } = await getAstAndSketchSelections(code)
+    const length = await getKclCommandValue('2')
+    const result = addExtrude({ ast, sketches, length })
+    if (err(result)) throw result
+    const newCode = recast(result.modifiedAst)
+    expect(newCode).toContain(code)
+    expect(newCode).toContain(`extrude001 = extrude(profile001, length = 2)`)
+    await runNewAstAndCheckForSweep(result.modifiedAst)
+  })
+
+  it('should push a call with variable if selection was in variable pipe', async () => {
+    const code = `profile001 = startSketchOn(XY)
+  |> circle(center = [0, 0], radius = 1)
+`
+    const { ast, sketches } = await getAstAndSketchSelections(code)
+    const length = await getKclCommandValue('3')
+    const result = addExtrude({ ast, sketches, length })
+    if (err(result)) throw result
+    await runNewAstAndCheckForSweep(result.modifiedAst)
+    const newCode = recast(result.modifiedAst)
+    expect(newCode).toContain(code)
+    expect(newCode).toContain(`extrude001 = extrude(profile001, length = 3)`)
+  })
+
+  it('should push a call with many compatible optional args if asked', async () => {
+    const code = `sketch001 = startSketchOn(XY)
+profile001 = circle(sketch001, center = [0, 0], radius = 1)
+`
+    const { ast, sketches } = await getAstAndSketchSelections(code)
+    const length = await getKclCommandValue('10')
+    const bidirectionalLength = await getKclCommandValue('20')
+    const twistAngle = await getKclCommandValue('30')
+    const result = addExtrude({
+      ast,
+      sketches,
+      length,
+      bidirectionalLength,
+      twistAngle,
+    })
+    if (err(result)) throw result
+    await runNewAstAndCheckForSweep(result.modifiedAst)
+    const newCode = recast(result.modifiedAst)
+    expect(newCode).toContain(code)
+    expect(newCode).toContain(`extrude001 = extrude(
+  profile001,
+  length = 10,
+  bidirectionalLength = 20,
+  twistAngle = 30,
+)`)
+  })
+
+  // TODO: missing edit flow test
+
+  // TODO: missing multi-profile test
+})
+
+describe('Testing addSweep', () => {
+  it('should push a call with variable and all compatible optional args', async () => {
+    const code = `sketch001 = startSketchOn(XY)
+profile001 = circle(sketch001, center = [0, 0], radius = 1)
+sketch002 = startSketchOn(XZ)
+profile002 = startProfile(sketch002, at = [0, 0])
+  |> xLine(length = -5)
+  |> tangentialArc(endAbsolute = [-20, 5])
+`
+    const { ast, artifactGraph } = await getAstAndArtifactGraph(code)
+    const artifact1 = artifactGraph.values().find((a) => a.type === 'path')
+    const artifact2 = [...artifactGraph.values()].findLast(
+      (a) => a.type === 'path'
+    )
+    if (!artifact1 || !artifact2) {
+      throw new Error('Artifact not found in the graph')
+    }
+
+    const sketches = createSelectionFromPathArtifact([artifact1])
+    const path = createSelectionFromPathArtifact([artifact2])
+    const sectional = true
+    const relativeTo = 'sketchPlane'
+    const result = addSweep({
+      ast,
+      sketches,
+      path,
+      sectional,
+      relativeTo,
+    })
+    if (err(result)) throw result
+    await runNewAstAndCheckForSweep(result.modifiedAst)
+    const newCode = recast(result.modifiedAst)
+    expect(newCode).toContain(code)
+    expect(newCode).toContain(`sweep001 = sweep(
+  profile001,
+  path = profile002,
+  sectional = true,
+  relativeTo = 'sketchPlane',
+)`)
+  })
+
+  // TODO: missing edit flow test
+
+  // TODO: missing multi-profile test
+})
+
+describe('Testing addLoft', () => {
+  it('should push a call with variable and all optional args if asked', async () => {
+    const code = `sketch001 = startSketchOn(XZ)
+profile001 = circle(sketch001, center = [0, 0], radius = 30)
+plane001 = offsetPlane(XZ, offset = 50)
+sketch002 = startSketchOn(plane001)
+profile002 = circle(sketch002, center = [0, 0], radius = 20)
+`
+    const { ast, sketches } = await getAstAndSketchSelections(code)
+    expect(sketches.graphSelections).toHaveLength(2)
+    const vDegree = await getKclCommandValue('3')
+    const result = addLoft({
+      ast,
+      sketches,
+      vDegree,
+    })
+    if (err(result)) throw result
+    const newCode = recast(result.modifiedAst)
+    expect(newCode).toContain(code)
+    expect(newCode).toContain(
+      `loft001 = loft([profile001, profile002], vDegree = 3)`
+    )
+    // Don't think we can find the artifact here for loft?
+  })
+
+  // TODO: missing edit flow test
+})
+
+describe('Testing addRevolve', () => {
+  it('should push a call with variable and compatible optional args if asked', async () => {
+    const code = `sketch001 = startSketchOn(XZ)
+profile001 = circle(sketch001, center = [3, 0], radius = 1)
+`
+    const { ast, sketches } = await getAstAndSketchSelections(code)
+    expect(sketches.graphSelections).toHaveLength(1)
+    const result = addRevolve({
+      ast,
+      sketches,
+      angle: await getKclCommandValue('1'),
+      axisOrEdge: 'Axis',
+      axis: 'X',
+      edge: undefined,
+      symmetric: false,
+      bidirectionalAngle: await getKclCommandValue('2'),
+    })
+    if (err(result)) throw result
+    await runNewAstAndCheckForSweep(result.modifiedAst)
+    const newCode = recast(result.modifiedAst)
+    console.log(newCode)
+    expect(newCode).toContain(code)
+    expect(newCode).toContain(`revolve001 = revolve(
+  profile001,
+  angle = 1,
+  axis = X,
+  bidirectionalAngle = 2,
+)`)
+  })
+
+  // TODO: missing edit flow test
+
+  // TODO: missing multi-profile test
+})

src/lang/modifyAst/sweeps.ts

@@ -1,35 +1,28 @@
 import type { Node } from '@rust/kcl-lib/bindings/Node'
 
 import {
-  createArrayExpression,
   createCallExpressionStdLibKw,
   createLabeledArg,
   createLiteral,
   createLocalName,
-  createVariableDeclaration,
-  findUniqueName,
 } from '@src/lang/create'
-import { insertVariableAndOffsetPathToNode } from '@src/lang/modifyAst'
+import {
+  createVariableExpressionsArray,
+  insertVariableAndOffsetPathToNode,
+  setCallInAst,
+} from '@src/lang/modifyAst'
 import {
   getEdgeTagCall,
   mutateAstWithTagForSketchSegment,
 } from '@src/lang/modifyAst/addEdgeTreatment'
 import {
   getNodeFromPath,
-  getSketchExprsFromSelection,
+  getVariableExprsFromSelection,
   valueOrVariable,
 } from '@src/lang/queryAst'
-import { ARG_INDEX_FIELD, LABELED_ARG_FIELD } from '@src/lang/queryAstConstants'
 import { getNodePathFromSourceRange } from '@src/lang/queryAstNodePathUtils'
-import type {
-  CallExpressionKw,
-  Expr,
-  PathToNode,
-  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'
 
@@ -60,13 +53,13 @@ export function addExtrude({
 
   // 2. Prepare unlabeled and labeled arguments
   // Map the sketches selection into a list of kcl expressions to be passed as unlabelled argument
-  const sketchesExprList = getSketchExprsFromSelection(
+  const variableExpressions = getVariableExprsFromSelection(
     sketches,
     modifiedAst,
     nodeToEdit
   )
-  if (err(sketchesExprList)) {
-    return sketchesExprList
+  if (err(variableExpressions)) {
+    return variableExpressions
   }
 
   // Extra labeled args expressions
@@ -85,7 +78,7 @@ export function addExtrude({
     ? [createLabeledArg('twistAngle', valueOrVariable(twistAngle))]
     : []
 
-  const sketchesExpr = createSketchExpression(sketchesExprList)
+  const sketchesExpr = createVariableExpressionsArray(variableExpressions.exprs)
   const call = createCallExpressionStdLibKw('extrude', sketchesExpr, [
     createLabeledArg('length', valueOrVariable(length)),
     ...symmetricExpr,
@@ -114,27 +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 name = findUniqueName(
-      modifiedAst,
-      KCL_DEFAULT_CONSTANT_PREFIXES.EXTRUDE
-    )
-    const declaration = createVariableDeclaration(name, call)
-    modifiedAst.body.push(declaration)
-    pathToNode = createPathToNode(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 {
@@ -168,13 +144,13 @@ export function addSweep({
 
   // 2. Prepare unlabeled and labeled arguments
   // Map the sketches selection into a list of kcl expressions to be passed as unlabelled argument
-  const sketchesExprList = getSketchExprsFromSelection(
+  const variableExprs = getVariableExprsFromSelection(
     sketches,
     modifiedAst,
     nodeToEdit
   )
-  if (err(sketchesExprList)) {
-    return sketchesExprList
+  if (err(variableExprs)) {
+    return variableExprs
   }
 
   // Find the path declaration for the labeled argument
@@ -196,7 +172,7 @@ export function addSweep({
     ? [createLabeledArg('relativeTo', createLiteral(relativeTo))]
     : []
 
-  const sketchesExpr = createSketchExpression(sketchesExprList)
+  const sketchesExpr = createVariableExpressionsArray(variableExprs.exprs)
   const call = createCallExpressionStdLibKw('sweep', sketchesExpr, [
     createLabeledArg('path', pathExpr),
     ...sectionalExpr,
@@ -205,27 +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 name = findUniqueName(
-      modifiedAst,
-      KCL_DEFAULT_CONSTANT_PREFIXES.SWEEP
-    )
-    const declaration = createVariableDeclaration(name, call)
-    modifiedAst.body.push(declaration)
-    pathToNode = createPathToNode(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 {
@@ -255,13 +214,13 @@ export function addLoft({
 
   // 2. Prepare unlabeled and labeled arguments
   // Map the sketches selection into a list of kcl expressions to be passed as unlabelled argument
-  const sketchesExprList = getSketchExprsFromSelection(
+  const variableExprs = getVariableExprsFromSelection(
     sketches,
     modifiedAst,
     nodeToEdit
   )
-  if (err(sketchesExprList)) {
-    return sketchesExprList
+  if (err(variableExprs)) {
+    return variableExprs
   }
 
   // Extra labeled args expressions
@@ -269,7 +228,7 @@ export function addLoft({
     ? [createLabeledArg('vDegree', valueOrVariable(vDegree))]
     : []
 
-  const sketchesExpr = createSketchExpression(sketchesExprList)
+  const sketchesExpr = createVariableExpressionsArray(variableExprs.exprs)
   const call = createCallExpressionStdLibKw('loft', sketchesExpr, [
     ...vDegreeExpr,
   ])
@@ -281,25 +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 name = findUniqueName(modifiedAst, KCL_DEFAULT_CONSTANT_PREFIXES.LOFT)
-    const declaration = createVariableDeclaration(name, call)
-    modifiedAst.body.push(declaration)
-    const toFirstKwarg = !!vDegree
-    pathToNode = createPathToNode(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 {
@@ -339,13 +283,13 @@ export function addRevolve({
 
   // 2. Prepare unlabeled and labeled arguments
   // Map the sketches selection into a list of kcl expressions to be passed as unlabelled argument
-  const sketchesExprList = getSketchExprsFromSelection(
+  const variableExprs = getVariableExprsFromSelection(
     sketches,
     modifiedAst,
     nodeToEdit
   )
-  if (err(sketchesExprList)) {
-    return sketchesExprList
+  if (err(variableExprs)) {
+    return variableExprs
   }
 
   // Retrieve axis expression depending on mode
@@ -372,7 +316,7 @@ export function addRevolve({
       ]
     : []
 
-  const sketchesExpr = createSketchExpression(sketchesExprList)
+  const sketchesExpr = createVariableExpressionsArray(variableExprs.exprs)
   const call = createCallExpressionStdLibKw('revolve', sketchesExpr, [
     createLabeledArg('angle', valueOrVariable(angle)),
     createLabeledArg('axis', getAxisResult.generatedAxis),
@@ -399,27 +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 name = findUniqueName(
-      modifiedAst,
-      KCL_DEFAULT_CONSTANT_PREFIXES.REVOLVE
-    )
-    const declaration = createVariableDeclaration(name, call)
-    modifiedAst.body.push(declaration)
-    pathToNode = createPathToNode(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 {
@@ -430,40 +357,6 @@ export function addRevolve({
 
 // Utilities
 
-function createSketchExpression(sketches: Expr[]) {
-  let sketchesExpr: Expr | null = null
-  if (sketches.every((s) => s.type === 'PipeSubstitution')) {
-    // Keeping null so we don't even put it the % sign
-  } else if (sketches.length === 1) {
-    sketchesExpr = sketches[0]
-  } else {
-    sketchesExpr = createArrayExpression(sketches)
-  }
-  return sketchesExpr
-}
-
-function createPathToNode(
-  modifiedAst: Node<Program>,
-  toFirstKwarg = true
-): PathToNode {
-  const argIndex = 0 // first kwarg for all sweeps here
-  const pathToCall: PathToNode = [
-    ['body', ''],
-    [modifiedAst.body.length - 1, 'index'],
-    ['declaration', 'VariableDeclaration'],
-    ['init', 'VariableDeclarator'],
-  ]
-  if (toFirstKwarg) {
-    pathToCall.push(
-      ['arguments', 'CallExpressionKw'],
-      [argIndex, ARG_INDEX_FIELD],
-      ['arg', LABELED_ARG_FIELD]
-    )
-  }
-
-  return pathToCall
-}
-
 export function getAxisExpressionAndIndex(
   axisOrEdge: 'Axis' | 'Edge',
   axis: string | undefined,

src/lang/queryAst.test.ts

@@ -15,6 +15,7 @@ import {
   findAllPreviousVariables,
   findUsesOfTagInPipe,
   getNodeFromPath,
+  getVariableExprsFromSelection,
   hasSketchPipeBeenExtruded,
   isCursorInFunctionDefinition,
   isNodeSafeToReplace,
@@ -27,7 +28,7 @@ import { topLevelRange } from '@src/lang/util'
 import type { Identifier, PathToNode } from '@src/lang/wasm'
 import { assertParse, recast } from '@src/lang/wasm'
 import { initPromise } from '@src/lang/wasmUtils'
-import { type Selection } from '@src/lib/selections'
+import type { Selections, Selection } from '@src/lib/selections'
 import { enginelessExecutor } from '@src/lib/testHelpers'
 import { err } from '@src/lib/trap'
 
@@ -778,3 +779,184 @@ describe('Testing specific sketch getNodeFromPath workflow', () => {
     expect(result).toEqual(false)
   })
 })
+
+describe('Testing getVariableExprsFromSelection', () => {
+  it('should find the variable expr in a simple profile selection', async () => {
+    const circleProfileInVar = `sketch001 = startSketchOn(XY)
+profile001 = circle(sketch001, center = [0, 0], radius = 1)
+`
+    const ast = assertParse(circleProfileInVar)
+    const { artifactGraph } = await enginelessExecutor(ast)
+    const artifact = artifactGraph.values().find((a) => a.type === 'path')
+    if (!artifact) {
+      throw new Error('Artifact not found in the graph')
+    }
+    const selections: Selections = {
+      graphSelections: [
+        {
+          codeRef: artifact.codeRef,
+          artifact,
+        },
+      ],
+      otherSelections: [],
+    }
+    const variableExprs = getVariableExprsFromSelection(selections, ast)
+    if (err(variableExprs)) throw variableExprs
+
+    expect(variableExprs.exprs).toHaveLength(1)
+    if (variableExprs.exprs[0].type !== 'Name') {
+      throw new Error(`Expected Name, got ${variableExprs.exprs[0].type}`)
+    }
+
+    expect(variableExprs.exprs[0].name.name).toEqual('profile001')
+
+    expect(variableExprs.paths).toHaveLength(1)
+    expect(variableExprs.paths[0]).toEqual([
+      ['body', ''],
+      [1, 'index'],
+      ['declaration', 'VariableDeclaration'],
+      ['init', ''],
+    ])
+  })
+
+  it('should return the pipe substitution symbol in a variable-less simple profile selection', async () => {
+    const circleProfileInVar = `startSketchOn(XY)
+  |> circle(center = [0, 0], radius = 1)
+`
+    const ast = assertParse(circleProfileInVar)
+    const { artifactGraph } = await enginelessExecutor(ast)
+    const artifact = artifactGraph.values().find((a) => a.type === 'path')
+    if (!artifact) {
+      throw new Error('Artifact not found in the graph')
+    }
+    const selections: Selections = {
+      graphSelections: [
+        {
+          codeRef: artifact.codeRef,
+          artifact,
+        },
+      ],
+      otherSelections: [],
+    }
+    const variableExprs = getVariableExprsFromSelection(selections, ast)
+    if (err(variableExprs)) throw variableExprs
+
+    expect(variableExprs.exprs).toHaveLength(1)
+    expect(variableExprs.exprs[0].type).toEqual('PipeSubstitution')
+
+    expect(variableExprs.paths).toHaveLength(1)
+    expect(variableExprs.paths[0]).toEqual([
+      ['body', ''],
+      [0, 'index'],
+      ['expression', 'ExpressionStatement'],
+      ['body', 'PipeExpression'],
+      [1, 'index'],
+    ])
+  })
+
+  it('should find the variable exprs in a multi profile selection ', async () => {
+    const circleProfileInVar = `sketch001 = startSketchOn(XY)
+profile001 = circle(sketch001, center = [0, 0], radius = 1)
+profile002 = circle(sketch001, center = [2, 2], radius = 1)
+`
+    const ast = assertParse(circleProfileInVar)
+    const { artifactGraph } = await enginelessExecutor(ast)
+    const artifacts = [...artifactGraph.values()].filter(
+      (a) => a.type === 'path'
+    )
+    if (!artifacts || artifacts.length !== 2) {
+      throw new Error('Artifact not found in the graph')
+    }
+    const selections: Selections = {
+      graphSelections: artifacts.map((artifact) => {
+        return {
+          codeRef: artifact.codeRef,
+          artifact,
+        }
+      }),
+      otherSelections: [],
+    }
+    const variableExprs = getVariableExprsFromSelection(selections, ast)
+    if (err(variableExprs)) throw variableExprs
+
+    expect(variableExprs.exprs).toHaveLength(2)
+    if (variableExprs.exprs[0].type !== 'Name') {
+      throw new Error(`Expected Name, got ${variableExprs.exprs[0].type}`)
+    }
+
+    if (variableExprs.exprs[1].type !== 'Name') {
+      throw new Error(`Expected Name, got ${variableExprs.exprs[1].type}`)
+    }
+
+    expect(variableExprs.exprs[0].name.name).toEqual('profile001')
+    expect(variableExprs.exprs[1].name.name).toEqual('profile002')
+
+    expect(variableExprs.paths).toHaveLength(2)
+    expect(variableExprs.paths[0]).toEqual([
+      ['body', ''],
+      [1, 'index'],
+      ['declaration', 'VariableDeclaration'],
+      ['init', ''],
+    ])
+    expect(variableExprs.paths[1]).toEqual([
+      ['body', ''],
+      [2, 'index'],
+      ['declaration', 'VariableDeclaration'],
+      ['init', ''],
+    ])
+  })
+
+  it('should return the pipe substitution symbol and a variable name in a complex multi profile selection', async () => {
+    const circleProfileInVar = `startSketchOn(XY)
+  |> circle(center = [0, 0], radius = 1)
+profile002 = circle(sketch001, center = [2, 2], radius = 1)
+`
+    const ast = assertParse(circleProfileInVar)
+    const { artifactGraph } = await enginelessExecutor(ast)
+    const artifacts = [...artifactGraph.values()].filter(
+      (a) => a.type === 'path'
+    )
+    if (!artifacts || artifacts.length !== 2) {
+      throw new Error('Artifact not found in the graph')
+    }
+    const selections: Selections = {
+      graphSelections: artifacts.map((artifact) => {
+        return {
+          codeRef: artifact.codeRef,
+          artifact,
+        }
+      }),
+      otherSelections: [],
+    }
+    const variableExprs = getVariableExprsFromSelection(selections, ast)
+    if (err(variableExprs)) throw variableExprs
+
+    expect(variableExprs.exprs).toHaveLength(2)
+    if (variableExprs.exprs[0].type !== 'PipeSubstitution') {
+      throw new Error(
+        `Expected PipeSubstitution, got ${variableExprs.exprs[0].type}`
+      )
+    }
+
+    if (variableExprs.exprs[1].type !== 'Name') {
+      throw new Error(`Expected Name, got ${variableExprs.exprs[1].type}`)
+    }
+
+    expect(variableExprs.exprs[1].name.name).toEqual('profile002')
+
+    expect(variableExprs.paths).toHaveLength(2)
+    expect(variableExprs.paths[0]).toEqual([
+      ['body', ''],
+      [0, 'index'],
+      ['expression', 'ExpressionStatement'],
+      ['body', 'PipeExpression'],
+      [1, 'index'],
+    ])
+    expect(variableExprs.paths[1]).toEqual([
+      ['body', ''],
+      [1, 'index'],
+      ['declaration', 'VariableDeclaration'],
+      ['init', ''],
+    ])
+  })
+})

src/lang/queryAst.ts

@@ -1042,19 +1042,21 @@ export const valueOrVariable = (variable: KclCommandValue) => {
 
 // Go from a selection of sketches to a list of KCL expressions that
 // can be used to create KCL sweep call declarations.
-export function getSketchExprsFromSelection(
+export function getVariableExprsFromSelection(
   selection: Selections,
   ast: Node<Program>,
   nodeToEdit?: PathToNode
-): Error | Expr[] {
-  const sketches: Expr[] = selection.graphSelections.flatMap((s) => {
+): Error | { exprs: Expr[]; paths: PathToNode[] } {
+  const paths: PathToNode[] = []
+  const exprs: Expr[] = []
+  for (const s of selection.graphSelections) {
     const sketchVariable = getNodeFromPath<VariableDeclarator>(
       ast,
       s?.codeRef.pathToNode,
       'VariableDeclarator'
     )
     if (err(sketchVariable)) {
-      return []
+      continue
     }
 
     if (sketchVariable.node.id) {
@@ -1071,22 +1073,27 @@ export function getSketchExprsFromSelection(
           name === result.node.id.name
         ) {
           // Pointing to same variable case
-          return createPipeSubstitution()
+          paths.push(nodeToEdit)
+          exprs.push(createPipeSubstitution())
+          continue
         }
       }
       // Pointing to different variable case
-      return createLocalName(name)
-    } else {
-      // No variable case
-      return createPipeSubstitution()
+      paths.push(sketchVariable.deepPath)
+      exprs.push(createLocalName(name))
+      continue
     }
-  })
 
-  if (sketches.length === 0) {
+    // No variable case
+    paths.push(sketchVariable.deepPath)
+    exprs.push(createPipeSubstitution())
+  }
+
+  if (exprs.length === 0) {
     return new Error("Couldn't map selections to program references")
   }
 
-  return sketches
+  return { exprs, paths }
 }
 
 // Go from the sketches argument in a KCL sweep call declaration

src/machines/modelingMachine.ts

@@ -72,7 +72,7 @@ import {
   addRevolve,
   addSweep,
   getAxisExpressionAndIndex,
-} from '@src/lang/modifyAst/addSweep'
+} from '@src/lang/modifyAst/sweeps'
 import {
   applyIntersectFromTargetOperatorSelections,
   applySubtractFromTargetOperatorSelections,