Improve completion AST node accuracy

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

@@ -248,9 +248,14 @@ export class DefaultCompletionProvider implements CompletionProvider {
         }
         // 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;
+        let astNodeOffset = nextTokenStart;
+        if (offset <= nextTokenStart && previousTokenStart !== undefined) {
+            // This check indicates that the cursor is still before the next token, so we should use the previous AST node (if it exists)
+            astNodeOffset = previousTokenStart;
+        }
+        const astNode = findLeafNodeAtOffset(cst, astNodeOffset)?.astNode;
+        let performNextCompletion = true;
         if (previousTokenStart !== undefined && previousTokenEnd !== undefined && previousTokenEnd === offset) {
-            astNode = findLeafNodeAtOffset(cst, previousTokenStart)?.astNode;
             // This context aims to complete the current feature
             yield {
                 ...partialContext,
@@ -259,18 +264,26 @@ export class DefaultCompletionProvider implements CompletionProvider {
                 tokenEndOffset: previousTokenEnd,
                 features: this.findFeaturesAt(textDocument, previousTokenStart),
             };
-            // This context aims to complete the immediate next feature (if one exists at the current cursor position)
-            // It uses the previous AST node for that.
-            yield {
-                ...partialContext,
-                node: astNode,
-                tokenOffset: previousTokenEnd,
-                tokenEndOffset: previousTokenEnd,
-                features: this.findFeaturesAt(textDocument, previousTokenEnd),
-            };
+            // The completion after the current token should be prevented in case we find out that the current token definitely isn't completed yet
+            // This is usually the case when the current token ends on a letter.
+            performNextCompletion = this.performNextTokenCompletion(
+                document,
+                text.substring(previousTokenStart, previousTokenEnd),
+                previousTokenStart,
+                previousTokenEnd
+            );
+            if (performNextCompletion) {
+                // This context aims to complete the immediate next feature (if one exists at the current cursor position)
+                // It uses the previous cst start/offset for that.
+                yield {
+                    ...partialContext,
+                    node: astNode,
+                    tokenOffset: previousTokenEnd,
+                    tokenEndOffset: previousTokenEnd,
+                    features: this.findFeaturesAt(textDocument, previousTokenEnd),
+                };
+            }
         }
-        astNode = findLeafNodeAtOffset(cst, nextTokenStart)?.astNode
-            ?? (previousTokenStart === undefined ? undefined : findLeafNodeAtOffset(cst, previousTokenStart)?.astNode);
 
         if (!astNode) {
             const parserRule = getEntryRule(this.grammar);
@@ -284,8 +297,8 @@ export class DefaultCompletionProvider implements CompletionProvider {
                 tokenEndOffset: nextTokenEnd,
                 features: findFirstFeatures(parserRule.definition)
             };
-        } else {
-            // This context aims to complete the next feature, using the next ast node
+        } else if (performNextCompletion) {
+            // This context aims to complete the next feature, using the next cst start/end
             yield {
                 ...partialContext,
                 node: astNode,
@@ -296,6 +309,12 @@ export class DefaultCompletionProvider implements CompletionProvider {
         }
     }
 
+    protected performNextTokenCompletion(document: LangiumDocument, text: string, _offset: number, _end: number): boolean {
+        // This regex returns false if the text ends with a letter
+        // We don't want to complete new text immediately after a keyword, ID etc.
+        return /\P{L}$/u.test(text);
+    }
+
     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
@@ -454,7 +473,7 @@ export class DefaultCompletionProvider implements CompletionProvider {
 
     protected filterKeyword(context: CompletionContext, keyword: ast.Keyword): boolean {
         // Filter out keywords that do not contain any word character
-        return keyword.value.match(/[\w]/) !== null;
+        return /\p{L}/u.test(keyword.value);
     }
 
     protected fillCompletionItem(context: CompletionContext, item: CompletionValueItem): CompletionItem | undefined {

packages/langium/src/utils/cst-util.ts

@@ -155,30 +155,53 @@ export function isCommentNode(cstNode: CstNode, commentNames: string[]): boolean
     return isLeafCstNode(cstNode) && commentNames.includes(cstNode.tokenType.name);
 }
 
+/**
+ * Finds the leaf CST node at the specified 0-based string offset.
+ * If no CST node exists at the specified position, it will return the leaf node before it.
+ *
+ * @param node The CST node to search through
+ * @param offset The specified offset
+ * @returns The CST node closest to the specified output.
+ */
 export function findLeafNodeAtOffset(node: CstNode, offset: number): LeafCstNode | undefined {
     if (isLeafCstNode(node)) {
         return node;
     } else if (isCompositeCstNode(node)) {
-        let firstChild = 0;
-        let lastChild = node.content.length - 1;
-        while (firstChild < lastChild) {
-            const middleChild = Math.floor((firstChild + lastChild) / 2);
-            const n = node.content[middleChild];
-            if (n.offset > offset) {
-                lastChild = middleChild - 1;
-            } else if (n.end <= offset) {
-                firstChild = middleChild + 1;
-            } else {
-                return findLeafNodeAtOffset(n, offset);
-            }
-        }
-        if (firstChild === lastChild) {
-            return findLeafNodeAtOffset(node.content[firstChild], offset);
+        const searchResult = binarySearch(node, offset);
+        if (searchResult) {
+            return findLeafNodeAtOffset(searchResult, offset);
         }
     }
     return undefined;
 }
 
+function binarySearch(node: CompositeCstNode, offset: number): CstNode | undefined {
+    let left = 0;
+    let right = node.content.length - 1;
+    let closest: CstNode | undefined = undefined;
+
+    while (left <= right) {
+        const middle = Math.floor((left + right) / 2);
+        const middleNode = node.content[middle];
+
+        if (middleNode.offset === offset) {
+            // Found an exact match
+            return middleNode;
+        }
+
+        if (middleNode.offset < offset) {
+            // Update the closest node (less than offset) and move to the right half
+            closest = middleNode;
+            left = middle + 1;
+        } else {
+            // Move to the left half
+            right = middle - 1;
+        }
+    }
+
+    return closest;
+}
+
 export function getPreviousNode(node: CstNode, hidden = true): CstNode | undefined {
     while (node.container) {
         const parent = node.container;

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

@@ -75,6 +75,25 @@ describe('Langium completion provider', () => {
             ]
         });
     });
+
+    test('Should not complete next feature directly after ID', async () => {
+        const model = `
+        grammar g<|>
+        A: name<|>="A";
+        B: "B";`;
+        // The completion provider used to propose any referencable element here
+        // It should not propose anything
+        await completion({
+            text: model,
+            index: 0,
+            expectedItems: []
+        });
+        await completion({
+            text: model,
+            index: 1,
+            expectedItems: []
+        });
+    });
 });
 
 describe('Completion within alternatives', () => {

packages/langium/test/utils/cst-utils.test.ts

@@ -55,4 +55,18 @@ describe('CST Utils', () => {
         const leafnode = findLeafNodeAtOffset(grammar.parseResult.value.$cstNode!, offset!);
         expect(leafnode!.text).toBe(';');
     });
+
+    test('Find Leaf Node at Offset: Main: value=AB <|> ;', async () => {
+        const text = `
+        grammar test
+        Main: value=AB  ;
+        terminal fragment Frag: 'B';
+        terminal AB: 'A' Frag;
+        `;
+
+        const grammar = await parser(text);
+        const offset = grammar.textDocument.getText().indexOf('AB') + 3;
+        const leafnode = findLeafNodeAtOffset(grammar.parseResult.value.$cstNode!, offset!);
+        expect(leafnode!.text).toBe('AB');
+    });
 });