introduced dedicated interface PrecomputedScopes, updated implementations of DefaultScopeComputation and DefaultScopeProvider

examples/domainmodel/src/language-server/domain-model-scope.ts

@@ -4,11 +4,11 @@
  * terms of the MIT License, which is available in the project root.
  ******************************************************************************/
 
-import type { AstNode, AstNodeDescription, LangiumDocument, PrecomputedScopes } from 'langium';
+import type { AstNodeDescription, LangiumDocument, PrecomputedScopes } from 'langium';
 import type { DomainModelServices } from './domain-model-module.js';
 import type { QualifiedNameProvider } from './domain-model-naming.js';
 import type { Domainmodel, PackageDeclaration } from './generated/ast.js';
-import { AstUtils, Cancellation, DefaultScopeComputation, interruptAndCheck, MultiMap } from 'langium';
+import { AstUtils, Cancellation, DefaultScopeComputation, interruptAndCheck } from 'langium';
 import { isType, isPackageDeclaration } from './generated/ast.js';
 
 export class DomainModelScopeComputation extends DefaultScopeComputation {
@@ -40,9 +40,9 @@ export class DomainModelScopeComputation extends DefaultScopeComputation {
         return descr;
     }
 
-    override async computeLocalScopes(document: LangiumDocument, cancelToken = Cancellation.CancellationToken.None): Promise<PrecomputedScopes> {
-        const model = document.parseResult.value as Domainmodel;
-        const scopes = new MultiMap<AstNode, AstNodeDescription>();
+    override async computeLocalScopes(document: LangiumDocument<Domainmodel>, cancelToken = Cancellation.CancellationToken.None): Promise<PrecomputedScopes> {
+        const model = document.parseResult.value;
+        const scopes = this.newPrecomputedScopes(document);
         await this.processContainer(model, scopes, document, cancelToken);
         return scopes;
     }

packages/langium/src/references/scope-computation.ts

@@ -13,6 +13,7 @@ import { CancellationToken } from '../utils/cancellation.js';
 import { streamAllContents, streamContents } from '../utils/ast-utils.js';
 import { MultiMap } from '../utils/collections.js';
 import { interruptAndCheck } from '../utils/promise-utils.js';
+import { stream } from '../utils/stream.js';
 
 /**
  * Language-specific service for precomputing global and local scopes. The service methods are executed
@@ -113,7 +114,7 @@ export class DefaultScopeComputation implements ScopeComputation {
 
     async computeLocalScopes(document: LangiumDocument, cancelToken = CancellationToken.None): Promise<PrecomputedScopes> {
         const rootNode = document.parseResult.value;
-        const scopes = new MultiMap<AstNode, AstNodeDescription>();
+        const scopes = this.newPrecomputedScopes(document);
         // Here we navigate the full AST - local scopes shall be available in the whole document
         for (const node of streamAllContents(rootNode)) {
             await interruptAndCheck(cancelToken);
@@ -122,6 +123,19 @@ export class DefaultScopeComputation implements ScopeComputation {
         return scopes;
     }
 
+    /**
+     * @returns A new precomputed scopes container instance for the given document.
+     */
+    protected newPrecomputedScopes(_document: LangiumDocument): PrecomputedScopes {
+        const map = new MultiMap<AstNode, AstNodeDescription>();
+        return {
+            add:       map.add.bind(map),
+            addAll:    map.addAll.bind(map),
+            get:       map.get.bind(map),
+            getStream: (key: AstNode) => stream(map.get(key))
+        };
+    }
+
     /**
      * Process a single node during scopes computation. The default implementation makes the node visible
      * in the subtree of its container (if the node has a name). Override this method to change this,

packages/langium/src/references/scope-provider.ts

@@ -56,11 +56,11 @@ export class DefaultScopeProvider implements ScopeProvider {
         if (precomputed) {
             let currentNode: AstNode | undefined = context.container;
             do {
-                const allDescriptions = precomputed.get(currentNode);
-                if (allDescriptions.length > 0) {
-                    scopes.push(stream(allDescriptions).filter(
-                        desc => this.reflection.isSubtype(desc.type, referenceType)));
-                }
+                scopes.push(
+                    precomputed.getStream(currentNode).filter(
+                        desc => this.reflection.isSubtype(desc.type, referenceType)
+                    )
+                );
                 currentNode = currentNode.$container;
             } while (currentNode);
         }

packages/langium/src/workspace/documents.ts

@@ -19,7 +19,6 @@ import type { ParseResult, ParserOptions } from '../parser/langium-parser.js';
 import type { ServiceRegistry } from '../service-registry.js';
 import type { LangiumSharedCoreServices } from '../services.js';
 import type { AstNode, AstNodeDescription, Mutable, Reference } from '../syntax-tree.js';
-import type { MultiMap } from '../utils/collections.js';
 import type { Stream } from '../utils/stream.js';
 import { TextDocument } from './documents.js';
 import { CancellationToken } from '../utils/cancellation.js';
@@ -95,10 +94,15 @@ export enum DocumentState {
 }
 
 /**
- * Result of the scope precomputation phase (`ScopeComputation` service).
+ * Result of the scope pre-computation phase (`ScopeComputation` service).
  * It maps every AST node to the set of symbols that are visible in the subtree of that node.
  */
-export type PrecomputedScopes = MultiMap<AstNode, AstNodeDescription>
+export interface PrecomputedScopes {
+    add(key: AstNode, value: AstNodeDescription): void
+    addAll(key: AstNode, values: AstNodeDescription[]): void
+    get(key: AstNode): readonly AstNodeDescription[]
+    getStream(key: AstNode): Stream<AstNodeDescription>
+}
 
 export interface DocumentSegment {
     readonly range: Range