Fix async parser implementation (#1389)

msujew · web-flow · commit c87b3149998b · 2024-02-28T17:39:24.000+01:00
diff --git a/packages/langium/src/parser/async-parser.ts b/packages/langium/src/parser/async-parser.ts
@@ -92,8 +92,8 @@ export abstract class AbstractThreadedAsyncParser implements AsyncParser {
             }, this.terminationDelay);
         });
         worker.parse(text).then(result => {
-            result.value = this.hydrator.hydrate(result.value);
-            deferred.resolve(result as ParseResult<T>);
+            const hydrated = this.hydrator.hydrate<T>(result);
+            deferred.resolve(hydrated);
         }).catch(err => {
             deferred.reject(err);
         }).finally(() => {
diff --git a/packages/langium/src/serializer/hydrator.ts b/packages/langium/src/serializer/hydrator.ts
@@ -12,25 +12,26 @@ import { isAbstractElement, type AbstractElement, type Grammar } from '../langua
 import type { Linker } from '../references/linker.js';
 import type { Lexer } from '../parser/lexer.js';
 import type { LangiumCoreServices } from '../services.js';
-import type { Reference, AstNode, CstNode, LeafCstNode, GenericAstNode, Mutable } from '../syntax-tree.js';
+import type { ParseResult } from '../parser/langium-parser.js';
+import type { Reference, AstNode, CstNode, LeafCstNode, GenericAstNode, Mutable, RootCstNode } from '../syntax-tree.js';
 import { isRootCstNode, isCompositeCstNode, isLeafCstNode, isAstNode, isReference } from '../syntax-tree.js';
 import { streamAst } from '../utils/ast-utils.js';
 import { BiMap } from '../utils/collections.js';
 import { streamCst } from '../utils/cst-utils.js';
 
 /**
- * The hydrator service is responsible for allowing AST nodes to be sent across worker threads.
+ * The hydrator service is responsible for allowing AST parse results to be sent across worker threads.
  */
 export interface Hydrator {
     /**
-     * Converts an AST node to a plain object. The resulting object can be sent across worker threads.
+     * Converts a parse result to a plain object. The resulting object can be sent across worker threads.
      */
-    dehydrate(node: AstNode): object;
+    dehydrate(result: ParseResult<AstNode>): ParseResult<object>;
     /**
-     * Converts a plain object to an AST node. The resulting AST node can be used in the main thread.
-     * Calling this method on non-plain objects will result in undefined behavior.
+     * Converts a plain object to a parse result. The included AST node can then be used in the main thread.
+     * Calling this method on objects that have not been dehydrated first will result in undefined behavior.
      */
-    hydrate(node: object): AstNode;
+    hydrate<T extends AstNode = AstNode>(result: ParseResult<object>): ParseResult<T>;
 }
 
 export interface DehydrateContext {
@@ -58,8 +59,14 @@ export class DefaultHydrator implements Hydrator {
         this.linker = services.references.Linker;
     }
 
-    dehydrate(node: AstNode): object {
-        return this.dehydrateAstNode(node, this.createDehyrationContext(node));
+    dehydrate(result: ParseResult<AstNode>): ParseResult<object> {
+        return {
+            // We need to create shallow copies of the errors
+            // The original errors inherit from the `Error` class, which is not transferable across worker threads
+            lexerErrors: result.lexerErrors.map(e => ({ ...e })),
+            parserErrors: result.parserErrors.map(e => ({ ...e })),
+            value: this.dehydrateAstNode(result.value, this.createDehyrationContext(result.value))
+        };
     }
 
     protected createDehyrationContext(node: AstNode): DehydrateContext {
@@ -128,6 +135,7 @@ export class DefaultHydrator implements Hydrator {
         if (isRootCstNode(node)) {
             cstNode.fullText = node.fullText;
         } else {
+            // Note: This returns undefined for hidden nodes (i.e. comments)
             cstNode.grammarSource = this.getGrammarElementId(node.grammarSource);
         }
         cstNode.hidden = node.hidden;
@@ -146,12 +154,17 @@ export class DefaultHydrator implements Hydrator {
         return cstNode;
     }
 
-    hydrate(node: object): AstNode {
+    hydrate<T extends AstNode = AstNode>(result: ParseResult<object>): ParseResult<T> {
+        const node = result.value;
         const context = this.createHydrationContext(node);
         if ('$cstNode' in node) {
             this.hydrateCstNode(node.$cstNode, context);
         }
-        return this.hydrateAstNode(node, context);
+        return {
+            lexerErrors: result.lexerErrors,
+            parserErrors: result.parserErrors,
+            value: this.hydrateAstNode(node, context) as T
+        };
     }
 
     protected createHydrationContext(node: any): HydrateContext {
@@ -160,18 +173,21 @@ export class DefaultHydrator implements Hydrator {
         for (const astNode of streamAst(node)) {
             astNodes.set(astNode, {} as AstNode);
         }
+        let root: RootCstNode;
         if (node.$cstNode) {
             for (const cstNode of streamCst(node.$cstNode)) {
-                let cst: CstNode | undefined;
+                let cst: Mutable<CstNode> | undefined;
                 if ('fullText' in cstNode) {
                     cst = new RootCstNodeImpl(cstNode.fullText as string);
+                    root = cst as RootCstNode;
                 } else if ('content' in cstNode) {
                     cst = new CompositeCstNodeImpl();
                 } else if ('tokenType' in cstNode) {
                     cst = this.hydrateCstLeafNode(cstNode);
                 }
                 if (cst) {
                     cstNodes.set(cstNode, cst);
+                    cst.root = root!;
                 }
             }
         }
@@ -198,15 +214,15 @@ export class DefaultHydrator implements Hydrator {
                 astNode[name] = arr;
                 for (const item of value) {
                     if (isAstNode(item)) {
-                        arr.push(this.setParent(this.hydrate(item), astNode));
+                        arr.push(this.setParent(this.hydrateAstNode(item, context), astNode));
                     } else if (isReference(item)) {
                         arr.push(this.hydrateReference(item, astNode, name, context));
                     } else {
                         arr.push(item);
                     }
                 }
             } else if (isAstNode(value)) {
-                astNode[name] = this.setParent(this.hydrate(value), astNode);
+                astNode[name] = this.setParent(this.hydrateAstNode(value, context), astNode);
             } else if (isReference(value)) {
                 astNode[name] = this.hydrateReference(value, astNode, name, context);
             } else if (value !== undefined) {
@@ -272,11 +288,11 @@ export class DefaultHydrator implements Hydrator {
         return this.lexer.definition[name];
     }
 
-    protected getGrammarElementId(node: AbstractElement): number {
+    protected getGrammarElementId(node: AbstractElement): number | undefined {
         if (this.grammarElementIdMap.size === 0) {
             this.createGrammarElementIdMap();
         }
-        return this.grammarElementIdMap.get(node) ?? -1;
+        return this.grammarElementIdMap.get(node);
     }
 
     protected getGrammarElement(id: number): AbstractElement {
diff --git a/packages/langium/test/parser/worker-thread-async-parser.test.ts b/packages/langium/test/parser/worker-thread-async-parser.test.ts
@@ -8,7 +8,8 @@ import { describe, expect, test } from 'vitest';
 import { WorkerThreadAsyncParser } from 'langium/node';
 import { createLangiumGrammarServices } from 'langium/grammar';
 import type { Grammar, LangiumCoreServices, ParseResult } from 'langium';
-import { EmptyFileSystem, GrammarUtils, isOperationCancelled } from 'langium';
+import type { LangiumServices } from 'langium/lsp';
+import { EmptyFileSystem, GrammarUtils, CstUtils, GrammarAST, isOperationCancelled } from 'langium';
 import { CancellationToken, CancellationTokenSource } from 'vscode-languageserver';
 import { fail } from 'node:assert';
 import { fileURLToPath } from 'node:url';
@@ -17,18 +18,18 @@ class TestAsyncParser extends WorkerThreadAsyncParser {
     constructor(services: LangiumCoreServices) {
         super(services, () => fileURLToPath(new URL('.', import.meta.url)) + '/worker-thread.js');
     }
+    setThreadCount(threadCount: number): void {
+        this.threadCount = threadCount;
+    }
 }
 
 describe('WorkerThreadAsyncParser', () => {
 
     test('performs async parsing in parallel', async () => {
-        const services = createLangiumGrammarServices(EmptyFileSystem, undefined, {
-            parser: {
-                AsyncParser: (services) => new TestAsyncParser(services)
-            }
-        }).grammar;
+        const services = getServices();
         const file = createLargeFile(10);
-        const asyncParser = services.parser.AsyncParser;
+        const asyncParser = services.parser.AsyncParser as TestAsyncParser;
+        asyncParser.setThreadCount(4);
         const promises: Array<Promise<ParseResult<Grammar>>> = [];
         for (let i = 0; i < 16; i++) {
             promises.push(asyncParser.parse<Grammar>(file, CancellationToken.None));
@@ -41,11 +42,7 @@ describe('WorkerThreadAsyncParser', () => {
     });
 
     test('async parsing can be cancelled', async () => {
-        const services = createLangiumGrammarServices(EmptyFileSystem, undefined, {
-            parser: {
-                AsyncParser: (services) => new TestAsyncParser(services)
-            }
-        }).grammar;
+        const services = getServices();
         // This file should take a few seconds to parse
         const file = createLargeFile(100000);
         const asyncParser = services.parser.AsyncParser;
@@ -63,11 +60,72 @@ describe('WorkerThreadAsyncParser', () => {
         expect(end - start).toBeLessThan(1000);
     });
 
+    test('async parsing can be cancelled and then restarted', async () => {
+        const services = getServices();
+        // This file should take a few seconds to parse
+        const file = createLargeFile(100000);
+        const asyncParser = services.parser.AsyncParser;
+        const cancellationTokenSource = new CancellationTokenSource();
+        setTimeout(() => cancellationTokenSource.cancel(), 50);
+        try {
+            await asyncParser.parse<Grammar>(file, cancellationTokenSource.token);
+            fail('Parsing should have been cancelled');
+        } catch (err) {
+            expect(isOperationCancelled(err)).toBe(true);
+        }
+        // Calling this method should recreate the worker and parse the file correctly
+        const result = await asyncParser.parse<Grammar>(createLargeFile(10), CancellationToken.None);
+        expect(result.value.name).toBe('Test');
+    });
+
+    test('async parsing yields correct CST', async () => {
+        const services = getServices();
+        const file = createLargeFile(10);
+        const result = await services.parser.AsyncParser.parse<Grammar>(file, CancellationToken.None);
+        const index = file.indexOf('TestRule');
+        // Assert that the CST can be found at all from the root node
+        // This indicates that the CST is correctly linked to itself
+        const node = CstUtils.findLeafNodeAtOffset(result.value.$cstNode!, index)!;
+        expect(node).toBeDefined();
+        expect(node.text).toBe('TestRule0');
+        // Assert that the CST node is correctly linked to its container elements
+        expect(node.container?.container).toBeDefined();
+        expect(node.container!.container!.text).toBe('TestRule0: name="Hello";');
+        // Assert that the CST node has a reference to the root
+        expect(node.root).toBeDefined();
+        expect(node.root.fullText).toBe(file);
+        // Assert that the CST node has a reference to the correct AST node
+        const astNode = node?.astNode as GrammarAST.ParserRule;
+        expect(astNode).toBeDefined();
+        expect(astNode.$type).toBe(GrammarAST.ParserRule);
+        expect(astNode.name).toBe('TestRule0');
+    });
+
+    test('parser errors are correctly transmitted', async () => {
+        const services = getServices();
+        const file = 'grammar Test Rule: name="Hello" // missing semicolon';
+        const result = await services.parser.AsyncParser.parse<Grammar>(file, CancellationToken.None);
+        expect(result.parserErrors).toHaveLength(1);
+        expect(result.parserErrors[0].name).toBe('MismatchedTokenException');
+        expect(result.parserErrors[0]).toHaveProperty('previousToken');
+    });
+
     function createLargeFile(size: number): string {
-        let result = 'grammar Test;\n';
+        let result = 'grammar Test\n';
         for (let i = 0; i < size; i++) {
             result += 'TestRule' + i + ': name="Hello";\n';
         }
         return result;
     }
+
+    function getServices(): LangiumServices {
+        const services = createLangiumGrammarServices(EmptyFileSystem, undefined, {
+            parser: {
+                AsyncParser: (services) => new TestAsyncParser(services)
+            }
+        }).grammar;
+        // We usually only need one thread for testing
+        (services.parser.AsyncParser as TestAsyncParser).setThreadCount(1);
+        return services;
+    }
 });
diff --git a/packages/langium/test/parser/worker-thread.js b/packages/langium/test/parser/worker-thread.js
@@ -14,9 +14,6 @@ const hydrator = services.serializer.Hydrator;
 
 parentPort.on('message', text => {
     const result = parser.parse(text);
-    const dehydrated = hydrator.dehydrate(result.value);
-    parentPort.postMessage({
-        ...result,
-        value: dehydrated
-    });
+    const dehydrated = hydrator.dehydrate(result);
+    parentPort.postMessage(dehydrated);
 });
