WIP: Add edge and segment selection in point-and-click Helix flow

Fixes #5393

e2e/playwright/point-click.spec.ts

@@ -1093,17 +1093,17 @@ openSketch = startSketchOn('XY')
       await toolbar.helixButton.click()
       await cmdBar.expectState({
         stage: 'arguments',
-        currentArgKey: 'revolutions',
+        currentArgKey: 'axisOrEdge',
-        currentArgValue: '1',
+        currentArgValue: '',
         headerArguments: {
           AngleStart: '',
-          Axis: '',
+          AxisOrEdge: '',
           CounterClockWise: '',
           Length: '',
           Radius: '',
           Revolutions: '',
         },
-        highlightedHeaderArg: 'revolutions',
+        highlightedHeaderArg: 'axisOrEdge',
         commandName: 'Helix',
       })
       await cmdBar.progressCmdBar()
@@ -1113,6 +1113,7 @@ openSketch = startSketchOn('XY')
       await cmdBar.progressCmdBar()
       await cmdBar.progressCmdBar()
       await cmdBar.progressCmdBar()
+      await cmdBar.progressCmdBar()
     })
 
     await test.step(`Confirm code is added to the editor, scene has changed`, async () => {

src/lang/modifyAst.ts

@@ -795,7 +795,7 @@ export function addHelix({
   angleStart: Expr
   counterClockWise: boolean
   radius: Expr
-  axis: string
+  axis: Node<Literal> | Node<Identifier | CallExpression | CallExpressionKw>
   length: Expr
   insertIndex?: number
   variableName?: string
@@ -813,7 +813,7 @@ export function addHelix({
         createLabeledArg('angleStart', angleStart),
         createLabeledArg('counterClockWise', createLiteral(counterClockWise)),
         createLabeledArg('radius', radius),
-        createLabeledArg('axis', createLiteral(axis)),
+        createLabeledArg('axis', axis),
         createLabeledArg('length', length),
       ]
     )

src/lib/commandBarConfigs/modelingCommandConfig.ts

@@ -84,12 +84,15 @@ export type ModelingCommandSchema = {
   Helix: {
     // Enables editing workflow
     nodeToEdit?: PathToNode
+    // Flow arg
+    axisOrEdge: 'Axis' | 'Edge'
     // KCL stdlib arguments
+    axis: string
+    edge: Selections
     revolutions: KclCommandValue
     angleStart: KclCommandValue
     counterClockWise: boolean
     radius: KclCommandValue
-    axis: string
     length: KclCommandValue
   }
   'change tool': {
@@ -495,6 +498,38 @@ export const modelingMachineCommandConfig: StateMachineCommandSetConfig<
         required: false,
         hidden: true,
       },
+      axisOrEdge: {
+        inputType: 'options',
+        required: true,
+        defaultValue: 'Axis',
+        options: [
+          { name: 'Axis', isCurrent: true, value: 'Axis' },
+          { name: 'Edge', isCurrent: false, value: 'Edge' },
+        ],
+        hidden: (context) => Boolean(context.argumentsToSubmit.nodeToEdit),
+      },
+      axis: {
+        inputType: 'options',
+        required: (commandContext) =>
+          ['Axis'].includes(
+            commandContext.argumentsToSubmit.axisOrEdge as string
+          ),
+        options: [
+          { name: 'X Axis', value: 'X' },
+          { name: 'Y Axis', value: 'Y' },
+          { name: 'Z Axis', value: 'Z' },
+        ],
+      },
+      edge: {
+        required: (commandContext) =>
+          ['Edge'].includes(
+            commandContext.argumentsToSubmit.axisOrEdge as string
+          ),
+        inputType: 'selection',
+        selectionTypes: ['segment', 'sweepEdge', 'edgeCutEdge'],
+        multiple: false,
+        hidden: (context) => Boolean(context.argumentsToSubmit.nodeToEdit),
+      },
       revolutions: {
         inputType: 'kcl',
         defaultValue: '1',
@@ -521,16 +556,6 @@ export const modelingMachineCommandConfig: StateMachineCommandSetConfig<
         defaultValue: KCL_DEFAULT_LENGTH,
         required: true,
       },
-      axis: {
-        inputType: 'options',
-        required: true,
-        defaultValue: 'X',
-        options: [
-          { name: 'X Axis', value: 'X' },
-          { name: 'Y Axis', value: 'Y' },
-          { name: 'Z Axis', value: 'Z' },
-        ],
-      },
       length: {
         inputType: 'kcl',
         defaultValue: KCL_DEFAULT_LENGTH,

src/lib/operations.ts

@@ -311,7 +311,10 @@ const prepareToEditOffsetPlane: PrepareToEditCallback = async ({
   }
 }
 
-const prepareToEditHelix: PrepareToEditCallback = async ({ operation }) => {
+const prepareToEditHelix: PrepareToEditCallback = async ({
+  operation,
+  artifact,
+}) => {
   const baseCommand = {
     name: 'Helix',
     groupId: 'modeling',
@@ -321,6 +324,25 @@ const prepareToEditHelix: PrepareToEditCallback = async ({ operation }) => {
   }
 
   // TODO: find a way to loop over the arguments while keeping it safe
+  // axis options string arg
+  if (!('axis' in operation.labeledArgs) || !operation.labeledArgs.axis)
+    return baseCommand
+  const axisValue = codeManager.code
+    .slice(
+      operation.labeledArgs.axis.sourceRange[0],
+      operation.labeledArgs.axis.sourceRange[1]
+    )
+    .replaceAll("'", '') // TODO: fix this crap
+  console.log('preparedToEditHelix axis', axisValue)
+  console.log('preparedToEditHelix artifact', artifact)
+  console.log('preparedToEditHelix operation', operation)
+  // TODO: replace with constants
+  const isAxis = axisValue === 'X' || axisValue === 'Y' || axisValue === 'Z'
+  const axisOrEdge = isAxis ? 'Axis' : 'Edge'
+  let axis = isAxis ? axisValue : undefined
+  // TODO: massage selection in for edge edits
+  let edge: Selections | undefined = undefined
+
   // revolutions kcl arg
   if (
     !('revolutions' in operation.labeledArgs) ||
@@ -372,16 +394,6 @@ const prepareToEditHelix: PrepareToEditCallback = async ({ operation }) => {
   )
   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
@@ -397,11 +409,13 @@ const prepareToEditHelix: PrepareToEditCallback = async ({ operation }) => {
   // with `nodeToEdit` set, which will let the Offset Plane actor know
   // to edit the node that corresponds to the StdLibCall.
   const argDefaultValues: ModelingCommandSchema['Helix'] = {
+    axisOrEdge,
+    axis,
+    edge,
     revolutions,
     angleStart,
     counterClockWise,
     radius,
-    axis,
     length,
     nodeToEdit: getNodePathFromSourceRange(
       kclManager.ast,

src/machines/modelingMachine.ts

@@ -1,4 +1,6 @@
 import {
+  CallExpression,
+  CallExpressionKw,
   Expr,
   PathToNode,
   VariableDeclaration,
@@ -57,6 +59,7 @@ import {
   ChamferParameters,
   EdgeTreatmentType,
   FilletParameters,
+  getEdgeTagCall,
   getPathToExtrudeForSegmentSelection,
   mutateAstWithTagForSketchSegment,
 } from 'lang/modifyAst/addEdgeTreatment'
@@ -1892,11 +1895,13 @@ export const modelingMachine = setup({
         // Extract inputs
         const ast = kclManager.ast
         const {
+          axisOrEdge,
+          axis,
+          edge,
           revolutions,
           angleStart,
           counterClockWise,
           radius,
-          axis,
           length,
           nodeToEdit,
         } = input
@@ -1924,6 +1929,47 @@ export const modelingMachine = setup({
           opInsertIndex = nodeToEdit[1][0]
         }
 
+        let generatedAxis
+        let axisDeclaration: PathToNode | null = null
+
+        if (axisOrEdge === 'Edge') {
+          const pathToAxisSelection = getNodePathFromSourceRange(
+            ast,
+            edge.graphSelections[0]?.codeRef.range
+          )
+          const lineNode = getNodeFromPath<CallExpression | CallExpressionKw>(
+            ast,
+            pathToAxisSelection,
+            ['CallExpression', 'CallExpressionKw']
+          )
+          if (err(lineNode)) return lineNode
+
+          const tagResult = mutateAstWithTagForSketchSegment(
+            ast,
+            pathToAxisSelection
+          )
+
+          // Have the tag whether it is already created or a new one is generated
+          if (err(tagResult)) return tagResult
+          const { tag } = tagResult
+          const axisSelection = edge?.graphSelections[0]?.artifact
+          if (!axisSelection)
+            return new Error('Generated axis selection is missing.')
+          generatedAxis = getEdgeTagCall(tag, axisSelection)
+          if (
+            axisSelection.type === 'segment' ||
+            axisSelection.type === 'path' ||
+            axisSelection.type === 'edgeCut'
+          ) {
+            axisDeclaration = axisSelection.codeRef.pathToNode
+          }
+        } else {
+          generatedAxis = createLiteral(axis)
+        }
+
+        if (!generatedAxis)
+          return new Error('Generated axis selection is missing.')
+
         for (const variable of [revolutions, angleStart, radius, length]) {
           // Insert the variable if it exists
           if (
@@ -1952,7 +1998,7 @@ export const modelingMachine = setup({
           angleStart: valueOrVariable(angleStart),
           counterClockWise,
           radius: valueOrVariable(radius),
-          axis,
+          axis: generatedAxis,
           length: valueOrVariable(length),
           insertIndex: opInsertIndex,
           variableName: opVariableName,