Merge branch 'main' into pierremtb/issue7349-transform-codemods-out-of-pipe

Fixes #7338

src/lang/modifyAst/addSweep.ts

@@ -331,7 +331,7 @@ function createSketchExpression(sketches: Expr[]) {
   return sketchesExpr
 }
 
-function createPathToNode(
+export function createPathToNode(
   modifiedAst: Node<Program>,
   toFirstKwarg = true
 ): PathToNode {

src/lang/modifyAst/setTransform.ts

@@ -6,8 +6,10 @@ import {
   createLabeledArg,
   createLocalName,
   createPipeExpression,
+  createVariableDeclaration,
+  findUniqueName,
 } from '@src/lang/create'
-import { getNodeFromPath } from '@src/lang/queryAst'
+import { getNodeFromPath, valueOrVariable } from '@src/lang/queryAst'
 import type {
   ArtifactGraph,
   CallExpressionKw,
@@ -16,6 +18,7 @@ import type {
   PathToNode,
   PipeExpression,
   Program,
+  VariableDeclaration,
   VariableDeclarator,
 } from '@src/lang/wasm'
 import { err } from '@src/lib/trap'
@@ -24,6 +27,95 @@ import {
   getLastVariable,
 } from '@src/lang/modifyAst/boolean'
 import type { Selections } from '@src/lib/selections'
+import { ImportStatement } from '@rust/kcl-lib/bindings/ImportStatement'
+import { KclCommandValue } from '@src/lib/commandTypes'
+import { insertVariableAndOffsetPathToNode } from '../modifyAst'
+import { KCL_DEFAULT_CONSTANT_PREFIXES } from '@src/lib/constants'
+import { createPathToNode } from './addSweep'
+
+export function addTranslate({
+  ast,
+  artifactGraph,
+  selection,
+  x,
+  y,
+  z,
+  nodeToEdit,
+}: {
+  ast: Node<Program>
+  artifactGraph: ArtifactGraph
+  selection: Selections
+  x: KclCommandValue
+  y: KclCommandValue
+  z: KclCommandValue
+  nodeToEdit?: PathToNode
+}):
+  | {
+      modifiedAst: Node<Program>
+      pathToNode: PathToNode
+    }
+  | Error {
+  // 1. Clone the ast so we can edit it
+  const modifiedAst = structuredClone(ast)
+
+  // 2. Prepare unlabeled and labeled arguments
+  // Map the sketches selection into a list of kcl expressions to be passed as unlabelled argument
+  const geometryNode = retrievePathToNodeFromTransformSelection(
+    selection,
+    artifactGraph,
+    ast
+  )
+  if (err(geometryNode)) {
+    return geometryNode
+  }
+
+  const geometryName = retrieveGeometryNameFromPath(ast, geometryNode)
+  if (err(geometryName)) {
+    return geometryName
+  }
+
+  const geometryExpr = createLocalName(geometryName)
+  const call = createCallExpressionStdLibKw('translate', geometryExpr, [
+    createLabeledArg('x', valueOrVariable(x)),
+    createLabeledArg('y', valueOrVariable(y)),
+    createLabeledArg('z', valueOrVariable(z)),
+  ])
+
+  // Insert variables for labeled arguments if provided
+  insertVariableAndOffsetPathToNode(x, modifiedAst, nodeToEdit)
+  insertVariableAndOffsetPathToNode(y, modifiedAst, nodeToEdit)
+  insertVariableAndOffsetPathToNode(z, modifiedAst, nodeToEdit)
+
+  // 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<CallExpressionKw>(
+      modifiedAst,
+      nodeToEdit,
+      'CallExpressionKw'
+    )
+    if (err(result)) {
+      return result
+    }
+
+    Object.assign(result.node, call)
+    pathToNode = nodeToEdit
+  } else {
+    const name = findUniqueName(
+      modifiedAst,
+      KCL_DEFAULT_CONSTANT_PREFIXES.TRANSLATE
+    )
+    const declaration = createVariableDeclaration(name, call)
+    modifiedAst.body.push(declaration)
+    pathToNode = createPathToNode(modifiedAst)
+  }
+
+  return {
+    modifiedAst,
+    pathToNode,
+  }
+}
 
 export function setTranslate({
   modifiedAst,
@@ -190,3 +282,36 @@ export function retrievePathToNodeFromTransformSelection(
 
   return pathToNode
 }
+
+export function retrieveGeometryNameFromPath(
+  ast: Node<Program>,
+  pathToNode: PathToNode
+) {
+  const returnEarly = true
+  const geometryNode = getNodeFromPath<
+    VariableDeclaration | ImportStatement | PipeExpression
+  >(
+    ast,
+    pathToNode,
+    ['VariableDeclaration', 'ImportStatement', 'PipeExpression'],
+    returnEarly
+  )
+  if (err(geometryNode)) {
+    return new Error("Couldn't find corresponding path to node")
+  }
+
+  let geometryName: string | undefined
+  if (geometryNode.node.type === 'VariableDeclaration') {
+    geometryName = geometryNode.node.declaration.id.name
+  } else if (
+    geometryNode.node.type === 'ImportStatement' &&
+    geometryNode.node.selector.type === 'None' &&
+    geometryNode.node.selector.alias
+  ) {
+    geometryName = geometryNode.node.selector.alias?.name
+  } else {
+    return new Error("Couldn't find corresponding geometry")
+  }
+
+  return geometryName
+}

src/lib/constants.ts

@@ -53,6 +53,7 @@ export const KCL_DEFAULT_CONSTANT_PREFIXES = {
   PLANE: 'plane',
   HELIX: 'helix',
   CLONE: 'clone',
+  TRANSLATE: 'translate',
 } as const
 /** The default KCL length expression */
 export const KCL_DEFAULT_LENGTH = `5`

src/lib/operations.ts

@@ -11,6 +11,7 @@ import type { Artifact } from '@src/lang/std/artifactGraph'
 import {
   getArtifactOfTypes,
   getCapCodeRef,
+  getCodeRefsByArtifactId,
   getEdgeCutConsumedCodeRef,
   getSweepEdgeCodeRef,
   getWallCodeRef,
@@ -1324,55 +1325,75 @@ async function prepareToEditTranslate({ operation }: EnterEditFlowProps) {
     name: 'Translate',
     groupId: 'modeling',
   }
-  const isModuleImport = operation.type === 'GroupBegin'
-  const isSupportedStdLibCall =
-    operation.type === 'StdLibCall' &&
-    stdLibMap[operation.name]?.supportsTransform
-  if (!isModuleImport && !isSupportedStdLibCall) {
-    return {
-      reason: 'Unsupported operation type. Please edit in the code editor.',
-    }
+  if (operation.type !== 'StdLibCall') {
+    return { reason: 'Wrong operation type' }
   }
 
-  const nodeToEdit = getNodePathFromSourceRange(
+  // 1. Map the unlabeled arguments to solid2d selections
+  console.log('operation', operation)
+  if (operation.unlabeledArg?.value.type !== 'Solid') {
+    return { reason: "Couldn't retrieve geometry selection" }
+  }
+
+  const refs = getCodeRefsByArtifactId(
+    operation.unlabeledArg?.value?.value?.artifactId,
+    kclManager.artifactGraph
+  )
+  if (!refs || refs.length === 0) {
+    return { reason: "Couldn't retrieve geometry code refs" }
+  }
+  const selection: Selections = {
+    graphSelections: refs.map((ref) => ({
+      codeRef: ref,
+    })),
+    otherSelections: [],
+  }
+
+  // 2. Convert the x, y, z argument from a string to a KCL expression
+  const x = await stringToKclExpression(
+    codeManager.code.slice(
+      operation.labeledArgs?.['x']?.sourceRange[0],
+      operation.labeledArgs?.['x']?.sourceRange[1]
+    )
+  )
+  if (err(x) || 'errors' in x) {
+    return { reason: "Couldn't retrieve x argument" }
+  }
+
+  const y = await stringToKclExpression(
+    codeManager.code.slice(
+      operation.labeledArgs?.['y']?.sourceRange[0],
+      operation.labeledArgs?.['y']?.sourceRange[1]
+    )
+  )
+  if (err(y) || 'errors' in y) {
+    return { reason: "Couldn't retrieve y argument" }
+  }
+
+  const z = await stringToKclExpression(
+    codeManager.code.slice(
+      operation.labeledArgs?.['z']?.sourceRange[0],
+      operation.labeledArgs?.['z']?.sourceRange[1]
+    )
+  )
+  if (err(z) || 'errors' in z) {
+    return { reason: "Couldn't retrieve z argument" }
+  }
+
+  // 3. Assemble the default argument values for the command,
+  // with `nodeToEdit` set, which will let the actor know
+  // to edit the node that corresponds to the StdLibCall.
+  const argDefaultValues: ModelingCommandSchema['Translate'] = {
+    selection,
+    x,
+    y,
+    z,
+    nodeToEdit: getNodePathFromSourceRange(
       kclManager.ast,
       sourceRangeFromRust(operation.sourceRange)
-  )
-  let x: KclExpression | undefined = undefined
-  let y: KclExpression | undefined = undefined
-  let z: KclExpression | undefined = undefined
-  const pipeLookupFromOperation = getNodeFromPath<PipeExpression>(
-    kclManager.ast,
-    nodeToEdit,
-    'PipeExpression'
-  )
-  let pipe: PipeExpression | undefined
-  const ast = kclManager.ast
-  if (
-    err(pipeLookupFromOperation) ||
-    pipeLookupFromOperation.node.type !== 'PipeExpression'
-  ) {
-    // Look for the last pipe with the import alias and a call to translate
-    const pipes = findPipesWithImportAlias(ast, nodeToEdit, 'translate')
-    pipe = pipes.at(-1)?.expression
-  } else {
-    pipe = pipeLookupFromOperation.node
+    ),
   }
 
-  if (pipe) {
-    const translate = pipe.body.find(
-      (n) => n.type === 'CallExpressionKw' && n.callee.name.name === 'translate'
-    )
-    if (translate?.type === 'CallExpressionKw') {
-      x = await retrieveArgFromPipedCallExpression(translate, 'x')
-      y = await retrieveArgFromPipedCallExpression(translate, 'y')
-      z = await retrieveArgFromPipedCallExpression(translate, 'z')
-    }
-  }
-
-  // Won't be used since we provide nodeToEdit
-  const selection: Selections = { graphSelections: [], otherSelections: [] }
-  const argDefaultValues = { nodeToEdit, selection, x, y, z }
   return {
     ...baseCommand,
     argDefaultValues,

src/machines/modelingMachine.ts

@@ -88,6 +88,8 @@ import {
   setRotate,
   insertExpressionNode,
   retrievePathToNodeFromTransformSelection,
+  retrieveGeometryNameFromPath,
+  addTranslate,
 } from '@src/lang/modifyAst/setTransform'
 import {
   getNodeFromPath,
@@ -3293,58 +3295,16 @@ export const modelingMachine = setup({
           return Promise.reject(new Error(NO_INPUT_PROVIDED_MESSAGE))
         }
 
-        const ast = kclManager.ast
-        const modifiedAst = structuredClone(ast)
+        const { ast, artifactGraph } = kclManager
         const { x, y, z, nodeToEdit, selection } = input
-        let pathToNode = nodeToEdit
-        if (!(pathToNode && typeof pathToNode[1][0] === 'number')) {
-          const result = retrievePathToNodeFromTransformSelection(
-            selection,
-            kclManager.artifactGraph,
-            ast
-          )
-          if (err(result)) {
-            return Promise.reject(result)
-          }
-
-          pathToNode = result
-        }
-
-        // Look for the last pipe with the import alias and a call to translate, with a fallback to rotate.
-        // Otherwise create one
-        const importNodeAndAlias = findImportNodeAndAlias(ast, pathToNode)
-        if (importNodeAndAlias) {
-          const pipes = findPipesWithImportAlias(ast, pathToNode, 'translate')
-          const lastPipe = pipes.at(-1)
-          if (lastPipe && lastPipe.pathToNode) {
-            pathToNode = lastPipe.pathToNode
-          } else {
-            const otherRelevantPipes = findPipesWithImportAlias(
+        const result = addTranslate({
           ast,
-              pathToNode,
-              'rotate'
-            )
-            const lastRelevantPipe = otherRelevantPipes.at(-1)
-            if (lastRelevantPipe && lastRelevantPipe.pathToNode) {
-              pathToNode = lastRelevantPipe.pathToNode
-            } else {
-              pathToNode = insertExpressionNode(
-                modifiedAst,
-                importNodeAndAlias.alias
-              )
-            }
-          }
-        }
-
-        insertVariableAndOffsetPathToNode(x, modifiedAst, pathToNode)
-        insertVariableAndOffsetPathToNode(y, modifiedAst, pathToNode)
-        insertVariableAndOffsetPathToNode(z, modifiedAst, pathToNode)
-        const result = setTranslate({
-          pathToNode,
-          modifiedAst,
-          x: valueOrVariable(x),
-          y: valueOrVariable(y),
-          z: valueOrVariable(z),
+          artifactGraph,
+          selection,
+          x,
+          y,
+          z,
+          nodeToEdit,
         })
         if (err(result)) {
           return Promise.reject(result)
@@ -3471,34 +3431,9 @@ export const modelingMachine = setup({
           pathToNode = result
         }
 
-        const returnEarly = true
-        const geometryNode = getNodeFromPath<
-          VariableDeclaration | ImportStatement | PipeExpression
-        >(
-          ast,
-          pathToNode,
-          ['VariableDeclaration', 'ImportStatement', 'PipeExpression'],
-          returnEarly
-        )
-        if (err(geometryNode)) {
-          return Promise.reject(
-            new Error("Couldn't find corresponding path to node")
-          )
-        }
-
-        let geometryName: string | undefined
-        if (geometryNode.node.type === 'VariableDeclaration') {
-          geometryName = geometryNode.node.declaration.id.name
-        } else if (
-          geometryNode.node.type === 'ImportStatement' &&
-          geometryNode.node.selector.type === 'None' &&
-          geometryNode.node.selector.alias
-        ) {
-          geometryName = geometryNode.node.selector.alias?.name
-        } else {
-          return Promise.reject(
-            new Error("Couldn't find corresponding geometry")
-          )
+        const geometryName = retrieveGeometryNameFromPath(ast, pathToNode)
+        if (err(geometryName)) {
+          return Promise.reject(geometryName)
         }
 
         const result = addClone({