diff --git a/src/lang/modifyAst.test.ts b/src/lang/modifyAst.test.ts index a420704ad..d0e74b0ef 100644 --- a/src/lang/modifyAst.test.ts +++ b/src/lang/modifyAst.test.ts @@ -2,7 +2,9 @@ import type { Node } from '@rust/kcl-lib/bindings/Node' import { createArrayExpression, + createCallExpressionStdLibKw, createIdentifier, + createLabeledArg, createLiteral, createLiteralMaybeSuffix, createLocalName, @@ -19,6 +21,7 @@ import { createVariableExpressionsArray, deleteSegmentFromPipeExpression, moveValueIntoNewVariable, + setCallInAst, sketchOnExtrudedFace, splitPipedProfile, } from '@src/lang/modifyAst' @@ -1037,3 +1040,45 @@ extrude001 = extrude(profile001, length = 123) expect(node.node.end).toEqual(startOfKwargIndex + 3) }) }) + +describe('Testing setCallInAst', () => { + it('should push an extrude call with variable on variable profile', () => { + const code = `sketch001 = startSketchOn(XY) +profile001 = circle(sketch001, center = [0, 0], radius = 1) +` + const ast = assertParse(code) + const exprs = createVariableExpressionsArray([ + createLocalName('profile001'), + ]) + const call = createCallExpressionStdLibKw('extrude', exprs, [ + createLabeledArg('length', createLiteral(5)), + ]) + const pathToNode = setCallInAst(ast, call) + if (err(pathToNode)) { + throw pathToNode + } + const newCode = recast(ast) + expect(newCode).toContain(code) + expect(newCode).toContain(`extrude001 = extrude(profile001, length = 5)`) + }) + + it('should push an extrude call with variable on variable profile', () => { + const code = `startSketchOn(XY) + |> circle(sketch001, center = [0, 0], radius = 1) +` + const ast = assertParse(code) + const exprs = createVariableExpressionsArray([ + createLocalName('profile001'), + ]) + const call = createCallExpressionStdLibKw('extrude', exprs, [ + createLabeledArg('length', createLiteral(5)), + ]) + const pathToNode = setCallInAst(ast, call) + if (err(pathToNode)) { + throw pathToNode + } + const newCode = recast(ast) + expect(newCode).toContain(code) + expect(newCode).toContain(`|> extrude(length = 5)`) + }) +}) diff --git a/src/lang/modifyAst.ts b/src/lang/modifyAst.ts index da64d8d7c..f12d32664 100644 --- a/src/lang/modifyAst.ts +++ b/src/lang/modifyAst.ts @@ -1247,3 +1247,45 @@ export function createPathToNodeForLastVariable( return pathToCall } + +export function setCallInAst( + ast: Node, + call: Node, + nodeToEdit?: PathToNode, + lastPathToNode?: PathToNode +): Error | PathToNode { + let pathToNode: PathToNode | undefined + if (nodeToEdit) { + const result = getNodeFromPath( + ast, + nodeToEdit, + 'CallExpressionKw' + ) + if (err(result)) { + return result + } + + Object.assign(result.node, call) + pathToNode = nodeToEdit + } else { + if (!call.unlabeled && lastPathToNode) { + const pipe = getNodeFromPath( + ast, + lastPathToNode, + 'PipeExpression' + ) + if (err(pipe)) { + return pipe + } + pipe.node.body.push(call) + pathToNode = lastPathToNode + } else { + const name = findUniqueName(ast, call.callee.name.name) + const declaration = createVariableDeclaration(name, call) + ast.body.push(declaration) + pathToNode = createPathToNodeForLastVariable(ast) + } + } + + return pathToNode +} diff --git a/src/lang/modifyAst/addSweep.ts b/src/lang/modifyAst/addSweep.ts index 3f39c2c50..00a6cc1a3 100644 --- a/src/lang/modifyAst/addSweep.ts +++ b/src/lang/modifyAst/addSweep.ts @@ -5,13 +5,11 @@ import { createLabeledArg, createLiteral, createLocalName, - createVariableDeclaration, - findUniqueName, } from '@src/lang/create' import { - createPathToNodeForLastVariable, createVariableExpressionsArray, insertVariableAndOffsetPathToNode, + setCallInAst, } from '@src/lang/modifyAst' import { getEdgeTagCall, @@ -23,15 +21,8 @@ import { valueOrVariable, } from '@src/lang/queryAst' import { getNodePathFromSourceRange } from '@src/lang/queryAstNodePathUtils' -import type { - CallExpressionKw, - PathToNode, - PipeExpression, - Program, - VariableDeclaration, -} from '@src/lang/wasm' +import type { PathToNode, Program, VariableDeclaration } from '@src/lang/wasm' import type { KclCommandValue } from '@src/lib/commandTypes' -import { KCL_DEFAULT_CONSTANT_PREFIXES } from '@src/lib/constants' import type { Selections } from '@src/lib/selections' import { err } from '@src/lib/trap' @@ -116,42 +107,10 @@ export function addExtrude({ // 3. If edit, we assign the new function call declaration to the existing node, // otherwise just push to the end - let pathToNode: PathToNode | undefined - if (nodeToEdit) { - const result = getNodeFromPath( - modifiedAst, - nodeToEdit, - 'CallExpressionKw' - ) - if (err(result)) { - return result - } - - Object.assign(result.node, call) - pathToNode = nodeToEdit - } else { - const lastPathToNode: PathToNode | undefined = - variableExpressions.paths.pop() - if (sketchesExpr === null && lastPathToNode) { - const pipe = getNodeFromPath( - modifiedAst, - lastPathToNode, - 'PipeExpression' - ) - if (err(pipe)) { - return pipe - } - pipe.node.body.push(call) - pathToNode = lastPathToNode - } else { - const name = findUniqueName( - modifiedAst, - KCL_DEFAULT_CONSTANT_PREFIXES.EXTRUDE - ) - const declaration = createVariableDeclaration(name, call) - modifiedAst.body.push(declaration) - pathToNode = createPathToNodeForLastVariable(modifiedAst) - } + const lastPath = variableExpressions.paths.pop() // TODO: check if this is correct + const pathToNode = setCallInAst(modifiedAst, call, nodeToEdit, lastPath) + if (err(pathToNode)) { + return pathToNode } return { @@ -222,41 +181,10 @@ export function addSweep({ // 3. If edit, we assign the new function call declaration to the existing node, // otherwise just push to the end - let pathToNode: PathToNode | undefined - if (nodeToEdit) { - const result = getNodeFromPath( - modifiedAst, - nodeToEdit, - 'CallExpressionKw' - ) - if (err(result)) { - return result - } - - Object.assign(result.node, call) - pathToNode = nodeToEdit - } else { - const lastPathToNode: PathToNode | undefined = variableExprs.paths.pop() - if (sketchesExpr === null && lastPathToNode) { - const pipe = getNodeFromPath( - modifiedAst, - lastPathToNode, - 'PipeExpression' - ) - if (err(pipe)) { - return pipe - } - pipe.node.body.push(call) - pathToNode = lastPathToNode - } else { - const name = findUniqueName( - modifiedAst, - KCL_DEFAULT_CONSTANT_PREFIXES.SWEEP - ) - const declaration = createVariableDeclaration(name, call) - modifiedAst.body.push(declaration) - pathToNode = createPathToNodeForLastVariable(modifiedAst) - } + const lastPath = variableExprs.paths.pop() // TODO: check if this is correct + const pathToNode = setCallInAst(modifiedAst, call, nodeToEdit, lastPath) + if (err(pathToNode)) { + return pathToNode } return { @@ -312,42 +240,10 @@ export function addLoft({ // 3. If edit, we assign the new function call declaration to the existing node, // otherwise just push to the end - let pathToNode: PathToNode | undefined - if (nodeToEdit) { - const result = getNodeFromPath( - modifiedAst, - nodeToEdit, - 'CallExpressionKw' - ) - if (err(result)) { - return result - } - - Object.assign(result.node, call) - pathToNode = nodeToEdit - } else { - const lastPathToNode: PathToNode | undefined = variableExprs.paths.pop() - if (sketchesExpr === null && lastPathToNode) { - const pipe = getNodeFromPath( - modifiedAst, - lastPathToNode, - 'PipeExpression' - ) - if (err(pipe)) { - return pipe - } - pipe.node.body.push(call) - pathToNode = lastPathToNode - } else { - const name = findUniqueName( - modifiedAst, - KCL_DEFAULT_CONSTANT_PREFIXES.LOFT - ) - const declaration = createVariableDeclaration(name, call) - modifiedAst.body.push(declaration) - const toFirstKwarg = !!vDegree - pathToNode = createPathToNodeForLastVariable(modifiedAst, toFirstKwarg) - } + const lastPath = variableExprs.paths.pop() // TODO: check if this is correct + const pathToNode = setCallInAst(modifiedAst, call, nodeToEdit, lastPath) + if (err(pathToNode)) { + return pathToNode } return { @@ -447,41 +343,10 @@ export function addRevolve({ // 3. If edit, we assign the new function call declaration to the existing node, // otherwise just push to the end - let pathToNode: PathToNode | undefined - if (nodeToEdit) { - const result = getNodeFromPath( - modifiedAst, - nodeToEdit, - 'CallExpressionKw' - ) - if (err(result)) { - return result - } - - Object.assign(result.node, call) - pathToNode = nodeToEdit - } else { - const lastPathToNode: PathToNode | undefined = variableExprs.paths.pop() - if (sketchesExpr === null && lastPathToNode) { - const pipe = getNodeFromPath( - modifiedAst, - lastPathToNode, - 'PipeExpression' - ) - if (err(pipe)) { - return pipe - } - pipe.node.body.push(call) - pathToNode = lastPathToNode - } else { - const name = findUniqueName( - modifiedAst, - KCL_DEFAULT_CONSTANT_PREFIXES.REVOLVE - ) - const declaration = createVariableDeclaration(name, call) - modifiedAst.body.push(declaration) - pathToNode = createPathToNodeForLastVariable(modifiedAst) - } + const lastPath = variableExprs.paths.pop() // TODO: check if this is correct + const pathToNode = setCallInAst(modifiedAst, call, nodeToEdit, lastPath) + if (err(pathToNode)) { + return pathToNode } return {