Introduce SourceDocument for text edit support (#8957)

Introduce `SourceDocument`, a reactive source code representation that can be synced from a `MutableModule`. `SourceDocument` replaces various logic for tracking source code and spans--most importantly, `ReactiveModule`. There is no longer any reactively-tracked `Module`, per-se: Changes to the `MutableModule` attached to the synchronized `ydoc` are pushed as Y.Js events to the `SourceDocument` and `GraphDb`, which are reactively tracked. This avoids a problem in the upcoming text-synchronization (next PR) that was caused by a reactive back channel bypassing the `GraphDb` and resulting in observation of inconsistent states.

Stacked on #8956. Part of #8238.

app/gui2/shared/ast/mutableModule.ts

@@ -2,8 +2,8 @@ import * as random from 'lib0/random'
 import * as Y from 'yjs'
 import type { AstId, Owned, SyncTokenId } from '.'
 import { Token, asOwned, isTokenId, newExternalId } from '.'
-import { assert, assertDefined } from '../util/assert'
-import type { ExternalId, SourceRange } from '../yjsModel'
+import { assert } from '../util/assert'
+import type { ExternalId } from '../yjsModel'
 import type { AstFields, FixedMap, Mutable } from './tree'
 import {
   Ast,
@@ -30,7 +30,6 @@ export interface Module {
   getToken(token: SyncTokenId | undefined): Token | undefined
   getAny(node: AstId | SyncTokenId): Ast | Token
   has(id: AstId): boolean
-  getSpan(id: AstId): SourceRange | undefined
 }
 
 export interface ModuleUpdate {
@@ -130,7 +129,16 @@ export class MutableModule implements Module {
   }
 
   observe(observer: (update: ModuleUpdate) => void) {
-    this.nodes.observeDeep((events) => observer(this.observeEvents(events)))
+    const handle = (events: Y.YEvent<any>[]) => observer(this.observeEvents(events))
+    // Attach the observer first, so that if an update hook causes changes in reaction to the initial state update, we
+    // won't miss them.
+    this.nodes.observeDeep(handle)
+    observer(this.getStateAsUpdate())
+    return handle
+  }
+
+  unobserve(handle: ReturnType<typeof this.observe>) {
+    this.nodes.unobserveDeep(handle)
   }
 
   getStateAsUpdate(): ModuleUpdate {
@@ -237,10 +245,6 @@ export class MutableModule implements Module {
 
   /////////////////////////////////////////////
 
-  getSpan(id: AstId) {
-    return undefined
-  }
-
   constructor(doc: Y.Doc) {
     this.nodes = doc.getMap<YNode>('nodes')
   }

app/gui2/shared/ast/sourceDocument.ts

@@ -0,0 +1,48 @@
+import { print, type AstId, type Module, type ModuleUpdate } from '.'
+import { rangeEquals, sourceRangeFromKey, type SourceRange } from '../yjsModel'
+
+/** Provides a view of the text representation of a module,
+ *  and information about the correspondence between the text and the ASTs,
+ *  that can be kept up-to-date by applying AST changes.
+ */
+export class SourceDocument {
+  private text_: string
+  private readonly spans: Map<AstId, SourceRange>
+
+  private constructor(text: string, spans: Map<AstId, SourceRange>) {
+    this.text_ = text
+    this.spans = spans
+  }
+
+  static Empty() {
+    return new this('', new Map())
+  }
+
+  clear() {
+    if (this.text_ !== '') this.text_ = ''
+    if (this.spans.size !== 0) this.spans.clear()
+  }
+
+  applyUpdate(module: Module, update: ModuleUpdate) {
+    for (const id of update.nodesDeleted) this.spans.delete(id)
+    const root = module.root()
+    if (!root) return
+    const printed = print(root)
+    for (const [key, nodes] of printed.info.nodes) {
+      const range = sourceRangeFromKey(key)
+      for (const node of nodes) {
+        const oldSpan = this.spans.get(node.id)
+        if (!oldSpan || !rangeEquals(range, oldSpan)) this.spans.set(node.id, range)
+      }
+    }
+    if (printed.code !== this.text_) this.text_ = printed.code
+  }
+
+  get text(): string {
+    return this.text_
+  }
+
+  getSpan(id: AstId): SourceRange | undefined {
+    return this.spans.get(id)
+  }
+}

app/gui2/shared/ast/tree.ts

@@ -27,7 +27,7 @@ import {
 import { assert, assertDefined, assertEqual, bail } from '../util/assert'
 import type { Result } from '../util/data/result'
 import { Err, Ok } from '../util/data/result'
-import type { ExternalId, SourceRange, VisualizationMetadata } from '../yjsModel'
+import type { ExternalId, VisualizationMetadata } from '../yjsModel'
 import * as RawAst from './generated/ast'
 
 declare const brandAstId: unique symbol
@@ -84,11 +84,6 @@ export abstract class Ast {
     return this.id === other.id
   }
 
-  /** Return this node's span, if it belongs to a module with an associated span map. */
-  get span(): SourceRange | undefined {
-    return this.module.getSpan(this.id)
-  }
-
   innerExpression(): Ast {
     // TODO: Override this in `Documented`, `Annotated`, `AnnotatedBuiltin`
     return this

app/gui2/src/components/CodeEditor.vue

@@ -39,8 +39,8 @@ const rootElement = ref<HTMLElement>()
 useAutoBlur(rootElement)
 
 const executionContextDiagnostics = computed(() =>
-  projectStore.module && graphStore.moduleCode
-    ? lsDiagnosticsToCMDiagnostics(graphStore.moduleCode, projectStore.diagnostics)
+  projectStore.module && graphStore.moduleSource.text
+    ? lsDiagnosticsToCMDiagnostics(graphStore.moduleSource.text, projectStore.diagnostics)
     : [],
 )
 
@@ -57,8 +57,9 @@ const expressionUpdatesDiagnostics = computed(() => {
     if (!externalId) continue
     const node = nodeMap.get(asNodeId(externalId))
     if (!node) continue
-    if (!node.rootSpan.span) continue
-    const [from, to] = node.rootSpan.span
+    const rootSpan = graphStore.moduleSource.getSpan(node.rootSpan.id)
+    if (!rootSpan) continue
+    const [from, to] = rootSpan
     switch (update.payload.type) {
       case 'Panic': {
         diagnostics.push({ from, to, message: update.payload.message, severity: 'error' })
@@ -88,10 +89,10 @@ watchEffect(() => {
   const awareness = projectStore.awareness.internal
   extensions: [yCollab(yText, awareness, { undoManager }), ...]
    */
-  if (!graphStore.moduleCode) return
+  if (!graphStore.moduleSource.text) return
   editorView.setState(
     EditorState.create({
-      doc: graphStore.moduleCode,
+      doc: graphStore.moduleSource.text,
       extensions: [
         minimalSetup,
         syntaxHighlighting(defaultHighlightStyle as Highlighter),
@@ -105,7 +106,8 @@ watchEffect(() => {
           const astSpan = ast.span()
           let foundNode: NodeId | undefined
           for (const [id, node] of graphStore.db.nodeIdToNode.entries()) {
-            if (node.rootSpan.span && rangeEncloses(node.rootSpan.span, astSpan)) {
+            const rootSpan = graphStore.moduleSource.getSpan(node.rootSpan.id)
+            if (rootSpan && rangeEncloses(rootSpan, astSpan)) {
               foundNode = id
               break
             }

app/gui2/src/components/GraphEditor/NodeWidget.vue

@@ -10,6 +10,7 @@ import {
   provideWidgetUsageInfo,
   usageKeyForInput,
 } from '@/providers/widgetUsageInfo'
+import { useGraphStore } from '@/stores/graph'
 import { Ast } from '@/util/ast'
 import { computed, proxyRefs } from 'vue'
 
@@ -30,6 +31,7 @@ defineOptions({
 
 type UpdateHandler = (update: WidgetUpdate) => boolean
 
+const graph = useGraphStore()
 const registry = injectWidgetRegistry()
 const tree = injectWidgetTree()
 const parentUsageInfo = injectWidgetUsageInfo(true)
@@ -85,8 +87,9 @@ provideWidgetUsageInfo(
 
 const spanStart = computed(() => {
   if (!(props.input instanceof Ast.Ast)) return undefined
-  if (props.input.span == null) return undefined
-  return props.input.span[0] - tree.nodeSpanStart
+  const span = graph.moduleSource.getSpan(props.input.id)
+  if (span == null) return undefined
+  return span[0] - tree.nodeSpanStart
 })
 </script>
 

app/gui2/src/providers/widgetTree.ts

@@ -1,4 +1,5 @@
 import { createContextStore } from '@/providers'
+import { useGraphStore } from '@/stores/graph'
 import { type NodeId } from '@/stores/graph/graphDatabase'
 import { Ast } from '@/util/ast'
 import { computed, proxyRefs, type Ref } from 'vue'
@@ -7,7 +8,8 @@ export { injectFn as injectWidgetTree, provideFn as provideWidgetTree }
 const { provideFn, injectFn } = createContextStore(
   'Widget tree',
   (astRoot: Ref<Ast.Ast>, nodeId: Ref<NodeId>, hasActiveAnimations: Ref<boolean>) => {
-    const nodeSpanStart = computed(() => astRoot.value.span![0])
+    const graph = useGraphStore()
+    const nodeSpanStart = computed(() => graph.moduleSource.getSpan(astRoot.value.id)![0])
     return proxyRefs({ astRoot, nodeId, nodeSpanStart, hasActiveAnimations })
   },
 )