#7690 Plane selection improvements (#7707)

* selectOffsetSketchPlane now returns an Error potentially

* move findOperationArtifact to queryAst

* add unit test for findOperationArtifact - not working yet

* fix test

* lint

* lint

* move getCurrentPlaneId to queryAst and rename it to getSelectedPlaneId

* move getSelectedPlane to queryAst

* cleanups, renaming

* add tests for getSelectedPlaneId, getSelectedPlaneAsNode

* getSelectedPlaneAsNode should take variables to avoid kclManager dependency in queryAst

* lint is now happy without this..

* rename findOperationArtifact -> findOperationPlaneArtifact

* PR feedback

Author: andrewvarga

src/components/ModelingSidebar/ModelingPanes/FeatureTreePane.tsx

@@ -40,12 +40,13 @@ import {
   kclEditorActor,
   selectionEventSelector,
 } from '@src/machines/kclEditorMachine'
-import type { Plane } from '@rust/kcl-lib/bindings/Artifact'
 import {
   selectDefaultSketchPlane,
   selectOffsetSketchPlane,
 } from '@src/lib/selections'
 import type { DefaultPlaneStr } from '@src/lib/planes'
+import { findOperationPlaneArtifact, isOffsetPlane } from '@src/lang/queryAst'
+import { err } from '@src/lib/trap'
 
 export const FeatureTreePane = () => {
   const isEditorMounted = useSelector(kclEditorActor, editorIsMountedSelector)
@@ -391,11 +392,17 @@ const OperationItem = (props: {
     )
   }, [kclContext.diagnostics.length])
 
