Track unknown documents via language id (#1492)

Author: msujew

packages/langium/src/default-module.ts

@@ -107,7 +107,7 @@ export interface DefaultSharedCoreModuleContext {
  */
 export function createDefaultSharedCoreModule(context: DefaultSharedCoreModuleContext): Module<LangiumSharedCoreServices, LangiumDefaultSharedCoreServices> {
     return {
-        ServiceRegistry: () => new DefaultServiceRegistry(),
+        ServiceRegistry: (services) => new DefaultServiceRegistry(services),
         workspace: {
             LangiumDocuments: (services) => new DefaultLangiumDocuments(services),
             LangiumDocumentFactory: (services) => new DefaultLangiumDocumentFactory(services),

packages/langium/src/service-registry.ts

@@ -4,7 +4,8 @@
  * terms of the MIT License, which is available in the project root.
  ******************************************************************************/
 
-import type { LangiumCoreServices } from './services.js';
+import type { LangiumCoreServices, LangiumSharedCoreServices } from './services.js';
+import type { TextDocumentProvider } from './workspace/documents.js';
 import { UriUtils, type URI } from './utils/uri-utils.js';
 
 /**
@@ -41,44 +42,60 @@ export interface ServiceRegistry {
 export class DefaultServiceRegistry implements ServiceRegistry {
 
     protected singleton?: LangiumCoreServices;
-    protected map?: Record<string, LangiumCoreServices>;
+    protected readonly languageIdMap = new Map<string, LangiumCoreServices>();
+    protected readonly fileExtensionMap = new Map<string, LangiumCoreServices>();
+
+    /**
+     * @deprecated Use the new `fileExtensionMap` (or `languageIdMap`) property instead.
+     */
+    protected get map(): Map<string, LangiumCoreServices> | undefined {
+        return this.fileExtensionMap;
+    }
+
+    protected readonly textDocuments?: TextDocumentProvider;
+
+    constructor(services?: LangiumSharedCoreServices) {
+        this.textDocuments = services?.workspace.TextDocuments;
+    }
 
     register(language: LangiumCoreServices): void {
-        if (!this.singleton && !this.map) {
-            // This is the first language to be registered; store it as singleton.
-            this.singleton = language;
-            return;
-        }
-        if (!this.map) {
-            this.map = {};
-            if (this.singleton) {
-                // Move the previous singleton instance to the new map.
-                for (const ext of this.singleton.LanguageMetaData.fileExtensions) {
-                    this.map[ext] = this.singleton;
-                }
-                this.singleton = undefined;
+        const data = language.LanguageMetaData;
+        for (const ext of data.fileExtensions) {
+            if (this.fileExtensionMap.has(ext)) {
+                console.warn(`The file extension ${ext} is used by multiple languages. It is now assigned to '${data.languageId}'.`);
             }
+            this.fileExtensionMap.set(ext, language);
         }
-        // Store the language services in the map.
-        for (const ext of language.LanguageMetaData.fileExtensions) {
-            if (this.map[ext] !== undefined && this.map[ext] !== language) {
-                console.warn(`The file extension ${ext} is used by multiple languages. It is now assigned to '${language.LanguageMetaData.languageId}'.`);
-            }
-            this.map[ext] = language;
+        this.languageIdMap.set(data.languageId, language);
+        if (this.languageIdMap.size === 1) {
+            this.singleton = language;
+        } else {
+            this.singleton = undefined;
         }
     }
 
     getServices(uri: URI): LangiumCoreServices {
         if (this.singleton !== undefined) {
             return this.singleton;
         }
-        if (this.map === undefined) {
+        if (this.languageIdMap.size === 0) {
             throw new Error('The service registry is empty. Use `register` to register the services of a language.');
         }
+        const languageId = this.textDocuments?.get(uri.toString())?.languageId;
+        if (languageId !== undefined) {
+            const services = this.languageIdMap.get(languageId);
+            if (services) {
+                return services;
+            }
+        }
         const ext = UriUtils.extname(uri);
-        const services = this.map[ext];
+        const services = this.fileExtensionMap.get(ext);
         if (!services) {
-            throw new Error(`The service registry contains no services for the extension '${ext}'.`);
+            if (languageId) {
+                throw new Error(`The service registry contains no services for the extension '${ext}' for language '${languageId}'.`);
+            } else {
+                throw new Error(`The service registry contains no services for the extension '${ext}'.`);
+            }
         }
         return services;
     }
@@ -93,12 +110,6 @@ export class DefaultServiceRegistry implements ServiceRegistry {
     }
 
     get all(): readonly LangiumCoreServices[] {
-        if (this.singleton !== undefined) {
-            return [this.singleton];
-        }
-        if (this.map !== undefined) {
-            return Object.values(this.map);
-        }
-        return [];
+        return Array.from(this.languageIdMap.values());
     }
 }

packages/langium/test/service-registry.test.ts

@@ -6,29 +6,57 @@
 
 /* eslint-disable @typescript-eslint/no-explicit-any */
 
-import type { LangiumCoreServices } from 'langium';
+import type { LangiumCoreServices, LangiumSharedCoreServices, Module, PartialLangiumSharedCoreServices, TextDocumentProvider } from 'langium';
 import { describe, expect, test } from 'vitest';
-import { DefaultServiceRegistry, URI } from 'langium';
+import { DefaultServiceRegistry, EmptyFileSystem, URI, createDefaultSharedCoreModule, inject, TextDocument } from 'langium';
 
 describe('DefaultServiceRegistry', () => {
 
     test('should work with a single language', () => {
-        const language: LangiumCoreServices = { foo: 'bar' } as any;
-        const registry = new DefaultServiceRegistry();
+        const language: LangiumCoreServices = { LanguageMetaData: { fileExtensions: ['.foo'], languageId: 'foo' } } as any;
+        const registry = new DefaultServiceRegistry(createSharedCoreServices());
         registry.register(language);
         expect(registry.getServices(URI.parse('file:/foo.bar'))).toBe(language);
         expect(registry.all).toHaveLength(1);
     });
 
     test('should work with two languages', () => {
-        const language1: LangiumCoreServices = { LanguageMetaData: { fileExtensions: ['.foo'] } } as any;
-        const language2: LangiumCoreServices = { LanguageMetaData: { fileExtensions: ['.bar'] } } as any;
-        const registry = new DefaultServiceRegistry();
+        const language1: LangiumCoreServices = { LanguageMetaData: { fileExtensions: ['.foo'], languageId: 'foo' } } as any;
+        const language2: LangiumCoreServices = { LanguageMetaData: { fileExtensions: ['.bar'], languageId: 'bar' } } as any;
+        const registry = new DefaultServiceRegistry(createSharedCoreServices());
         registry.register(language1);
         registry.register(language2);
         expect(registry.getServices(URI.parse('file:/test.foo'))).toBe(language1);
         expect(registry.getServices(URI.parse('file:/test.bar'))).toBe(language2);
         expect(registry.all).toHaveLength(2);
     });
 
+    test('should work based on language ids', () => {
+        const language1: LangiumCoreServices = { LanguageMetaData: { fileExtensions: [], languageId: 'foo' } } as any;
+        const language2: LangiumCoreServices = { LanguageMetaData: { fileExtensions: [], languageId: 'bar' } } as any;
+        const registry = new DefaultServiceRegistry(createSharedCoreServices('bar'));
+        registry.register(language1);
+        registry.register(language2);
+        expect(registry.getServices(URI.parse('file:/test.x'))).toBe(language2);
+        expect(registry.getServices(URI.parse('file:/test.y'))).toBe(language2);
+        expect(registry.all).toHaveLength(2);
+    });
+
+    function createSharedCoreServices(id?: string): LangiumSharedCoreServices {
+        const textDocumentsModule: Module<LangiumSharedCoreServices, PartialLangiumSharedCoreServices> = {
+            workspace: {
+                TextDocuments: (id ? (() => createTextDocuments(id)) : undefined)
+            }
+        };
+        return inject(createDefaultSharedCoreModule(EmptyFileSystem), textDocumentsModule);
+    }
+
+    function createTextDocuments(id: string): TextDocumentProvider {
+        return {
+            get(uri: string): TextDocument | undefined {
+                return TextDocument.create(uri, id, 0, '');
+            }
+        };
+    }
+
 });