From dc18342393e55e57a7a4979f7e950e56a1927c12 Mon Sep 17 00:00:00 2001
From: Mark Sujew <mark.sujew@typefox.io>
Date: Wed, 2 Oct 2024 08:47:57 +0200
Subject: [PATCH] Enable cache eviction on specific document state (#1659)

---
 packages/langium/src/utils/caching.ts       | 60 +++++++++++++++++----
 packages/langium/test/utils/caching.test.ts | 26 ++++++++-
 2 files changed, 74 insertions(+), 12 deletions(-)

diff --git a/packages/langium/src/utils/caching.ts b/packages/langium/src/utils/caching.ts
index 6c26227f9..093a3ba3f 100644
--- a/packages/langium/src/utils/caching.ts
+++ b/packages/langium/src/utils/caching.ts
@@ -7,6 +7,7 @@
 import type { Disposable } from './disposable.js';
 import type { URI } from './uri-utils.js';
 import type { LangiumSharedCoreServices } from '../services.js';
+import type { DocumentState } from '../workspace/documents.js';
 
 export abstract class DisposableCache implements Disposable {
 
@@ -141,14 +142,35 @@ export class ContextCache<Context, Key, Value, ContextKey = Context> extends Dis
  * If this document is changed or deleted, all associated key/value pairs are deleted.
  */
 export class DocumentCache<K, V> extends ContextCache<URI | string, K, V, string> {
-    constructor(sharedServices: LangiumSharedCoreServices) {
+
+    /**
+     * Creates a new document cache.
+     *
+     * @param sharedServices Service container instance to hook into document lifecycle events.
+     * @param state Optional document state on which the cache should evict.
+     * If not provided, the cache will evict on `DocumentBuilder#onUpdate`.
+     * Note that only *changed* documents are considered in this case.
+     *
+     * Providing a state here will use `DocumentBuilder#onDocumentPhase` instead,
+     * which triggers on all documents that have been affected by this change, assuming that the
+     * state is `DocumentState.Linked` or a later state.
+     */
+    constructor(sharedServices: LangiumSharedCoreServices, state?: DocumentState) {
         super(uri => uri.toString());
-        this.onDispose(sharedServices.workspace.DocumentBuilder.onUpdate((changed, deleted) => {
-            const allUris = changed.concat(deleted);
-            for (const uri of allUris) {
-                this.clear(uri);
-            }
-        }));
+        let disposable: Disposable;
+        if (state) {
+            disposable = sharedServices.workspace.DocumentBuilder.onDocumentPhase(state, document => {
+                this.clear(document.uri.toString());
+            });
+        } else {
+            disposable = sharedServices.workspace.DocumentBuilder.onUpdate((changed, deleted) => {
+                const allUris = changed.concat(deleted);
+                for (const uri of allUris) {
+                    this.clear(uri);
+                }
+            });
+        }
+        this.toDispose.push(disposable);
     }
 }
 
@@ -157,10 +179,26 @@ export class DocumentCache<K, V> extends ContextCache<URI | string, K, V, string
  * If any document in the workspace changes, the whole cache is evicted.
  */
 export class WorkspaceCache<K, V> extends SimpleCache<K, V> {
-    constructor(sharedServices: LangiumSharedCoreServices) {
+
+    /**
+     * Creates a new workspace cache.
+     *
+     * @param sharedServices Service container instance to hook into document lifecycle events.
+     * @param state Optional document state on which the cache should evict.
+     * If not provided, the cache will evict on `DocumentBuilder#onUpdate`.
+     */
+    constructor(sharedServices: LangiumSharedCoreServices, state?: DocumentState) {
         super();
-        this.onDispose(sharedServices.workspace.DocumentBuilder.onUpdate(() => {
-            this.clear();
-        }));
+        let disposable: Disposable;
+        if (state) {
+            disposable = sharedServices.workspace.DocumentBuilder.onBuildPhase(state, () => {
+                this.clear();
+            });
+        } else {
+            disposable = sharedServices.workspace.DocumentBuilder.onUpdate(() => {
+                this.clear();
+            });
+        }
+        this.toDispose.push(disposable);
     }
 }
diff --git a/packages/langium/test/utils/caching.test.ts b/packages/langium/test/utils/caching.test.ts
index 97ab192a0..f73f2a2aa 100644
--- a/packages/langium/test/utils/caching.test.ts
+++ b/packages/langium/test/utils/caching.test.ts
@@ -8,7 +8,7 @@
 
 import { beforeEach, describe, expect, test } from 'vitest';
 import type { DefaultDocumentBuilder} from 'langium';
-import { DocumentCache, EmptyFileSystem, URI, WorkspaceCache } from 'langium';
+import { DocumentCache, DocumentState, EmptyFileSystem, URI, WorkspaceCache } from 'langium';
 import { createLangiumGrammarServices } from 'langium/grammar';
 
 const services = createLangiumGrammarServices(EmptyFileSystem);
@@ -83,6 +83,18 @@ describe('Document Cache', () => {
         expect(() => cache.clear()).toThrow();
     });
 
+    test('Document cache can trigger on document state', async () => {
+        const cache = new DocumentCache<string, string>(services.shared, DocumentState.Linked);
+        cache.set(document1.uri, 'key', 'value');
+        const documentPhase = workspace.DocumentBuilder.onDocumentPhase(DocumentState.ComputedScopes, () => {
+            expect(cache.get(document1.uri, 'key')).toBe('value');
+        });
+        await workspace.DocumentBuilder.update([document1.uri], []);
+        expect(cache.has(document1.uri, 'key')).toBe(false);
+        documentPhase.dispose();
+        cache.dispose();
+    });
+
 });
 
 describe('Workspace Cache', () => {
@@ -129,4 +141,16 @@ describe('Workspace Cache', () => {
         expect(() => cache.delete('key')).toThrow();
         expect(() => cache.clear()).toThrow();
     });
+
+    test('Workspace cache can trigger on document state', async () => {
+        const cache = new WorkspaceCache<string, string>(services.shared, DocumentState.Linked);
+        cache.set('key', 'value');
+        const documentPhase = workspace.DocumentBuilder.onBuildPhase(DocumentState.ComputedScopes, () => {
+            expect(cache.get('key')).toBe('value');
+        });
+        await workspace.DocumentBuilder.update([document1.uri], []);
+        expect(cache.has('key')).toBe(false);
+        documentPhase.dispose();
+        cache.dispose();
+    });
 });