diff --git a/src/lang/modifyAst/addRevolve.ts b/src/lang/modifyAst/addRevolve.ts index 2ae982a51..8fd5fc227 100644 --- a/src/lang/modifyAst/addRevolve.ts +++ b/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') - return new Error('expected sketchIndexInBody to be a number') - clonedAst.body.splice(sketchIndexInBody + 1, 0, VariableDeclaration) + let insertIndex = sketchIndexInBody + + 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']],