WIP: Add edit flow for Revolve

Fixes #5504

src/lang/modifyAst/addRevolve.ts

@@ -35,10 +35,11 @@ import { kclManager } from 'lib/singletons'
 export function revolveSketch(
   ast: Node<Program>,
   pathToSketchNode: PathToNode,
+  variableName: string | undefined,
   angle: Expr = createLiteral(4),
   axisOrEdge: string,
-  axis: string,
-  edge: Selections,
+  axis: string | undefined,
+  edge: Selections | undefined,
   artifactGraph: ArtifactGraph,
   artifact?: Artifact
 ):
@@ -62,7 +63,7 @@ export function revolveSketch(
   let generatedAxis
   let axisDeclaration: PathToNode | null = null
 
-  if (axisOrEdge === 'Edge') {
+  if (axisOrEdge === 'Edge' && edge) {
     const pathToAxisSelection = getNodePathFromSourceRange(
       clonedAst,
       edge.graphSelections[0]?.codeRef.range
@@ -92,7 +93,7 @@ export function revolveSketch(
     ) {
       axisDeclaration = axisSelection.codeRef.pathToNode
     }
-  } else {
+  } else if (axisOrEdge === 'Axis' && axis) {
     generatedAxis = createLiteral(axis)
   }
 
@@ -114,7 +115,9 @@ export function revolveSketch(
 
   // We're not creating a pipe expression,
   // but rather a separate constant for the extrusion
-  const name = findUniqueName(clonedAst, KCL_DEFAULT_CONSTANT_PREFIXES.REVOLVE)
+  const name =
+    variableName ??
+    findUniqueName(clonedAst, KCL_DEFAULT_CONSTANT_PREFIXES.REVOLVE)
   const VariableDeclaration = createVariableDeclaration(name, revolveCall)
   const lastSketchNodePath =
     orderedSketchNodePaths[orderedSketchNodePaths.length - 1]

src/lib/commandBarConfigs/modelingCommandConfig.ts

@@ -73,11 +73,14 @@ export type ModelingCommandSchema = {
     thickness: KclCommandValue
   }
   Revolve: {
+    // Enables editing workflow
+    nodeToEdit?: PathToNode
+    // KCL stdlib arguments
     selection: Selections
     angle: KclCommandValue
     axisOrEdge: 'Axis' | 'Edge'
-    axis: string
-    edge: Selections
+    axis?: string
+    edge?: Selections
   }
   Fillet: {
     selection: Selections
@@ -445,6 +448,13 @@ export const modelingMachineCommandConfig: StateMachineCommandSetConfig<
     icon: 'revolve',
     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,
+      },
       selection: {
         inputType: 'selection',
         selectionTypes: ['solid2d', 'segment'],
@@ -473,6 +483,7 @@ export const modelingMachineCommandConfig: StateMachineCommandSetConfig<
           { name: 'X Axis', isCurrent: true, value: 'X' },
           { name: 'Y Axis', isCurrent: false, value: 'Y' },
         ],
+        hidden: (context) => Boolean(context.argumentsToSubmit.nodeToEdit),
       },
       edge: {
         required: (commandContext) =>

src/lib/operations.ts

@@ -3,6 +3,7 @@ import {
   Artifact,
   getArtifactOfTypes,
   getCapCodeRef,
+  getSweepEdgeCodeRef,
 } from 'lang/std/artifactGraph'
 import { Operation } from '@rust/kcl-lib/bindings/Operation'
 import { codeManager, engineCommandManager, kclManager } from './singletons'
@@ -552,6 +553,166 @@ const prepareToEditHelix: PrepareToEditCallback = async ({ operation }) => {
   }
 }
 
+const prepareToEditRevolve: PrepareToEditCallback = async ({
+  operation,
+  artifact,
+}) => {
+  const baseCommand = {
+    name: 'Revolve',
+    groupId: 'modeling',
+  }
+  if (
+    !artifact ||
+    !('pathId' in artifact) ||
+    operation.type !== 'StdLibCall' ||
+    !operation.labeledArgs
+  ) {
+    return baseCommand
+  }
+
+  // We have to go a little roundabout to get from the original artifact
+  // to the solid2DId that we need to pass to the Extrude command.
+  const pathArtifact = getArtifactOfTypes(
+    {
+      key: artifact.pathId,
+      types: ['path'],
+    },
+    engineCommandManager.artifactGraph
+  )
+  if (
+    err(pathArtifact) ||
+    pathArtifact.type !== 'path' ||
+    !pathArtifact.solid2dId
+  )
+    return baseCommand
+  const solid2DArtifact = getArtifactOfTypes(
+    {
+      key: pathArtifact.solid2dId,
+      types: ['solid2d'],
+    },
+    engineCommandManager.artifactGraph
+  )
+  if (err(solid2DArtifact) || solid2DArtifact.type !== 'solid2d') {
+    return baseCommand
+  }
+
+  const selection = {
+    graphSelections: [
+      {
+        artifact: solid2DArtifact,
+        codeRef: pathArtifact.codeRef,
+      },
+    ],
+    otherSelections: [],
+  }
+
+  // axis options string arg
+  if (!('axis' in operation.labeledArgs) || !operation.labeledArgs.axis) {
+    return baseCommand
+  }
+
+  const axisValue = operation.labeledArgs.axis.value
+  let axisOrEdge: 'Axis' | 'Edge' | undefined
+  let axis: string | undefined
+  let edge: Selections | undefined
+  if (axisValue.type === 'String') {
+    // default axis casee
+    axisOrEdge = 'Axis'
+    axis = axisValue.value
+  } else if (axisValue.type === 'TagIdentifier' && axisValue.artifact_id) {
+    // segment case
+    axisOrEdge = 'Edge'
+    const artifact = getArtifactOfTypes(
+      {
+        key: axisValue.artifact_id,
+        types: ['segment'],
+      },
+      engineCommandManager.artifactGraph
+    )
+    if (err(artifact)) {
+      return baseCommand
+    }
+
+    edge = {
+      graphSelections: [
+        {
+          artifact,
+          codeRef: artifact.codeRef,
+        },
+      ],
+      otherSelections: [],
+    }
+  } else if (axisValue.type === 'Uuid') {
+    // sweepEdge case
+    axisOrEdge = 'Edge'
+    const artifact = getArtifactOfTypes(
+      {
+        key: axisValue.value,
+        types: ['sweepEdge'],
+      },
+      engineCommandManager.artifactGraph
+    )
+    if (err(artifact)) {
+      return baseCommand
+    }
+
+    const codeRef = getSweepEdgeCodeRef(
+      artifact,
+      engineCommandManager.artifactGraph
+    )
+    if (err(codeRef)) {
+      return baseCommand
+    }
+
+    edge = {
+      graphSelections: [
+        {
+          artifact,
+          codeRef,
+        },
+      ],
+      otherSelections: [],
+    }
+  } else {
+    return baseCommand
+  }
+
+  // angle kcl arg
+  if (!('angle' in operation.labeledArgs) || !operation.labeledArgs.angle) {
+    return baseCommand
+  }
+  const angle = await stringToKclExpression(
+    codeManager.code.slice(
+      operation.labeledArgs.angle.sourceRange[0],
+      operation.labeledArgs.angle.sourceRange[1]
+    )
+  )
+  if (err(angle) || 'errors' in angle) {
+    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['Revolve'] = {
+    axisOrEdge,
+    axis,
+    edge,
+    selection,
+    angle,
+    nodeToEdit: getNodePathFromSourceRange(
+      kclManager.ast,
+      sourceRangeFromRust(operation.sourceRange)
+    ),
+  }
+  console.log(argDefaultValues)
+
+  return {
+    ...baseCommand,
+    argDefaultValues,
+  }
+}
+
 /**
  * A map of standard library calls to their corresponding information
  * for use in the feature tree UI.
@@ -618,6 +779,7 @@ export const stdLibMap: Record<string, StdLibCallInfo> = {
   revolve: {
     label: 'Revolve',
     icon: 'revolve',
+    prepareToEdit: prepareToEditRevolve,
     supportsAppearance: true,
   },
   shell: {

src/machines/modelingMachine.ts

@@ -715,8 +715,31 @@ export const modelingMachine = setup({
       if (event.type !== 'Revolve') return
       ;(async () => {
         if (!event.data) return
-        const { selection, angle, axis, edge, axisOrEdge } = event.data
+        const { nodeToEdit, selection, angle, axis, edge, axisOrEdge } =
+          event.data
         let ast = kclManager.ast
+        let variableName: string | undefined = undefined
+
+        // If this is an edit flow, first we're going to remove the old extrusion
+        if (nodeToEdit && typeof nodeToEdit[1][0] === 'number') {
+          // Extract the plane name from the node to edit
+          const nameNode = getNodeFromPath<VariableDeclaration>(
+            ast,
+            nodeToEdit,
+            'VariableDeclaration'
+          )
+          if (err(nameNode)) {
+            console.error('Error extracting plane name')
+          } else {
+            variableName = nameNode.node.declaration.id.name
+          }
+
+          // Removing the old extrusion statement
+          const newBody = [...ast.body]
+          newBody.splice(nodeToEdit[1][0] as number, 1)
+          ast.body = newBody
+        }
+
         if (
           'variableName' in angle &&
           angle.variableName &&
@@ -736,6 +759,7 @@ export const modelingMachine = setup({
         const revolveSketchRes = revolveSketch(
           ast,
           pathToNode,
+          variableName,
           'variableName' in angle
             ? angle.variableIdentifierAst
             : angle.valueAst,