Fix: inserting revolve command accounts for axis dependency (#5231)

fix: inserting revolve with axis dependency

src/lang/modifyAst/addRevolve.ts

@@ -46,6 +46,7 @@ export function revolveSketch(
   if (err(sketchNode)) return sketchNode
 
   let generatedAxis
+  let axisDeclaration: PathToNode | null = null
 
   if (axisOrEdge === 'Edge') {
     const pathToAxisSelection = getNodePathFromSourceRange(
@@ -70,6 +71,13 @@ export function revolveSketch(
     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)
   }
@@ -139,18 +147,34 @@ export function revolveSketch(
   const sketchIndexInPathToNode =
     sketchPathToDecleration.findIndex((a) => a[0] === 'body') + 1
   const sketchIndexInBody = sketchPathToDecleration[sketchIndexInPathToNode][0]
-  if (typeof sketchIndexInBody !== 'number')
+  let insertIndex = sketchIndexInBody
-    return new Error('expected sketchIndexInBody to be a number')
+
-  clonedAst.body.splice(sketchIndexInBody + 1, 0, VariableDeclaration)
+  if (typeof insertIndex !== 'number')
+    return new Error('expected insertIndex to be a number')
+
+  // If an axis was selected in KCL, find the max index to insert the revolve command
+  if (axisDeclaration) {
+    const axisIndexInPathToNode =
+      axisDeclaration.findIndex((a) => a[0] === 'body') + 1
+    const axisIndex = axisDeclaration[axisIndexInPathToNode][0]
+
+    if (typeof axisIndex !== 'number')
+      return new Error('expected axisIndex to be a number')
+
+    insertIndex = Math.max(insertIndex, axisIndex)
+  }
+
+  clonedAst.body.splice(insertIndex + 1, 0, VariableDeclaration)
 
   const pathToRevolveArg: PathToNode = [
     ['body', ''],
-    [sketchIndexInBody + 1, 'index'],
+    [insertIndex + 1, 'index'],
     ['declaration', 'VariableDeclaration'],
     ['init', 'VariableDeclarator'],
     ['arguments', 'CallExpression'],
     [0, 'index'],
   ]
+
   return {
     modifiedAst: clonedAst,
     pathToSketchNode: [...pathToSketchNode.slice(0, -1), [-1, 'index']],