-  function selectOperation() {
+  async function selectOperation() {
     if (props.sketchNoFace) {
       if (isOffsetPlane(props.item)) {
-        const artifact = findOperationArtifact(props.item)
-        void selectOffsetSketchPlane(artifact)
+        const artifact = findOperationPlaneArtifact(
+          props.item,
+          kclManager.artifactGraph
+        )
+        const result = await selectOffsetSketchPlane(artifact)
+        if (err(result)) {
+          console.error(result)
+        }
       }
     } else {
       if (props.item.type === 'GroupEnd') {
@@ -502,7 +509,10 @@ const OperationItem = (props: {
 
   function startSketchOnOffsetPlane() {
     if (isOffsetPlane(props.item)) {
-      const artifact = findOperationArtifact(props.item)
+      const artifact = findOperationPlaneArtifact(
+        props.item,
+        kclManager.artifactGraph
+      )
       if (artifact?.id) {
         sceneInfra.modelingSend({
           type: 'Enter sketch',
@@ -665,7 +675,9 @@ const OperationItem = (props: {
       variableName={variableName}
       valueDetail={valueDetail}
       menuItems={menuItems}
-      onClick={selectOperation}
+      onClick={() => {
+        void selectOperation()
+      }}
       onDoubleClick={props.sketchNoFace ? undefined : enterEditFlow} // no double click in "Sketch no face" mode
       errors={errors}
       greyedOut={!enabled}
@@ -774,17 +786,3 @@ const DefaultPlanes = () => {
     </div>
   )
 }
-
-type StdLibCallOp = Extract<Operation, { type: 'StdLibCall' }>
-
-const isOffsetPlane = (item: Operation): item is StdLibCallOp => {
-  return item.type === 'StdLibCall' && item.name === 'offsetPlane'
-}
-
-const findOperationArtifact = (item: StdLibCallOp) => {
-  const nodePath = JSON.stringify(item.nodePath)
-  const artifact = [...kclManager.artifactGraph.values()].find(
-    (a) => JSON.stringify((a as Plane).codeRef?.nodePath) === nodePath
-  )
-  return artifact
-}

src/components/ViewControlMenu.tsx

@@ -14,15 +14,15 @@ import { kclManager, sceneInfra } from '@src/lib/singletons'
 import { err, reportRejection } from '@src/lib/trap'
 import { useSettings } from '@src/lib/singletons'
 import { resetCameraPosition } from '@src/lib/resetCameraPosition'
-import type { Selections } from '@src/lib/selections'
 import {
   selectDefaultSketchPlane,
   selectOffsetSketchPlane,
 } from '@src/lib/selections'
+import { getSelectedPlaneId } from '@src/lang/queryAst'
 
 export function useViewControlMenuItems() {
   const { state: modelingState, send: modelingSend } = useModelingContext()
-  const selectedPlaneId = getCurrentPlaneId(
+  const selectedPlaneId = getSelectedPlaneId(
     modelingState.context.selectionRanges
   )
 
@@ -112,21 +112,3 @@ export function ViewControlContextMenu({
     />
   )
 }
-
-function getCurrentPlaneId(selectionRanges: Selections): string | null {
-  const defaultPlane = selectionRanges.otherSelections.find(
-    (selection) => typeof selection === 'object' && 'name' in selection
-  )
-  if (defaultPlane) {
-    return defaultPlane.id
-  }
-
-  const planeSelection = selectionRanges.graphSelections.find(
-    (selection) => selection.artifact?.type === 'plane'
-  )
-  if (planeSelection) {
-    return planeSelection.artifact?.id || null
-  }
-
-  return null
-}

src/hooks/useEngineConnectionSubscriptions.ts

@@ -109,7 +109,9 @@ export function useEngineConnectionSubscriptions() {
               }
 
               const artifact = kclManager.artifactGraph.get(planeOrFaceId)
-              if (await selectOffsetSketchPlane(artifact)) {
+              const offsetPlaneSelected =
+                await selectOffsetSketchPlane(artifact)
+              if (!err(offsetPlaneSelected) && offsetPlaneSelected) {
                 return
               }
 

src/lang/queryAst.test.ts

@@ -13,12 +13,16 @@ import {
   doesSceneHaveExtrudedSketch,
   doesSceneHaveSweepableSketch,
   findAllPreviousVariables,
+  findOperationPlaneArtifact,
   findUsesOfTagInPipe,
   getNodeFromPath,
+  getSelectedPlaneId,
+  getSelectedPlaneAsNode,
   getVariableExprsFromSelection,
   hasSketchPipeBeenExtruded,
   isCursorInFunctionDefinition,
   isNodeSafeToReplace,
+  isOffsetPlane,
   retrieveSelectionsFromOpArg,
   traverse,
 } from '@src/lang/queryAst'
@@ -32,6 +36,7 @@ import { initPromise } from '@src/lang/wasmUtils'
 import type { Selections, Selection } from '@src/lib/selections'
 import { enginelessExecutor } from '@src/lib/testHelpers'
 import { err } from '@src/lib/trap'
+import type { Plane } from '@rust/kcl-lib/bindings/Artifact'
 
 beforeAll(async () => {
   await initPromise
@@ -781,6 +786,223 @@ describe('Testing specific sketch getNodeFromPath workflow', () => {
   })
 })
 
+describe('Testing findOperationArtifact', () => {
+  it('should find the correct artifact for a given operation', async () => {
+    const code = `sketch001 = startSketchOn(XY)
+  |> startProfile(at = [0, 0])
+  |> line(end = [10, 0])
+  |> line(end = [10, 10])
+  |> line(end = [0, 10])
+  |> close()
+plane001 = offsetPlane(YZ, offset = 10)
+part001 = startSketchOn(plane001)
+  |> startProfile(at = [0, 0])
+  |> line(end = [5, 0])
+  |> line(end = [5, 5])
+  |> line(end = [0, 5])
+  |> close()
+`
+
+    const ast = assertParse(code)
+    const execState = await enginelessExecutor(ast, false)
+    const { operations, artifactGraph } = execState
+
+    expect(operations).toBeTruthy()
+    expect(operations.length).toBeGreaterThan(0)
+
+    // Find an offsetPlane operation
+    const offsetPlaneOp = operations.find(
+      (op) => op.type === 'StdLibCall' && op.name === 'offsetPlane'
+    )
+    expect(offsetPlaneOp).toBeTruthy()
+
+    if (offsetPlaneOp && isOffsetPlane(offsetPlaneOp)) {
+      const artifact = findOperationPlaneArtifact(offsetPlaneOp, artifactGraph)
+
+      expect(artifact).toBeTruthy()
+      expect(artifact?.type).toBe('plane')
+
+      const artifactNodePath = JSON.stringify(
+        (artifact as Plane)?.codeRef?.nodePath
+      )
+      const operationNodePath = JSON.stringify(offsetPlaneOp.nodePath)
+      expect(artifactNodePath).toBe(operationNodePath)
+    }
+  })
+})
+
+describe('Testing getSelectedPlaneId', () => {
+  it('should return the id of the selected default plane', () => {
+    const selections: Selections = {
+      otherSelections: [
+        {
+          name: 'XY', // actually, this is lowercase during runtime ("xy")!
+          id: 'default-plane-xy-id',
+        },
+      ],
+      graphSelections: [],
+    }
+
+    const result = getSelectedPlaneId(selections)
+    expect(result).toBe('default-plane-xy-id')
+  })
+
+  it('should return the id of the selected offset plane', () => {
+    const selections: Selections = {
+      otherSelections: [],
+      graphSelections: [
+        {
+          artifact: {
+            type: 'plane' as const,
+            id: 'offset-plane-id',
+            pathIds: [],
+            codeRef: {
+              nodePath: {
+                steps: [],
+              },
+              range: [0, 10, 0] as [number, number, number],
+              pathToNode: [
+                ['body', ''],
+                [0, 'index'],
+              ] as PathToNode,
+            },
+          },
+          codeRef: {
+            range: [0, 10, 0] as [number, number, number],
+            pathToNode: [
+              ['body', ''],
+              [0, 'index'],
+            ] as PathToNode,
+          },
+        },
+      ],
+    }
+
+    const result = getSelectedPlaneId(selections)
+    expect(result).toBe('offset-plane-id')
+  })
+
+  it('should prioritize default plane over offset plane when both are selected', () => {
+    const mockPlaneArtifact = {
+      type: 'plane' as const,
+      id: 'offset-plane-id',
+      pathIds: [],
+      codeRef: {
+        range: [0, 10, 0] as [number, number, number],
+        nodePath: { steps: [] },
+        pathToNode: [
+          ['body', ''],
+          [0, 'index'],
+        ] as PathToNode,
+      },
+    }
+
+    const selections: Selections = {
+      otherSelections: [
+        {
+          name: 'XY',
+          id: 'default-plane-xy-id',
+        },
+      ],
+      graphSelections: [
+        {
+          artifact: mockPlaneArtifact,
+          codeRef: {
+            range: [0, 10, 0] as [number, number, number],
+            pathToNode: [
+              ['body', ''],
+              [0, 'index'],
+            ] as PathToNode,
+          },
+        },
+      ],
+    }
+
+    const result = getSelectedPlaneId(selections)
+    expect(result).toBe('default-plane-xy-id')
+  })
+
+  it('should return null when no plane is selected', () => {
+    const selections: Selections = {
+      otherSelections: ['x-axis'],
+      graphSelections: [
+        {
+          artifact: {
+            type: 'startSketchOnFace' as const,
+            id: 'segment-id',
+            faceId: 'face-id',
+            codeRef: {
+              range: [0, 10, 0] as [number, number, number],
+              nodePath: { steps: [] },
+              pathToNode: [
+                ['body', ''],
+                [0, 'index'],
+              ] as PathToNode,
+            },
+          },
+          codeRef: {
+            range: [0, 10, 0] as [number, number, number],
+            pathToNode: [
+              ['body', ''],
+              [0, 'index'],
+            ] as PathToNode,
+          },
+        },
+      ],
+    }
+
+    const result = getSelectedPlaneId(selections)
+    expect(result).toBeNull()
+  })
+
+  it('should return null when selection is empty', () => {
+    const selections: Selections = {
+      otherSelections: [],
+      graphSelections: [],
+    }
+
+    const result = getSelectedPlaneId(selections)
+    expect(result).toBeNull()
+  })
+})
+
+describe('Testing getSelectedPlaneAsNode', () => {
+  it('should return a Node<Literal> for default plane selection', async () => {
+    const code = `sketch001 = startSketchOn(XY)
+  |> startProfile(at = [0, 0])
+  |> line(end = [10, 0])
+  |> line(end = [10, 10])
+  |> line(end = [0, 10])
+  |> close()
+plane001 = offsetPlane(YZ, offset = 10)
+`
+
+    const ast = assertParse(code)
+    const execState = await enginelessExecutor(ast, false)
+    const { variables } = execState
+
+    const selections: Selections = {
+      otherSelections: [
+        {
+          name: 'XY',
+          id: 'default-plane-xy-id',
+        },
+      ],
+      graphSelections: [],
+    }
+
+    const result = getSelectedPlaneAsNode(selections, variables)
+    expect(result).toBeTruthy()
+    expect(result?.type).toBe('Literal')
+    if (result?.type !== 'Literal') {
+      // To make TypeScript happy
+      throw new Error('Expected result to be a Literal node')
+    }
+
+    expect(result?.value).toBe('XY')
+  })
+})
+
 describe('Testing getVariableExprsFromSelection', () => {
   it('should find the variable expr in a simple profile selection', async () => {
     const circleProfileInVar = `sketch001 = startSketchOn(XY)