Fix Second-Body Extrude Selection (#5456)

* getSweepArtifactFromSelection

* update getPathToExtrudeForSegmentSelection

* update shell

* add tests and update selection

* A snapshot a day keeps the bugs away! 📷🐛 (OS: namespace-profile-ubuntu-8-cores)

* A snapshot a day keeps the bugs away! 📷🐛 (OS: namespace-profile-ubuntu-8-cores)

* A snapshot a day keeps the bugs away! 📷🐛 (OS: namespace-profile-ubuntu-8-cores)

* A snapshot a day keeps the bugs away! 📷🐛 (OS: namespace-profile-ubuntu-8-cores)

* A snapshot a day keeps the bugs away! 📷🐛 (OS: namespace-profile-ubuntu-8-cores)

* add support for wall and cap

* A snapshot a day keeps the bugs away! 📷🐛 (OS: namespace-profile-ubuntu-8-cores)

* fmt

* add CallExpressionKw

* A snapshot a day keeps the bugs away! 📷🐛 (OS: namespace-profile-ubuntu-8-cores)

* A snapshot a day keeps the bugs away! 📷🐛 (OS: namespace-profile-ubuntu-8-cores)

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

src/lang/modifyAst/addEdgeTreatment.test.ts

@@ -116,7 +116,11 @@ const runGetPathToExtrudeForSegmentSelectionTest = async (
     }
     if (!extrudeInSketchPipe) {
       const init = expectedExtrudeNode.init
-      if (init.type !== 'CallExpression' && init.type !== 'PipeExpression') {
+      if (
+        init.type !== 'CallExpression' &&
+        init.type !== 'CallExpressionKw' &&
+        init.type !== 'PipeExpression'
+      ) {
         return new Error(
           'Expected extrude expression is not a CallExpression or PipeExpression'
         )
@@ -129,25 +133,33 @@ const runGetPathToExtrudeForSegmentSelectionTest = async (
   // ast
   const ast = assertParse(code)
 
-  // selection
+  // range
   const segmentRange = topLevelRange(
     code.indexOf(selectedSegmentSnippet),
     code.indexOf(selectedSegmentSnippet) + selectedSegmentSnippet.length
   )
-  const selection: Selection = {
-    codeRef: codeRefFromRange(segmentRange, ast),
-  }
 
   // executeAst and artifactGraph
   await kclManager.executeAst({ ast })
   const artifactGraph = engineCommandManager.artifactGraph
 
+  // find artifact
+  const maybeArtifact = [...artifactGraph].find(([, artifact]) => {
+    if (!('codeRef' in artifact && artifact.codeRef)) return false
+    return isOverlap(artifact.codeRef.range, segmentRange)
+  })
+
+  // build selection
+  const selection: Selection = {
+    codeRef: codeRefFromRange(segmentRange, ast),
+    artifact: maybeArtifact ? maybeArtifact[1] : undefined,
+  }
+
   // get extrude expression
   const pathResult = getPathToExtrudeForSegmentSelection(
     ast,
     selection,
-    artifactGraph,
-    dependencies
+    artifactGraph
   )
   if (err(pathResult)) return pathResult
   const { pathToExtrudeNode } = pathResult
@@ -234,6 +246,56 @@ extrude003 = extrude(sketch003, length = -15)`
       expectedExtrudeSnippet
     )
   })
+  it('should return the correct paths for a (piped) extrude based on the other body (face)', async () => {
+    const code = `sketch001 = startSketchOn('XY')
+  |> startProfileAt([-25, -25], %)
+  |> yLine(50, %)
+  |> xLine(50, %)
+  |> yLine(-50, %)
+  |> line(endAbsolute = [profileStartX(%), profileStartY(%)])
+  |> close()
+  |> extrude(length = 50)
+sketch002 = startSketchOn(sketch001, 'END')
+  |> startProfileAt([-15, -15], %)
+  |> yLine(30, %)
+  |> xLine(30, %)
+  |> yLine(-30, %)
+  |> line(endAbsolute = [profileStartX(%), profileStartY(%)])
+  |> close()
+  |> extrude(length = 30)`
+    const selectedSegmentSnippet = `xLine(30, %)`
+    const expectedExtrudeSnippet = `extrude(length = 30)`
+    await runGetPathToExtrudeForSegmentSelectionTest(
+      code,
+      selectedSegmentSnippet,
+      expectedExtrudeSnippet
+    )
+  })
+  it('should return the correct paths for a (non-piped) extrude based on the other body (face)', async () => {
+    const code = `sketch001 = startSketchOn('XY')
+  |> startProfileAt([-25, -25], %)
+  |> yLine(50, %)
+  |> xLine(50, %)
+  |> yLine(-50, %)
+  |> line(endAbsolute = [profileStartX(%), profileStartY(%)])
+  |> close()
+extrude001 = extrude(sketch001, length = 50)
+sketch002 = startSketchOn(extrude001, 'END')
+  |> startProfileAt([-15, -15], %)
+  |> yLine(30, %)
+  |> xLine(30, %)
+  |> yLine(-30, %)
+  |> line(endAbsolute = [profileStartX(%), profileStartY(%)])
+  |> close()
+extrude002 = extrude(sketch002, length = 30)`
+    const selectedSegmentSnippet = `xLine(30, %)`
+    const expectedExtrudeSnippet = `extrude002 = extrude(sketch002, length = 30)`
+    await runGetPathToExtrudeForSegmentSelectionTest(
+      code,
+      selectedSegmentSnippet,
+      expectedExtrudeSnippet
+    )
+  })
   it('should not return any path for missing extrusion', async () => {
     const code = `sketch001 = startSketchOn('XY')
   |> startProfileAt([-30, 30], %)

src/lang/modifyAst/addEdgeTreatment.ts

@@ -10,7 +10,6 @@ import {
   Program,
   VariableDeclaration,
   VariableDeclarator,
-  sketchFromKclValue,
 } from '../wasm'
 import {
   createCallExpressionStdLib,
@@ -35,11 +34,11 @@ import {
   sketchLineHelperMap,
   sketchLineHelperMapKw,
 } from '../std/sketch'
-import { err, trap } from 'lib/trap'
+import { err } from 'lib/trap'
 import { Selection, Selections } from 'lib/selections'
 import { KclCommandValue } from 'lib/commandTypes'
 import { isArray } from 'lib/utils'
-import { Artifact, getSweepFromSuspectedPath } from 'lang/std/artifactGraph'
+import { Artifact, getSweepArtifactFromSelection } from 'lang/std/artifactGraph'
 import { Node } from 'wasm-lib/kcl/bindings/Node'
 import { findKwArg } from 'lang/util'
 import { KclManager } from 'lang/KclSingleton'
@@ -121,8 +120,7 @@ export function modifyAstWithEdgeTreatmentAndTag(
     const result = getPathToExtrudeForSegmentSelection(
       clonedAstForGetExtrude,
       selection,
-      artifactGraph,
-      dependencies
+      artifactGraph
     )
     if (err(result)) return result
     const { pathToSegmentNode, pathToExtrudeNode } = result
@@ -279,39 +277,19 @@ function insertParametersIntoAst(
 export function getPathToExtrudeForSegmentSelection(
   ast: Program,
   selection: Selection,
-  artifactGraph: ArtifactGraph,
-  dependencies: {
-    kclManager: KclManager
-    engineCommandManager: EngineCommandManager
-    editorManager: EditorManager
-    codeManager: CodeManager
-  }
+  artifactGraph: ArtifactGraph
 ): { pathToSegmentNode: PathToNode; pathToExtrudeNode: PathToNode } | Error {
   const pathToSegmentNode = getNodePathFromSourceRange(
     ast,
     selection.codeRef?.range
   )
 
-  const varDecNode = getNodeFromPath<VariableDeclaration>(
-    ast,
-    pathToSegmentNode,
-    'VariableDeclaration'
-  )
-  if (err(varDecNode)) return varDecNode
-  const sketchVar = varDecNode.node.declaration.id.name
-
-  const sketch = sketchFromKclValue(
-    dependencies.kclManager.variables[sketchVar],
-    sketchVar
-  )
-  if (trap(sketch)) return sketch
-
-  const extrusion = getSweepFromSuspectedPath(sketch.id, artifactGraph)
-  if (err(extrusion)) return extrusion
+  const sweepArtifact = getSweepArtifactFromSelection(selection, artifactGraph)
+  if (err(sweepArtifact)) return sweepArtifact
 
   const pathToExtrudeNode = getNodePathFromSourceRange(
     ast,
-    extrusion.codeRef.range
+    sweepArtifact.codeRef.range
   )
   if (err(pathToExtrudeNode)) return pathToExtrudeNode
 

src/lang/modifyAst/addShell.ts

@@ -13,36 +13,23 @@ import {
   createLiteral,
   createIdentifier,
   findUniqueName,
-  createCallExpressionStdLib,
-  createObjectExpression,
   createArrayExpression,
   createVariableDeclaration,
   createCallExpressionStdLibKw,
   createLabeledArg,
 } from 'lang/modifyAst'
 import { KCL_DEFAULT_CONSTANT_PREFIXES } from 'lib/constants'
-import { KclManager } from 'lang/KclSingleton'
-import { EngineCommandManager } from 'lang/std/engineConnection'
-import EditorManager from 'editor/manager'
-import CodeManager from 'lang/codeManager'
 
 export function addShell({
   node,
   selection,
   artifactGraph,
   thickness,
-  dependencies,
 }: {
   node: Node<Program>
   selection: Selections
   artifactGraph: ArtifactGraph
   thickness: Expr
-  dependencies: {
-    kclManager: KclManager
-    engineCommandManager: EngineCommandManager
-    editorManager: EditorManager
-    codeManager: CodeManager
-  }
 }): Error | { modifiedAst: Node<Program>; pathToNode: PathToNode } {
   const modifiedAst = structuredClone(node)
 
@@ -55,8 +42,7 @@ export function addShell({
     const extrudeLookupResult = getPathToExtrudeForSegmentSelection(
       clonedAstForGetExtrude,
       graphSelection,
-      artifactGraph,
-      dependencies
+      artifactGraph
     )
     if (err(extrudeLookupResult)) {
       return new Error("Couldn't find extrude")

src/lang/std/artifactGraph.ts

@@ -18,6 +18,7 @@ import {
 } from 'lang/wasm'
 import { Models } from '@kittycad/lib'
 import { getNodePathFromSourceRange } from 'lang/queryAstNodePathUtils'
+import { Selection } from 'lib/selections'
 import { err } from 'lib/trap'
 import { Cap, Plane, Wall } from 'wasm-lib/kcl/bindings/Artifact'
 import { CapSubType } from 'wasm-lib/kcl/bindings/Artifact'
@@ -455,6 +456,47 @@ export function getSweepFromSuspectedPath(
   )
 }
 
+export function getSweepArtifactFromSelection(
+  selection: Selection,
+  artifactGraph: ArtifactGraph
+): SweepArtifact | Error {
+  let sweepArtifact: Artifact | null = null
+  if (selection.artifact?.type === 'sweepEdge') {
+    const _artifact = getArtifactOfTypes(
+      { key: selection.artifact.sweepId, types: ['sweep'] },
+      artifactGraph
+    )
+    if (err(_artifact)) return _artifact
+    sweepArtifact = _artifact
+  } else if (selection.artifact?.type === 'segment') {
+    const _pathArtifact = getArtifactOfTypes(
+      { key: selection.artifact.pathId, types: ['path'] },
+      artifactGraph
+    )
+    if (err(_pathArtifact)) return _pathArtifact
+    if (!_pathArtifact.sweepId) return new Error('Path does not have a sweepId')
+    const _artifact = getArtifactOfTypes(
+      { key: _pathArtifact.sweepId, types: ['sweep'] },
+      artifactGraph
+    )
+    if (err(_artifact)) return _artifact
+    sweepArtifact = _artifact
+  } else if (
+    selection.artifact?.type === 'cap' ||
+    selection.artifact?.type === 'wall'
+  ) {
+    const _artifact = getArtifactOfTypes(
+      { key: selection.artifact.sweepId, types: ['sweep'] },
+      artifactGraph
+    )
+    if (err(_artifact)) return _artifact
+    sweepArtifact = _artifact
+  }
+  if (!sweepArtifact) return new Error('No sweep artifact found')
+
+  return sweepArtifact
+}
+
 export function getCodeRefsByArtifactId(
   id: string,
   artifactGraph: ArtifactGraph

src/machines/modelingMachine.ts

@@ -1995,12 +1995,6 @@ export const modelingMachine = setup({
         // Extract inputs
         const ast = kclManager.ast
         const { selection, thickness } = input
-        const dependencies = {
-          kclManager,
-          engineCommandManager,
-          editorManager,
-          codeManager,
-        }
 
         // Insert the thickness variable if it exists
         if (
@@ -2026,7 +2020,6 @@ export const modelingMachine = setup({
             'variableName' in thickness
               ? thickness.variableIdentifierAst
               : thickness.valueAst,
-          dependencies,
         })
         if (err(shellResult)) {
           return err(shellResult)