refactor codeToIdSelections (#5432)

* add test for original range to artifact conversion

* try naieve refactor

* types clean up

* typo

* break function into smaller functions

* optimizations

* better comments

* camera test tweak

* fmt

* 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)

* overflow fix

* update binary search to ignore end ranges

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

* break binary search into sub function

---------

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

e2e/playwright/testing-camera-movement.spec.ts

@@ -94,6 +94,8 @@ test.describe('Testing Camera Movement', { tag: ['@skipWin'] }, () => {
     await bakeInRetries(async () => {
       await page.mouse.move(700, 200)
       await page.mouse.down({ button: 'right' })
+      await page.waitForTimeout(100)
+
       const appLogoBBox = await page.getByTestId('app-logo').boundingBox()
       expect(appLogoBBox).not.toBeNull()
       if (!appLogoBBox) throw new Error('app logo not found')
@@ -101,7 +103,9 @@ test.describe('Testing Camera Movement', { tag: ['@skipWin'] }, () => {
         appLogoBBox.x + appLogoBBox.width / 2,
         appLogoBBox.y + appLogoBBox.height / 2
       )
+      await page.waitForTimeout(100)
       await page.mouse.move(600, 303)
+      await page.waitForTimeout(100)
       await page.mouse.up({ button: 'right' })
     }, [4, -10.5, -120])
 

src/components/AstExplorer.tsx

@@ -133,9 +133,11 @@ function DisplayObj({
       }}
       onClick={(e) => {
         const range = topLevelRange(obj?.start || 0, obj.end || 0)
-        const idInfo = codeToIdSelections([
-          { codeRef: codeRefFromRange(range, kclManager.ast) },
-        ])[0]
+        const idInfo = codeToIdSelections(
+          [{ codeRef: codeRefFromRange(range, kclManager.ast) }],
+          engineCommandManager.artifactGraph,
+          engineCommandManager.artifactIndex
+        )[0]
         const artifact = engineCommandManager.artifactGraph.get(
           idInfo?.id || ''
         )

src/components/ModelingMachineProvider.tsx

@@ -802,7 +802,7 @@ export const ModelingMachineProvider = ({
               engineCommandManager.artifactGraph
             )
             if (err(plane)) return Promise.reject(plane)
-            // if the user selected a segment, make sure we enter the right sketch as there can be multiple on a plan
+            // if the user selected a segment, make sure we enter the right sketch as there can be multiple on a plane
             // but still works if the user selected a plane/face by defaulting to the first path
             const mainPath =
               artifact?.type === 'segment' || artifact?.type === 'solid2d'

src/editor/manager.ts

@@ -374,6 +374,7 @@ export default class EditorManager {
       selectionRanges: this._selectionRanges,
       isShiftDown: this._isShiftDown,
       ast: kclManager.ast,
+      artifactGraph: engineCommandManager.artifactGraph,
     })
 
     if (!eventInfo) {

src/lang/std/engineConnection.ts

@@ -31,6 +31,8 @@ import { markOnce } from 'lib/performance'
 import { MachineManager } from 'components/MachineManagerProvider'
 import { DefaultPlaneStr } from 'lib/planes'
 import { defaultPlaneStrToKey } from 'lib/planes'
+import { buildArtifactIndex } from 'lib/artifactIndex'
+import { ArtifactIndex } from 'lib/artifactIndex'
 
 // TODO(paultag): This ought to be tweakable.
 const pingIntervalMs = 5_000
@@ -1407,6 +1409,7 @@ export class EngineCommandManager extends EventTarget {
    * see: src/lang/std/artifactGraph-README.md for a full explanation.
    */
   artifactGraph: ArtifactGraph = new Map()
+  artifactIndex: ArtifactIndex = []
   /**
    * The pendingCommands object is a map of the commands that have been sent to the engine that are still waiting on a reply
    */
@@ -2184,6 +2187,7 @@ export class EngineCommandManager extends EventTarget {
   }
   updateArtifactGraph(execStateArtifactGraph: ExecState['artifactGraph']) {
     this.artifactGraph = execStateArtifactGraph
+    this.artifactIndex = buildArtifactIndex(execStateArtifactGraph)
     // TODO check if these still need to be deferred once e2e tests are working again.
     if (this.artifactGraph.size) {
       this.deferredArtifactEmptied(null)

src/lib/artifactIndex.ts

@@ -0,0 +1,29 @@
+import { ArtifactGraph, ArtifactId, SourceRange, Artifact } from 'lang/wasm'
+import { getFaceCodeRef } from 'lang/std/artifactGraph'
+
+// Index artifacts in an ordered list for binary search
+export type ArtifactEntry = { artifact: Artifact; id: ArtifactId }
+/** Index artifacts by their codeRef range, ordered by start position */
+export type ArtifactIndex = Array<{
+  range: SourceRange
+  entry: ArtifactEntry
+}>
+
+/** Creates an array of artifacts, only those with codeRefs, orders them by start range,
+ * to be used later by binary search */
+export function buildArtifactIndex(
+  artifactGraph: ArtifactGraph
+): ArtifactIndex {
+  const index: ArtifactIndex = []
+
+  Array.from(artifactGraph).forEach(([id, artifact]) => {
+    const codeRef = getFaceCodeRef(artifact)
+    if (!codeRef?.range) return
+
+    const entry = { artifact, id }
+    index.push({ range: codeRef.range, entry })
+  })
+
+  // Sort by start position for binary search
+  return index.sort((a, b) => a.range[0] - b.range[0])
+}

src/lib/selections.ts

@@ -11,6 +11,7 @@ import {
   Expr,
   defaultSourceRange,
   topLevelRange,
+  ArtifactGraph,
 } from 'lang/wasm'
 import { ModelingMachineEvent } from 'machines/modelingMachine'
 import { isNonNullable, uuidv4 } from 'lib/utils'
@@ -31,19 +32,13 @@ import { PathToNodeMap } from 'lang/std/sketchcombos'
 import { err } from 'lib/trap'
 import {
   Artifact,
-  getArtifactOfTypes,
-  getArtifactsOfTypes,
-  getCapCodeRef,
-  getSweepEdgeCodeRef,
-  getSolid2dCodeRef,
-  getWallCodeRef,
   CodeRef,
   getCodeRefsByArtifactId,
   ArtifactId,
-  getFaceCodeRef,
 } from 'lang/std/artifactGraph'
 import { Node } from 'wasm-lib/kcl/bindings/Node'
 import { DefaultPlaneStr } from './planes'
+import { ArtifactEntry, ArtifactIndex } from './artifactIndex'
 
 export const X_AXIS_UUID = 'ad792545-7fd3-482a-a602-a93924e3055b'
 export const Y_AXIS_UUID = '680fd157-266f-4b8a-984f-cdf46b8bdf01'
@@ -54,38 +49,7 @@ export type DefaultPlaneSelection = {
   id: string
 }
 
-/** @deprecated Use {@link Artifact} instead. */
-type Selection__old =
-  | {
-      type:
-        | 'default'
-        | 'line-end'
-        | 'line-mid'
-        | 'extrude-wall'
-        | 'solid2d'
-        | 'start-cap'
-        | 'end-cap'
-        | 'point'
-        | 'edge'
-        | 'adjacent-edge'
-        | 'line'
-        | 'arc'
-        | 'all'
-      range: SourceRange
-    }
-  | {
-      type: 'opposite-edgeCut' | 'adjacent-edgeCut' | 'base-edgeCut'
-      range: SourceRange
-      // TODO this is a temporary measure that well be made redundant with: https://github.com/KittyCAD/modeling-app/pull/3836
-      secondaryRange: SourceRange
-    }
 export type NonCodeSelection = Axis | DefaultPlaneSelection
-
-/** @deprecated Use {@link Selection} instead. */
-export type Selections__old = {
-  otherSelections: NonCodeSelection[]
-  codeBasedSelections: Selection__old[]
-}
 export interface Selection {
   artifact?: Artifact
   codeRef: CodeRef
@@ -95,76 +59,6 @@ export type Selections = {
   graphSelections: Array<Selection>
 }
 
-/** @deprecated If you're writing a new function, it should use {@link Selection} and not {@link Selection__old}
- * this function should only be used for backwards compatibility with old functions.
- */
-function convertSelectionToOld(selection: Selection): Selection__old | null {
-  // return {} as Selection__old
-  // TODO implementation
-  const _artifact = selection.artifact
-  if (_artifact?.type === 'solid2d') {
-    const codeRef = getSolid2dCodeRef(
-      _artifact,
-      engineCommandManager.artifactGraph
-    )
-    if (err(codeRef)) return null
-    return { range: codeRef.range, type: 'solid2d' }
-  }
-  if (_artifact?.type === 'cap') {
-    const codeRef = getCapCodeRef(_artifact, engineCommandManager.artifactGraph)
-    if (err(codeRef)) return null
-    return {
-      range: codeRef.range,
-      type: _artifact?.subType === 'end' ? 'end-cap' : 'start-cap',
-    }
-  }
-  if (_artifact?.type === 'wall') {
-    const codeRef = getWallCodeRef(
-      _artifact,
-      engineCommandManager.artifactGraph
-    )
-    if (err(codeRef)) return null
-    return { range: codeRef.range, type: 'extrude-wall' }
-  }
-  if (_artifact?.type === 'segment' || _artifact?.type === 'path') {
-    return { range: _artifact.codeRef.range, type: 'default' }
-  }
-  if (_artifact?.type === 'sweepEdge') {
-    const codeRef = getSweepEdgeCodeRef(
-      _artifact,
-      engineCommandManager.artifactGraph
-    )
-    if (err(codeRef)) return null
-    if (_artifact?.subType === 'adjacent') {
-      return { range: codeRef.range, type: 'adjacent-edge' }
-    }
-    return { range: codeRef.range, type: 'edge' }
-  }
-  if (_artifact?.type === 'edgeCut') {
-    const codeRef = _artifact.codeRef
-    return { range: codeRef.range, type: 'default' }
-  }
-  if (selection?.codeRef?.range) {
-    return { range: selection.codeRef.range, type: 'default' }
-  }
-  return null
-}
-/** @deprecated If you're writing a new function, it should use {@link Selection} and not {@link Selection__old}
- * this function should only be used for backwards compatibility with old functions.
- */
-export function convertSelectionsToOld(selection: Selections): Selections__old {
-  const selections: Selection__old[] = []
-  for (const artifact of selection.graphSelections) {
-    const converted = convertSelectionToOld(artifact)
-    if (converted) selections.push(converted)
-  }
-  const selectionsOld: Selections__old = {
-    otherSelections: selection.otherSelections,
-    codeBasedSelections: selections,
-  }
-  return selectionsOld
-}
-
 export async function getEventForSelectWithPoint({
   data,
 }: Extract<
@@ -310,7 +204,6 @@ export function handleSelectionBatch({
   selections.graphSelections.forEach(({ artifact }) => {
     artifact?.id &&
       selectionToEngine.push({
-        type: 'default',
         id: artifact?.id,
         range:
           getCodeRefsByArtifactId(
@@ -350,7 +243,6 @@ export function handleSelectionBatch({
 }
 
 type SelectionToEngine = {
-  type: Selection__old['type']
   id?: string
   range: SourceRange
 }
@@ -360,11 +252,13 @@ export function processCodeMirrorRanges({
   selectionRanges,
   isShiftDown,
   ast,
+  artifactGraph,
 }: {
   codeMirrorRanges: readonly SelectionRange[]
   selectionRanges: Selections
   isShiftDown: boolean
   ast: Program
+  artifactGraph: ArtifactGraph
 }): null | {
   modelingEvent: ModelingMachineEvent
   engineEvents: Models['WebSocketRequest_type'][]
@@ -392,8 +286,11 @@ export function processCodeMirrorRanges({
         },
       }
     })
-  const idBasedSelections: SelectionToEngine[] =
-    codeToIdSelections(codeBasedSelections)
+  const idBasedSelections: SelectionToEngine[] = codeToIdSelections(
+    codeBasedSelections,
+    artifactGraph,
+    engineCommandManager.artifactIndex
+  )
   const selections: Selection[] = []
   for (const { id, range } of idBasedSelections) {
     if (!id) {
@@ -406,11 +303,8 @@ export function processCodeMirrorRanges({
       })
       continue
     }
-    const artifact = engineCommandManager.artifactGraph.get(id)
-    const codeRefs = getCodeRefsByArtifactId(
-      id,
-      engineCommandManager.artifactGraph
-    )
+    const artifact = artifactGraph.get(id)
+    const codeRefs = getCodeRefsByArtifactId(id, artifactGraph)
     if (artifact && codeRefs) {
       selections.push({ artifact, codeRef: codeRefs[0] })
     } else if (codeRefs) {
@@ -601,234 +495,150 @@ export function canSubmitSelectionArg(
   )
 }
 
-export function codeToIdSelections(
-  selections: Selection[]
-): SelectionToEngine[] {
-  const selectionsOld = convertSelectionsToOld({
-    graphSelections: selections,
-    otherSelections: [],
-  }).codeBasedSelections
-  return selectionsOld
-    .flatMap((selection): null | SelectionToEngine[] => {
-      const { type } = selection
-      // TODO #868: loops over all artifacts will become inefficient at a large scale
-      const overlappingEntries = Array.from(engineCommandManager.artifactGraph)
-        .map(([id, artifact]) => {
-          const codeRef = getFaceCodeRef(artifact)
-          if (!codeRef) return null
-          return isOverlap(codeRef.range, selection.range)
-            ? {
-                artifact,
-                selection,
-                id,
-              }
-            : null
-        })
-        .filter(isNonNullable)
-
-      /** TODO refactor
-       * selections in our app is a sourceRange plus some metadata
-       * The metadata is just a union type string of different types of artifacts or 3d features 'extrude-wall' 'segment' etc
-       * Because the source range is not enough to figure out what the user selected, so here we're using filtering through all the artifacts
-       * to find something that matches both the source range and the metadata.
-       *
-       * What we should migrate to is just storing what the user selected by what it matched in the artifactGraph it will simply the below a lot.
-       *
-       * In the case of a user moving the cursor them, we will still need to figure out what artifact from the graph matches best, but we will just need sane defaults
-       * and most of the time we can expect the user to be clicking in the 3d scene instead.
+/**
+ * Find the index of the last range where range[0] < targetStart
+ * This is used as a starting point for linear search of overlapping ranges
+ * @param index The sorted array of ranges to search through
+ * @param targetStart The start position to compare against
+ * @returns The index of the last range where range[0] < targetStart
  */
-      let bestCandidate:
-        | {
-            id: ArtifactId
-            artifact: unknown
-            selection: Selection__old
+export function findLastRangeStartingBefore(
+  index: ArtifactIndex,
+  targetStart: number
+): number {
+  let left = 0
+  let right = index.length - 1
+  let lastValidIndex = 0
+
+  while (left <= right) {
+    const mid = left + Math.floor((right - left) / 2)
+    const midRange = index[mid].range
+
+    if (midRange[0] < targetStart) {
+      // This range starts before our selection, look in right half for later ones
+      lastValidIndex = mid
+      left = mid + 1
+    } else {
+      // This range starts at or after our selection, look in left half
+      right = mid - 1
     }
-        | undefined
-      overlappingEntries.forEach((entry) => {
-        // TODO probably need to remove much of the `type === 'xyz'` below
-        if (type === 'default' && entry.artifact.type === 'segment') {
-          bestCandidate = entry
-          return
   }
+
+  return lastValidIndex
+}
+
+function findOverlappingArtifactsFromIndex(
+  selection: Selection,
+  index: ArtifactIndex
+): ArtifactEntry[] {
+  if (!selection.codeRef?.range) {
+    console.warn('Selection missing code reference range')
+    return []
+  }
+
+  const selectionRange = selection.codeRef.range
+  const results: ArtifactEntry[] = []
+
+  // Binary search to find the last range where range[0] < selectionRange[0]
+  // This search does not take into consideration the end range, so it's possible
+  // the index it finds dose not have any overlap (depending on the end range)
+  // but it's main purpose is to act as a starting point for the linear part of the search
+  // so a tiny loss in efficiency is acceptable to keep the code simple
+  const startIndex = findLastRangeStartingBefore(index, selectionRange[0])
+
+  // Check all potential overlaps from the found position
+  for (let i = startIndex; i < index.length; i++) {
+    const { range, entry } = index[i]
+    // Stop if we've gone past possible overlaps
+    if (range[0] > selectionRange[1]) break
+
+    if (isOverlap(range, selectionRange)) {
+      results.push(entry)
+    }
+  }
+
+  return results
+}
+
+function getBestCandidate(
+  entries: ArtifactEntry[],
+  artifactGraph: ArtifactGraph
+): ArtifactEntry | undefined {
+  if (!entries.length) {
+    return undefined
+  }
+
+  for (const entry of entries) {
+    // Segments take precedence
+    if (entry.artifact.type === 'segment') {
+      return entry
+    }
+
+    // Handle paths and their solid2d references
     if (entry.artifact.type === 'path') {
-          const artifact = engineCommandManager.artifactGraph.get(
-            entry.artifact.solid2dId || ''
-          )
-          if (artifact?.type !== 'solid2d') {
-            bestCandidate = {
-              artifact: entry.artifact,
-              selection,
-              id: entry.id,
-            }
-          }
-          if (!entry.artifact.solid2dId) {
-            console.error(
-              'Expected PathArtifact to have solid2dId, but none found'
-            )
-            return
-          }
-          bestCandidate = {
-            artifact: artifact,
-            selection,
-            id: entry.artifact.solid2dId,
-          }
-        }
-        if (entry.artifact.type === 'plane') {
-          bestCandidate = {
-            artifact: entry.artifact,
-            selection,
-            id: entry.id,
-          }
-        }
-        if (entry.artifact.type === 'cap') {
-          bestCandidate = {
-            artifact: entry.artifact,
-            selection,
-            id: entry.id,
-          }
-        }
-        if (entry.artifact.type === 'wall') {
-          bestCandidate = {
-            artifact: entry.artifact,
-            selection,
-            id: entry.id,
-          }
-        }
-        if (type === 'extrude-wall' && entry.artifact.type === 'segment') {
-          if (!entry.artifact.surfaceId) return
-          const wall = engineCommandManager.artifactGraph.get(
-            entry.artifact.surfaceId
-          )
-          if (wall?.type !== 'wall') return
-          bestCandidate = {
-            artifact: wall,
-            selection,
-            id: entry.artifact.surfaceId,
-          }
-          return
-        }
-        if (type === 'edge' && entry.artifact.type === 'segment') {
-          const edges = getArtifactsOfTypes(
-            { keys: entry.artifact.edgeIds, types: ['sweepEdge'] },
-            engineCommandManager.artifactGraph
-          )
-          const edge = [...edges].find(([_, edge]) => edge.type === 'sweepEdge')
-          if (!edge) return
-          bestCandidate = {
-            artifact: edge[1],
-            selection,
-            id: edge[0],
-          }
-        }
-        if (type === 'adjacent-edge' && entry.artifact.type === 'segment') {
-          const edges = getArtifactsOfTypes(
-            { keys: entry.artifact.edgeIds, types: ['sweepEdge'] },
-            engineCommandManager.artifactGraph
-          )
-          const edge = [...edges].find(
-            ([_, edge]) =>
-              edge.type === 'sweepEdge' && edge.subType === 'adjacent'
-          )
-          if (!edge) return
-          bestCandidate = {
-            artifact: edge[1],
-            selection,
-            id: edge[0],
-          }
-        }
-        if (
-          (type === 'end-cap' || type === 'start-cap') &&
-          entry.artifact.type === 'path'
-        ) {
-          if (!entry.artifact.sweepId) return
-          const extrusion = getArtifactOfTypes(
-            {
-              key: entry.artifact.sweepId,
-              types: ['sweep'],
-            },
-            engineCommandManager.artifactGraph
-          )
-          if (err(extrusion)) return
-          const caps = getArtifactsOfTypes(
-            { keys: extrusion.surfaceIds, types: ['cap'] },
-            engineCommandManager.artifactGraph
-          )
-          const cap = [...caps].find(
-            ([_, cap]) => cap.subType === (type === 'end-cap' ? 'end' : 'start')
-          )
-          if (!cap) return
-          bestCandidate = {
-            artifact: entry.artifact,
-            selection,
-            id: cap[0],
-          }
-          return
-        }
-        if (entry.artifact.type === 'edgeCut') {
-          const consumedEdge = getArtifactOfTypes(
-            {
-              key: entry.artifact.consumedEdgeId,
-              types: ['segment', 'sweepEdge'],
-            },
-            engineCommandManager.artifactGraph
-          )
-          if (err(consumedEdge)) return
-          if (
-            consumedEdge.type === 'segment' &&
-            type === 'base-edgeCut' &&
-            isOverlap(
-              consumedEdge.codeRef.range,
-              selection.secondaryRange || [0, 0]
-            )
-          ) {
-            bestCandidate = {
-              artifact: entry.artifact,
-              selection,
-              id: entry.id,
-            }
-          } else if (
-            consumedEdge.type === 'sweepEdge' &&
-            ((type === 'adjacent-edgeCut' &&
-              consumedEdge.subType === 'adjacent') ||
-              (type === 'opposite-edgeCut' &&
-                consumedEdge.subType === 'opposite'))
-          ) {
-            const seg = getArtifactOfTypes(
-              { key: consumedEdge.segId, types: ['segment'] },
-              engineCommandManager.artifactGraph
-            )
-            if (err(seg)) return
-            if (
-              isOverlap(seg.codeRef.range, selection.secondaryRange || [0, 0])
-            ) {
-              bestCandidate = {
-                artifact: entry.artifact,
-                selection,
-                id: entry.id,
-              }
+      const solid2dId = entry.artifact.solid2dId
+      if (!solid2dId) {
+        return entry
       }
+      const solid2d = artifactGraph.get(solid2dId)
+      if (solid2d?.type === 'solid2d') {
+        return { id: solid2dId, artifact: solid2d }
       }
+      continue
     }
 
-        if (entry.artifact.type === 'sweep') {
-          bestCandidate = {
-            artifact: entry.artifact,
-            selection,
-            id: entry.id,
+    // Other valid artifact types
+    if (['plane', 'cap', 'wall', 'sweep'].includes(entry.artifact.type)) {
+      return entry
     }
   }
-      })
+  return undefined
+}
 
-      if (bestCandidate) {
-        return [
-          {
-            type,
-            id: bestCandidate.id,
-            range: bestCandidate.selection.range,
-          },
-        ]
+function createSelectionToEngine(
+  selection: Selection,
+  candidateId?: ArtifactId
+): SelectionToEngine {
+  return {
+    ...(candidateId && { id: candidateId }),
+    range: selection.codeRef.range,
   }
-      return [selection]
+}
+
+export function codeToIdSelections(
+  selections: Selection[],
+  artifactGraph: ArtifactGraph,
+  artifactIndex: ArtifactIndex
+): SelectionToEngine[] {
+  if (!selections?.length) {
+    return []
+  }
+
+  if (!artifactGraph) {
+    console.warn('Artifact graph is missing or empty')
+    return selections.map((selection) => createSelectionToEngine(selection))
+  }
+
+  return selections
+    .flatMap((selection): SelectionToEngine[] => {
+      if (!selection) {
+        console.warn('Null or undefined selection encountered')
+        return []
+      }
+
+      // Direct artifact case
+      if (selection.artifact?.id) {
+        return [createSelectionToEngine(selection, selection.artifact.id)]
+      }
+
+      // Find matching artifacts by code range overlap
+      const overlappingEntries = findOverlappingArtifactsFromIndex(
+        selection,
+        artifactIndex
+      )
+      const bestCandidate = getBestCandidate(overlappingEntries, artifactGraph)
+
+      return [createSelectionToEngine(selection, bestCandidate?.id)]
     })
     .filter(isNonNullable)
 }

src/machines/commandBarMachine.ts

@@ -5,7 +5,6 @@ import {
   CommandArgumentWithName,
   KclCommandValue,
 } from 'lib/commandTypes'
-import { Selections__old } from 'lib/selections'
 import { getCommandArgumentKclValuesOnly } from 'lib/commandUtils'
 import { MachineManager } from 'components/MachineManagerProvider'
 import toast from 'react-hot-toast'
@@ -16,7 +15,6 @@ export type CommandBarContext = {
   commands: Command[]
   selectedCommand?: Command
   currentArgument?: CommandArgument<unknown> & { name: string }
-  selectionRanges: Selections__old
   argumentsToSubmit: { [x: string]: unknown }
   machineManager: MachineManager
 }