Add edit flow for Helix (#5394)

* WIP: Add edit flow for Helix
Fixes #5392

* Clean upp

* Add e2e test step

e2e/playwright/point-click.spec.ts

@@ -1125,7 +1125,49 @@ openSketch = startSketchOn('XY')
       await scene.expectPixelColor([250, 250, 250], testPoint, 15)
     })
 
-    await test.step('Delete offset plane via feature tree selection', async () => {
+    await test.step(`Edit helix through the feature tree`, async () => {
+      await editor.closePane()
+      const operationButton = await toolbar.getFeatureTreeOperation('Helix', 0)
+      await operationButton.dblclick()
+      const initialInput = '5'
+      const newInput = '50'
+      await cmdBar.expectState({
+        commandName: 'Helix',
+        stage: 'arguments',
+        currentArgKey: 'length',
+        currentArgValue: initialInput,
+        headerArguments: {
+          AngleStart: '360',
+          Axis: 'X',
+          CounterClockWise: '',
+          Length: initialInput,
+          Radius: '5',
+          Revolutions: '1',
+        },
+        highlightedHeaderArg: 'length',
+      })
+      await expect(cmdBar.currentArgumentInput).toBeVisible()
+      await cmdBar.currentArgumentInput.locator('.cm-content').fill(newInput)
+      await cmdBar.progressCmdBar()
+      await cmdBar.expectState({
+        stage: 'review',
+        headerArguments: {
+          AngleStart: '360',
+          Axis: 'X',
+          CounterClockWise: '',
+          Length: newInput,
+          Radius: '5',
+          Revolutions: '1',
+        },
+        commandName: 'Helix',
+      })
+      await cmdBar.progressCmdBar()
+      await toolbar.closeFeatureTreePane()
+      await editor.openPane()
+      await editor.expectEditor.toContain('length = ' + newInput)
+    })
+
+    await test.step('Delete helix via feature tree selection', async () => {
       await editor.closePane()
       const operationButton = await toolbar.getFeatureTreeOperation('Helix', 0)
       await operationButton.click({ button: 'left' })

src/lang/modifyAst.ts

@@ -717,6 +717,8 @@ export function addHelix({
   radius,
   axis,
   length,
+  insertIndex,
+  variableName,
 }: {
   node: Node<Program>
   revolutions: Expr
@@ -725,9 +727,12 @@ export function addHelix({
   radius: Expr
   axis: string
   length: Expr
+  insertIndex?: number
+  variableName?: string
 }): { modifiedAst: Node<Program>; pathToNode: PathToNode } {
   const modifiedAst = structuredClone(node)
-  const name = findUniqueName(node, KCL_DEFAULT_CONSTANT_PREFIXES.HELIX)
+  const name =
+    variableName ?? findUniqueName(node, KCL_DEFAULT_CONSTANT_PREFIXES.HELIX)
   const variable = createVariableDeclaration(
     name,
     createCallExpressionStdLibKw(
@@ -744,12 +749,20 @@ export function addHelix({
     )
   )
 
-  // TODO: figure out smart insertion than just appending at the end
+  const insertAt =
+    insertIndex !== undefined
+      ? insertIndex
+      : modifiedAst.body.length
+      ? modifiedAst.body.length
+      : 0
+
+  modifiedAst.body.length
+    ? modifiedAst.body.splice(insertAt, 0, variable)
+    : modifiedAst.body.push(variable)
   const argIndex = 0
-  modifiedAst.body.push(variable)
   const pathToNode: PathToNode = [
     ['body', ''],
-    [modifiedAst.body.length - 1, 'index'],
+    [insertAt, 'index'],
     ['declaration', 'VariableDeclaration'],
     ['init', 'VariableDeclarator'],
     ['arguments', 'CallExpressionKw'],

src/lib/commandBarConfigs/modelingCommandConfig.ts

@@ -83,6 +83,9 @@ export type ModelingCommandSchema = {
     distance: KclCommandValue
   }
   Helix: {
+    // Enables editing workflow
+    nodeToEdit?: PathToNode
+    // KCL stdlib arguments
     revolutions: KclCommandValue
     angleStart: KclCommandValue
     counterClockWise: boolean
@@ -472,6 +475,13 @@ export const modelingMachineCommandConfig: StateMachineCommandSetConfig<
     status: 'development',
     needsReview: true,
     args: {
+      nodeToEdit: {
+        description:
+          'Path to the node in the AST to edit. Never shown to the user.',
+        skip: true,
+        inputType: 'text',
+        required: false,
+      },
       revolutions: {
         inputType: 'kcl',
         defaultValue: '1',
@@ -487,9 +497,10 @@ export const modelingMachineCommandConfig: StateMachineCommandSetConfig<
       counterClockWise: {
         inputType: 'options',
         required: true,
+        defaultValue: false,
         options: [
-          { name: 'True', isCurrent: false, value: true },
-          { name: 'False', isCurrent: true, value: false },
+          { name: 'False', value: false },
+          { name: 'True', value: true },
         ],
       },
       radius: {
@@ -500,10 +511,11 @@ export const modelingMachineCommandConfig: StateMachineCommandSetConfig<
       axis: {
         inputType: 'options',
         required: true,
+        defaultValue: 'X',
         options: [
-          { name: 'X Axis', isCurrent: true, value: 'X' },
-          { name: 'Y Axis', isCurrent: false, value: 'Y' },
-          { name: 'Z Axis', isCurrent: false, value: 'Z' },
+          { name: 'X Axis', value: 'X' },
+          { name: 'Y Axis', value: 'Y' },
+          { name: 'Z Axis', value: 'Z' },
         ],
       },
       length: {

src/lib/operations.ts

@@ -191,6 +191,114 @@ const prepareToEditOffsetPlane: PrepareToEditCallback = async ({
   }
 }
 
+const prepareToEditHelix: PrepareToEditCallback = async ({ operation }) => {
+  const baseCommand = {
+    name: 'Helix',
+    groupId: 'modeling',
+  }
+  if (operation.type !== 'StdLibCall' || !operation.labeledArgs) {
+    return baseCommand
+  }
+
+  // TODO: find a way to loop over the arguments while keeping it safe
+  // revolutions kcl arg
+  if (
+    !('revolutions' in operation.labeledArgs) ||
+    !operation.labeledArgs.revolutions
+  )
+    return baseCommand
+  const revolutions = await stringToKclExpression(
+    codeManager.code.slice(
+      operation.labeledArgs.revolutions.sourceRange[0],
+      operation.labeledArgs.revolutions.sourceRange[1]
+    ),
+    {}
+  )
+  if (err(revolutions) || 'errors' in revolutions) return baseCommand
+
+  // angleStart kcl arg
+  if (
+    !('angleStart' in operation.labeledArgs) ||
+    !operation.labeledArgs.angleStart
+  )
+    return baseCommand
+  const angleStart = await stringToKclExpression(
+    codeManager.code.slice(
+      operation.labeledArgs.angleStart.sourceRange[0],
+      operation.labeledArgs.angleStart.sourceRange[1]
+    ),
+    {}
+  )
+  if (err(angleStart) || 'errors' in angleStart) return baseCommand
+
+  // counterClockWise options boolean arg
+  if (
+    !('counterClockWise' in operation.labeledArgs) ||
+    !operation.labeledArgs.counterClockWise
+  )
+    return baseCommand
+  const counterClockWise =
+    codeManager.code.slice(
+      operation.labeledArgs.counterClockWise.sourceRange[0],
+      operation.labeledArgs.counterClockWise.sourceRange[1]
+    ) === 'true'
+
+  // radius kcl arg
+  if (!('radius' in operation.labeledArgs) || !operation.labeledArgs.radius)
+    return baseCommand
+  const radius = await stringToKclExpression(
+    codeManager.code.slice(
+      operation.labeledArgs.radius.sourceRange[0],
+      operation.labeledArgs.radius.sourceRange[1]
+    ),
+    {}
+  )
+  if (err(radius) || 'errors' in radius) return baseCommand
+
+  // axis options string arg
+  if (!('axis' in operation.labeledArgs) || !operation.labeledArgs.axis)
+    return baseCommand
+  const axis = codeManager.code
+    .slice(
+      operation.labeledArgs.axis.sourceRange[0],
+      operation.labeledArgs.axis.sourceRange[1]
+    )
+    .replaceAll("'", '') // TODO: fix this crap
+
+  // length kcl arg
+  if (!('length' in operation.labeledArgs) || !operation.labeledArgs.length)
+    return baseCommand
+  const length = await stringToKclExpression(
+    codeManager.code.slice(
+      operation.labeledArgs.length.sourceRange[0],
+      operation.labeledArgs.length.sourceRange[1]
+    ),
+    {}
+  )
+  if (err(length) || 'errors' in length) return baseCommand
+
+  // Assemble the default argument values for the Offset Plane command,
+  // with `nodeToEdit` set, which will let the Offset Plane actor know
+  // to edit the node that corresponds to the StdLibCall.
+  const argDefaultValues: ModelingCommandSchema['Helix'] = {
+    revolutions,
+    angleStart,
+    counterClockWise,
+    radius,
+    axis,
+    length,
+    nodeToEdit: getNodePathFromSourceRange(
+      kclManager.ast,
+      sourceRangeFromRust(operation.sourceRange)
+    ),
+  }
+
+  return {
+    ...baseCommand,
+    argDefaultValues,
+  }
+}
+
 /**
  * A map of standard library calls to their corresponding information
  * for use in the feature tree UI.
@@ -214,6 +322,7 @@ export const stdLibMap: Record<string, StdLibCallInfo> = {
   helix: {
     label: 'Helix',
     icon: 'helix',
+    prepareToEdit: prepareToEditHelix,
   },
   hole: {
     label: 'Hole',

src/machines/modelingMachine.ts

@@ -1801,8 +1801,32 @@ export const modelingMachine = setup({
           radius,
           axis,
           length,
+          nodeToEdit,
         } = input
 
+        let opInsertIndex: number | undefined = undefined
+        let opVariableName: string | undefined = undefined
+
+        // If this is an edit flow, first we're going to remove the old one
+        if (nodeToEdit && typeof nodeToEdit[1][0] === 'number') {
+          // Extract the old name from the node to edit
+          const oldNode = getNodeFromPath<VariableDeclaration>(
+            ast,
+            nodeToEdit,
+            'VariableDeclaration'
+          )
+          if (err(oldNode)) {
+            console.error('Error extracting plane name')
+          } else {
+            opVariableName = oldNode.node.declaration.id.name
+          }
+
+          const newBody = [...ast.body]
+          newBody.splice(nodeToEdit[1][0], 1)
+          ast.body = newBody
+          opInsertIndex = nodeToEdit[1][0]
+        }
+
         for (const variable of [revolutions, angleStart, radius, length]) {
           // Insert the variable if it exists
           if (
@@ -1833,6 +1857,8 @@ export const modelingMachine = setup({
           radius: valueOrVariable(radius),
           axis,
           length: valueOrVariable(length),
+          insertIndex: opInsertIndex,
+          variableName: opVariableName,
         })
 
         const updateAstResult = await kclManager.updateAst(