From f9895accd8c1d8edc088c93e341e9fb35846507f Mon Sep 17 00:00:00 2001
From: Dmitry Bushev <bushevdv@gmail.com>
Date: Sat, 27 Jul 2024 12:15:57 +0100
Subject: [PATCH] Reapply "Persist a subset of IdMap (#10347)" (#10626)

This reverts commit 033e4ae32386d5fcad5ca5e3bfce8823e6aa2530.
---
 app/gui2/shared/languageServer.ts             | 10 +++++--
 app/gui2/shared/languageServerTypes.ts        |  9 ++++++
 app/gui2/ydoc-server/languageServerSession.ts | 30 +++++++++++++++----
 app/gui2/ydoc-server/serialization.ts         |  2 +-
 4 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/app/gui2/shared/languageServer.ts b/app/gui2/shared/languageServer.ts
index 935f94f5d59b..9fd73ae6c62d 100644
--- a/app/gui2/shared/languageServer.ts
+++ b/app/gui2/shared/languageServer.ts
@@ -14,6 +14,8 @@ import type {
   ExpressionId,
   FileEdit,
   FileSystemObject,
+  IdMapTriple,
+  IdMapTuple,
   Notifications,
   Path,
   RegisterOptions,
@@ -288,8 +290,12 @@ export class LanguageServer extends ObservableV2<Notifications & TransportEvents
   }
 
   /** [Documentation](https://github.com/enso-org/enso/blob/develop/docs/language-server/protocol-language-server.md#textapplyedit) */
-  applyEdit(edit: FileEdit, execute: boolean): Promise<LsRpcResult<void>> {
-    return this.request('text/applyEdit', { edit, execute })
+  applyEdit(
+    edit: FileEdit,
+    execute: boolean,
+    idMap?: IdMapTriple[] | IdMapTuple[],
+  ): Promise<LsRpcResult<void>> {
+    return this.request('text/applyEdit', { edit, execute, idMap })
   }
 
   /** [Documentation](https://github.com/enso-org/enso/blob/develop/docs/language-server/protocol-language-server.md#filewrite) */
diff --git a/app/gui2/shared/languageServerTypes.ts b/app/gui2/shared/languageServerTypes.ts
index 9aac47a2fb13..c72db55f3b1e 100644
--- a/app/gui2/shared/languageServerTypes.ts
+++ b/app/gui2/shared/languageServerTypes.ts
@@ -74,6 +74,15 @@ export interface Position {
   character: number
 }
 
+interface IdMapSpan {
+  index: { value: number }
+  size: { value: number }
+}
+
+export type IdMapTuple = [IdMapSpan, string]
+
+export type IdMapTriple = [number, number, string]
+
 export type RegisterOptions = { path: Path } | { contextId: ContextId } | {}
 
 export interface CapabilityRegistration {
diff --git a/app/gui2/ydoc-server/languageServerSession.ts b/app/gui2/ydoc-server/languageServerSession.ts
index 6e5a265084f3..d7d5d10997bf 100644
--- a/app/gui2/ydoc-server/languageServerSession.ts
+++ b/app/gui2/ydoc-server/languageServerSession.ts
@@ -34,7 +34,7 @@ import {
   translateVisualizationFromFile,
 } from './edits'
 import * as fileFormat from './fileFormat'
-import { deserializeIdMap, serializeIdMap } from './serialization'
+import { deserializeIdMap, idMapToArray, serializeIdMap } from './serialization'
 import { WSSharedDoc } from './ydoc'
 
 const SOURCE_DIR = 'src'
@@ -457,6 +457,18 @@ class ModulePersistence extends ObservableV2<{ removed: () => void }> {
     }
   }
 
+  private static getIdMapToPersist(
+    idMap: IdMap | undefined,
+    metadata: fileFormat.IdeMetadata['node'],
+  ): IdMap | undefined {
+    if (idMap === undefined) {
+      return
+    } else {
+      const entriesIntersection = idMap.entries().filter(([, id]) => id in metadata)
+      return new IdMap(entriesIntersection)
+    }
+  }
+
   private sendLsUpdate(
     synced: EnsoFileParts,
     newCode: string | undefined,
@@ -468,11 +480,17 @@ class ModulePersistence extends ObservableV2<{ removed: () => void }> {
     const code = newCode ?? synced.code
     const newMetadataJson =
       newMetadata &&
-      json.stringify({ ...this.syncedMeta, ide: { ...this.syncedMeta.ide, node: newMetadata } })
-    const newIdMapJson = newIdMap && serializeIdMap(newIdMap)
+      json.stringify({
+        ...this.syncedMeta,
+        ide: { ...this.syncedMeta.ide, node: newMetadata },
+      })
+    const idMapToPersist =
+      (newIdMap || newMetadata) &&
+      ModulePersistence.getIdMapToPersist(newIdMap, newMetadata ?? this.syncedMeta.ide.node)
+    const newIdMapToPersistJson = idMapToPersist && serializeIdMap(idMapToPersist)
     const newContent = combineFileParts({
       code,
-      idMapJson: newIdMapJson ?? synced.idMapJson ?? '[]',
+      idMapJson: newIdMapToPersistJson ?? synced.idMapJson ?? '[]',
       metadataJson: newMetadataJson ?? synced.metadataJson ?? '{}',
     })
 
@@ -502,7 +520,7 @@ class ModulePersistence extends ObservableV2<{ removed: () => void }> {
 
     const execute = newCode != null || newIdMap != null
     const edit: FileEdit = { path: this.path, edits, oldVersion: this.syncedVersion, newVersion }
-    const apply = this.ls.applyEdit(edit, execute)
+    const apply = this.ls.applyEdit(edit, execute, newIdMap && idMapToArray(newIdMap))
     const handleError = (error: unknown) => {
       console.error('Could not apply edit:', error)
       // Try to recover by reloading the file.
@@ -521,7 +539,7 @@ class ModulePersistence extends ObservableV2<{ removed: () => void }> {
       this.syncedVersion = newVersion
       if (newMetadata) this.syncedMeta.ide.node = newMetadata
       if (newCode) this.syncedCode = newCode
-      if (newIdMapJson) this.syncedIdMap = newIdMapJson
+      if (newIdMapToPersistJson) this.syncedIdMap = newIdMapToPersistJson
       if (newMetadataJson) this.syncedMetaJson = newMetadataJson
       this.setState(LsSyncState.Synchronized)
     }, handleError)
diff --git a/app/gui2/ydoc-server/serialization.ts b/app/gui2/ydoc-server/serialization.ts
index 913c6e3acfc6..f00bb47351ed 100644
--- a/app/gui2/ydoc-server/serialization.ts
+++ b/app/gui2/ydoc-server/serialization.ts
@@ -23,7 +23,7 @@ export function serializeIdMap(map: IdMap): string {
   return json.stringify(idMapToArray(map))
 }
 
-function idMapToArray(map: IdMap): fileFormat.IdMapEntry[] {
+export function idMapToArray(map: IdMap): fileFormat.IdMapEntry[] {
   const entries: fileFormat.IdMapEntry[] = []
   map.entries().forEach(([rangeBuffer, id]) => {
     const decoded = sourceRangeFromKey(rangeBuffer)