Fix id source range mapping with path_get_info (#517)

src/lang/std/engineConnection.ts

@@ -1,10 +1,12 @@
-import { SourceRange } from 'lang/executor'
+import { ProgramMemory, SourceRange } from 'lang/executor'
 import { Selections } from 'useStore'
 import { VITE_KC_API_WS_MODELING_URL, VITE_KC_CONNECTION_TIMEOUT_MS } from 'env'
 import { Models } from '@kittycad/lib'
 import { exportSave } from 'lib/exportSave'
 import { v4 as uuidv4 } from 'uuid'
 import * as Sentry from '@sentry/react'
+import { getNodeFromPath, getNodePathFromSourceRange } from 'lang/queryAst'
+import { Program, VariableDeclarator } from 'lang/abstractSyntaxTreeTypes'
 
 let lastMessage = ''
 
@@ -883,7 +885,10 @@ export class EngineCommandManager {
     }
     return command.promise
   }
-  async waitForAllCommands(): Promise<{
+  async waitForAllCommands(
+    ast?: Program,
+    programMemory?: ProgramMemory
+  ): Promise<{
     artifactMap: ArtifactMap
     sourceRangeMap: SourceRangeMap
   }> {
@@ -892,9 +897,94 @@ export class EngineCommandManager {
     ) as PendingCommand[]
     const proms = pendingCommands.map(({ promise }) => promise)
     await Promise.all(proms)
+    if (ast && programMemory) {
+      await this.fixIdMappings(ast, programMemory)
+    }
+
     return {
       artifactMap: this.artifactMap,
       sourceRangeMap: this.sourceRangeMap,
     }
   }
+  private async fixIdMappings(ast: Program, programMemory: ProgramMemory) {
+    /* This is a temporary solution since the cmd_ids that are sent through when
+    sending 'extend_path' ids are not used as the segment ids. 
+
+    We have a way to back fill them with 'path_get_info', however this relies on one
+    the sketchGroup array and the segements array returned from the server to be in
+    the same length and order. plus it's super hacky, we first use the path_id to get
+    the source range of the pipe expression then use the name of the variable to get
+    the sketchGroup from programMemory.
+    
+    I feel queezy about relying on all these steps to always line up.
+    We have also had to pollute this EngineCommandManager class with knowledge of both the ast and programMemory
+    We should get the cmd_ids to match with the segment ids and delete this method.
+    */
+    const pathInfoProms = []
+    for (const [id, artifact] of Object.entries(this.artifactMap)) {
+      if (artifact.commandType === 'start_path') {
+        pathInfoProms.push(
+          this.sendSceneCommand({
+            type: 'modeling_cmd_req',
+            cmd_id: uuidv4(),
+            cmd: {
+              type: 'path_get_info',
+              path_id: id,
+            },
+          }).then(({ data }) => ({
+            originalId: id,
+            segments: data?.data?.segments,
+          }))
+        )
+      }
+    }
+
+    const pathInfos = await Promise.all(pathInfoProms)
+    pathInfos.forEach(({ originalId, segments }) => {
+      const originalArtifact = this.artifactMap[originalId]
+      if (!originalArtifact || originalArtifact.type === 'pending') {
+        console.log('problem')
+        return
+      }
+      const pipeExpPath = getNodePathFromSourceRange(
+        ast,
+        originalArtifact.range
+      )
+      const pipeExp = getNodeFromPath<VariableDeclarator>(
+        ast,
+        pipeExpPath,
+        'VariableDeclarator'
+      ).node
+      if (pipeExp.type !== 'VariableDeclarator') {
+        console.log('problem', pipeExp, pipeExpPath, ast)
+        return
+      }
+      const variableName = pipeExp.id.name
+      const memoryItem = programMemory.root[variableName]
+      if (!memoryItem) {
+        console.log('problem', variableName, programMemory)
+        return
+      } else if (memoryItem.type !== 'SketchGroup') {
+        console.log('problem', memoryItem, programMemory)
+        return
+      }
+      const relevantSegments = segments.filter(
+        ({ command_id }: { command_id: string | null }) => command_id
+      )
+      if (memoryItem.value.length !== relevantSegments.length) {
+        console.log('problem', memoryItem.value, relevantSegments)
+        return
+      }
+      for (let i = 0; i < relevantSegments.length; i++) {
+        const engineSegment = relevantSegments[i]
+        const memorySegment = memoryItem.value[i]
+        const oldId = memorySegment.__geoMeta.id
+        const artifact = this.artifactMap[oldId]
+        delete this.artifactMap[oldId]
+        delete this.sourceRangeMap[oldId]
+        this.artifactMap[engineSegment.command_id] = artifact
+        this.sourceRangeMap[engineSegment.command_id] = artifact.range
+      }
+    })
+  }
 }