Improve completion results for data type rules (#1138)

packages/langium/src/lsp/completion/completion-provider.ts

@@ -10,19 +10,20 @@ import type { LangiumCompletionParser } from '../../parser/langium-parser';
 import type { NameProvider } from '../../references/name-provider';
 import type { ScopeProvider } from '../../references/scope-provider';
 import type { LangiumServices } from '../../services';
-import type { AstNode, AstNodeDescription, Reference, ReferenceInfo } from '../../syntax-tree';
+import type { AstNode, AstNodeDescription, CstNode, Reference, ReferenceInfo } from '../../syntax-tree';
 import type { MaybePromise } from '../../utils/promise-util';
 import type { LangiumDocument } from '../../workspace/documents';
 import type { NextFeature } from './follow-element-computation';
 import type { NodeKindProvider } from '../node-kind-provider';
 import type { FuzzyMatcher } from '../fuzzy-matcher';
+import type { GrammarConfig } from '../../grammar/grammar-config';
 import type { Lexer } from '../../parser/lexer';
 import type { IToken } from 'chevrotain';
 import { CompletionItemKind, CompletionList, Position } from 'vscode-languageserver';
 import * as ast from '../../grammar/generated/ast';
 import { getExplicitRuleType } from '../../grammar/internal-grammar-util';
 import { getContainerOfType } from '../../utils/ast-util';
-import { findLeafNodeAtOffset } from '../../utils/cst-util';
+import { findDeclarationNodeAtOffset, findLeafNodeAtOffset } from '../../utils/cst-util';
 import { getEntryRule } from '../../utils/grammar-util';
 import { stream } from '../../utils/stream';
 import { findFirstFeatures, findNextFeatures } from './follow-element-computation';
@@ -126,6 +127,7 @@ export class DefaultCompletionProvider implements CompletionProvider {
     protected readonly lexer: Lexer;
     protected readonly nodeKindProvider: NodeKindProvider;
     protected readonly fuzzyMatcher: FuzzyMatcher;
+    protected readonly grammarConfig: GrammarConfig;
 
     constructor(services: LangiumServices) {
         this.scopeProvider = services.references.ScopeProvider;
@@ -135,6 +137,7 @@ export class DefaultCompletionProvider implements CompletionProvider {
         this.lexer = services.parser.Lexer;
         this.nodeKindProvider = services.shared.lsp.NodeKindProvider;
         this.fuzzyMatcher = services.shared.lsp.FuzzyMatcher;
+        this.grammarConfig = services.parser.GrammarConfig;
     }
 
     async getCompletion(document: LangiumDocument, params: CompletionParams): Promise<CompletionList | undefined> {
@@ -222,13 +225,29 @@ export class DefaultCompletionProvider implements CompletionProvider {
         const textDocument = document.textDocument;
         const text = textDocument.getText();
         const offset = textDocument.offsetAt(position);
-        const { nextTokenStart, nextTokenEnd, previousTokenStart, previousTokenEnd } = this.backtrackToAnyToken(text, offset);
         const partialContext = {
             document,
             textDocument,
             offset,
             position
         };
+        // Data type rules need special handling, as their tokens are irrelevant for completion purposes.
+        // If we encounter a data type rule at the current offset, we jump to the start of the data type rule.
+        const dataTypeRuleOffsets = this.findDataTypeRuleStart(cst, offset);
+        if (dataTypeRuleOffsets) {
+            const [ruleStart, ruleEnd] = dataTypeRuleOffsets;
+            const parentNode = findLeafNodeAtOffset(cst, ruleStart)?.element;
+            const previousTokenFeatures = this.findFeaturesAt(textDocument, ruleStart);
+            yield {
+                ...partialContext,
+                node: parentNode,
+                tokenOffset: ruleStart,
+                tokenEndOffset: ruleEnd,
+                features: previousTokenFeatures,
+            };
+        }
+        // For all other purposes, it's enough to jump to the start of the current/previous token
+        const { nextTokenStart, nextTokenEnd, previousTokenStart, previousTokenEnd } = this.backtrackToAnyToken(text, offset);
         let astNode: AstNode | undefined;
         if (previousTokenStart !== undefined && previousTokenEnd !== undefined && previousTokenEnd === offset) {
             astNode = findLeafNodeAtOffset(cst, previousTokenStart)?.element;
@@ -265,6 +284,23 @@ export class DefaultCompletionProvider implements CompletionProvider {
         }
     }
 
+    protected findDataTypeRuleStart(cst: CstNode, offset: number): [number, number] | undefined {
+        let containerNode: CstNode | undefined = findDeclarationNodeAtOffset(cst, offset, this.grammarConfig.nameRegexp);
+        // Identify whether the element was parsed as part of a data type rule
+        let isDataTypeNode = Boolean(getContainerOfType(containerNode?.feature, ast.isParserRule)?.dataType);
+        if (isDataTypeNode) {
+            while (isDataTypeNode) {
+                // Use the container to find the correct parent element
+                containerNode = containerNode?.parent;
+                isDataTypeNode = Boolean(getContainerOfType(containerNode?.feature, ast.isParserRule)?.dataType);
+            }
+            if (containerNode) {
+                return [containerNode.offset, containerNode.end];
+            }
+        }
+        return undefined;
+    }
+
     /**
      * Indicates whether the completion should continue to process the next completion context.
      *

packages/langium/test/lsp/completion-provider.test.ts

@@ -202,3 +202,94 @@ describe('Path import completion', () => {
         });
     });
 });
+
+describe('Completion in data type rules', () => {
+
+    test('Can perform completion for fully qualified names', async () => {
+        const grammar = `
+        grammar FQNCompletionTest
+
+        entry Model:
+            (persons+=Person | greetings+=Greeting)*;
+
+        Person:
+            'person' name=FQN;
+
+        Greeting:
+            'Hello' person=[Person:FQN] '!';
+
+        FQN returns string: ID ('.' ID)*;
+
+        hidden terminal WS: /\\s+/;
+        terminal ID: /[_a-zA-Z][\\w_]*/;
+        hidden terminal ML_COMMENT: /\\/\\*[\\s\\S]*?\\*\\//;
+        `;
+
+        const services = await createServicesForGrammar({ grammar });
+        const completion = expectCompletion(services);
+
+        const text = `
+            person John.Miller
+            person John.Smith.Junior
+            person John.Smith.Senior
+
+            Hello <|>John<|>.Smi<|>th.Jun<|>ior
+            Hello <|>John./* Hello */ <|>Miller
+        `;
+
+        await completion({
+            text: text,
+            index: 0,
+            expectedItems: [
+                'John.Miller',
+                'John.Smith.Junior',
+                'John.Smith.Senior'
+            ]
+        });
+
+        await completion({
+            text: text,
+            index: 1,
+            expectedItems: [
+                'John.Miller',
+                'John.Smith.Junior',
+                'John.Smith.Senior'
+            ]
+        });
+
+        await completion({
+            text: text,
+            index: 2,
+            expectedItems: [
+                'John.Smith.Junior',
+                'John.Smith.Senior'
+            ]
+        });
+
+        await completion({
+            text: text,
+            index: 3,
+            expectedItems: [
+                'John.Smith.Junior'
+            ]
+        });
+
+        await completion({
+            text: text,
+            index: 4,
+            expectedItems: [
+                'John.Miller',
+                'John.Smith.Junior',
+                'John.Smith.Senior'
+            ]
+        });
+
+        // A comment within the FQN should prevent any completion from appearing
+        await completion({
+            text: text,
+            index: 5,
+            expectedItems: []
+        });
+    });
+
+});