Improve next element computation in completion (#1215)

msujew · web-flow · commit c2cd5c88a32c · 2023-11-01T17:22:07.000+01:00
diff --git a/packages/langium/src/lsp/completion/completion-provider.ts b/packages/langium/src/lsp/completion/completion-provider.ts
@@ -10,7 +10,7 @@ import type { LangiumCompletionParser } from '../../parser/langium-parser.js';
 import type { NameProvider } from '../../references/name-provider.js';
 import type { ScopeProvider } from '../../references/scope-provider.js';
 import type { LangiumServices } from '../../services.js';
-import type { AstNode, AstNodeDescription, CstNode, Reference, ReferenceInfo } from '../../syntax-tree.js';
+import type { AstNode, AstNodeDescription, AstReflection, CstNode, ReferenceInfo } from '../../syntax-tree.js';
 import type { MaybePromise } from '../../utils/promise-util.js';
 import type { LangiumDocument } from '../../workspace/documents.js';
 import type { NextFeature } from './follow-element-computation.js';
@@ -22,7 +22,7 @@ import type { IToken } from 'chevrotain';
 import { CompletionItemKind, CompletionList, Position } from 'vscode-languageserver';
 import * as ast from '../../grammar/generated/ast.js';
 import { getExplicitRuleType } from '../../grammar/internal-grammar-util.js';
-import { getContainerOfType } from '../../utils/ast-util.js';
+import { assignMandatoryAstProperties, getContainerOfType } from '../../utils/ast-util.js';
 import { findDeclarationNodeAtOffset, findLeafNodeAtOffset } from '../../utils/cst-util.js';
 import { getEntryRule } from '../../utils/grammar-util.js';
 import { stream } from '../../utils/stream.js';
@@ -128,6 +128,7 @@ export class DefaultCompletionProvider implements CompletionProvider {
     protected readonly nodeKindProvider: NodeKindProvider;
     protected readonly fuzzyMatcher: FuzzyMatcher;
     protected readonly grammarConfig: GrammarConfig;
+    protected readonly astReflection: AstReflection;
 
     constructor(services: LangiumServices) {
         this.scopeProvider = services.references.ScopeProvider;
@@ -138,6 +139,7 @@ export class DefaultCompletionProvider implements CompletionProvider {
         this.nodeKindProvider = services.shared.lsp.NodeKindProvider;
         this.fuzzyMatcher = services.shared.lsp.FuzzyMatcher;
         this.grammarConfig = services.parser.GrammarConfig;
+        this.astReflection = services.shared.AstReflection;
     }
 
     async getCompletion(document: LangiumDocument, params: CompletionParams): Promise<CompletionList | undefined> {
@@ -200,7 +202,6 @@ export class DefaultCompletionProvider implements CompletionProvider {
             const parserRule = getEntryRule(this.grammar)!;
             const firstFeatures = findFirstFeatures({
                 feature: parserRule.definition,
-                new: true,
                 type: getExplicitRuleType(parserRule)
             });
             if (tokens.length > 0) {
@@ -373,37 +374,36 @@ export class DefaultCompletionProvider implements CompletionProvider {
         };
     }
 
-    protected async completionForRule(context: CompletionContext, rule: ast.AbstractRule, acceptor: CompletionAcceptor): Promise<void> {
-        if (ast.isParserRule(rule)) {
-            const firstFeatures = findFirstFeatures(rule.definition);
-            await Promise.all(firstFeatures.map(next => this.completionFor(context, next, acceptor)));
-        }
-    }
-
     protected completionFor(context: CompletionContext, next: NextFeature, acceptor: CompletionAcceptor): MaybePromise<void> {
         if (ast.isKeyword(next.feature)) {
             return this.completionForKeyword(context, next.feature, acceptor);
         } else if (ast.isCrossReference(next.feature) && context.node) {
             return this.completionForCrossReference(context, next as NextFeature<ast.CrossReference>, acceptor);
         }
+        // Don't offer any completion for other elements (i.e. terminals, datatype rules)
+        // We - from a framework level - cannot reasonably assume their contents.
+        // Adopters can just override `completionFor` if they want to do that anyway.
     }
 
-    protected completionForCrossReference(context: CompletionContext, crossRef: NextFeature<ast.CrossReference>, acceptor: CompletionAcceptor): MaybePromise<void> {
-        const assignment = getContainerOfType(crossRef.feature, ast.isAssignment);
+    protected completionForCrossReference(context: CompletionContext, next: NextFeature<ast.CrossReference>, acceptor: CompletionAcceptor): MaybePromise<void> {
+        const assignment = getContainerOfType(next.feature, ast.isAssignment);
         let node = context.node;
         if (assignment && node) {
-            if (crossRef.type && (crossRef.new || node.$type !== crossRef.type)) {
+            if (next.type) {
+                // When `type` is set, it indicates that we have just entered a new parser rule.
+                // The cross reference that we're trying to complete is on a new element that doesn't exist yet.
+                // So we create a new synthetic element with the correct type information.
                 node = {
-                    $type: crossRef.type,
+                    $type: next.type,
                     $container: node,
-                    $containerProperty: crossRef.property
+                    $containerProperty: next.property
                 };
-            }
-            if (!context) {
-                return;
+                assignMandatoryAstProperties(this.astReflection, node);
             }
             const refInfo: ReferenceInfo = {
-                reference: {} as Reference,
+                reference: {
+                    $refText: ''
+                },
                 container: node,
                 property: assignment.feature
             };
@@ -452,7 +452,7 @@ export class DefaultCompletionProvider implements CompletionProvider {
         });
     }
 
-    protected filterKeyword(_context: CompletionContext, keyword: ast.Keyword): boolean {
+    protected filterKeyword(context: CompletionContext, keyword: ast.Keyword): boolean {
         // Filter out keywords that do not contain any word character
         return keyword.value.match(/[\w]/) !== null;
     }
diff --git a/packages/langium/src/lsp/completion/follow-element-computation.ts b/packages/langium/src/lsp/completion/follow-element-computation.ts
@@ -25,10 +25,6 @@ export interface NextFeature<T extends ast.AbstractElement = ast.AbstractElement
      * The container property for the new `type`
      */
     property?: string
-    /**
-     * Determines whether this `feature` is directly preceded by a new object declaration (such as an action or a rule call)
-     */
-    new?: boolean
 }
 
 /**
@@ -77,8 +73,7 @@ function findNextFeaturesInternal(options: { next: NextFeature, cardinalities: M
         const repeatingFeatures = findFirstFeaturesInternal({
             next: {
                 feature: item,
-                type: next.type,
-                new: false
+                type: next.type
             },
             cardinalities,
             visited,
@@ -95,8 +90,7 @@ function findNextFeaturesInternal(options: { next: NextFeature, cardinalities: M
         if (ownIndex !== undefined && ownIndex < parent.elements.length - 1) {
             features.push(...findNextFeaturesInGroup({
                 feature: parent,
-                type: next.type,
-                new: false
+                type: next.type
             }, ownIndex + 1, cardinalities, visited, plus));
         }
         // Try to find the next elements of the parent
@@ -105,8 +99,7 @@ function findNextFeaturesInternal(options: { next: NextFeature, cardinalities: M
             features.push(...findNextFeaturesInternal({
                 next: {
                     feature: parent,
-                    type: next.type,
-                    new: false
+                    type: next.type
                 },
                 cardinalities,
                 visited,
@@ -146,7 +139,11 @@ function findFirstFeaturesInternal(options: { next: NextFeature, cardinalities:
             .map(e => modifyCardinality(e, feature.cardinality, cardinalities));
     } else if (ast.isAlternatives(feature) || ast.isUnorderedGroup(feature)) {
         return feature.elements.flatMap(e => findFirstFeaturesInternal({
-            next: { feature: e, new: false, type },
+            next: {
+                feature: e,
+                type,
+                property: next.property
+            },
             cardinalities,
             visited,
             plus
@@ -155,7 +152,6 @@ function findFirstFeaturesInternal(options: { next: NextFeature, cardinalities:
     } else if (ast.isAssignment(feature)) {
         const assignmentNext = {
             feature: feature.terminal,
-            new: false,
             type,
             property: next.property ?? feature.feature
         };
@@ -165,7 +161,6 @@ function findFirstFeaturesInternal(options: { next: NextFeature, cardinalities:
         return findNextFeaturesInternal({
             next: {
                 feature,
-                new: true,
                 type: getTypeName(feature),
                 property: next.property ?? feature.feature
             },
@@ -177,8 +172,7 @@ function findFirstFeaturesInternal(options: { next: NextFeature, cardinalities:
         const rule = feature.rule.ref;
         const ruleCallNext = {
             feature: rule.definition,
-            new: true,
-            type: rule.fragment ? undefined : getExplicitRuleType(rule) ?? rule.name,
+            type: rule.fragment || rule.dataType ? undefined : (getExplicitRuleType(rule) ?? rule.name),
             property: next.property
         };
         return findFirstFeaturesInternal({ next: ruleCallNext, cardinalities, visited, plus })
@@ -204,7 +198,11 @@ function findNextFeaturesInGroup(next: NextFeature<ast.Group>, index: number, ca
     const features: NextFeature[] = [];
     let firstFeature: NextFeature;
     while (index < next.feature.elements.length) {
-        firstFeature = { feature: next.feature.elements[index++], new: false, type: next.type };
+        const feature = next.feature.elements[index++];
+        firstFeature = {
+            feature,
+            type: next.type
+        };
         features.push(...findFirstFeaturesInternal({
             next: firstFeature,
             cardinalities,
diff --git a/packages/langium/src/parser/langium-parser.ts b/packages/langium/src/parser/langium-parser.ts
@@ -17,7 +17,7 @@ import { defaultParserErrorProvider, EmbeddedActionsParser, LLkLookaheadStrategy
 import { LLStarLookaheadStrategy } from 'chevrotain-allstar';
 import { isAssignment, isCrossReference, isKeyword } from '../grammar/generated/ast.js';
 import { getTypeName, isDataTypeRule } from '../grammar/internal-grammar-util.js';
-import { getContainerOfType, linkContentToContainer } from '../utils/ast-util.js';
+import { assignMandatoryAstProperties, getContainerOfType, linkContentToContainer } from '../utils/ast-util.js';
 import { CstNodeBuilder } from './cst-node-builder.js';
 
 export type ParseResult<T = AstNode> = {
@@ -275,23 +275,11 @@ export class LangiumParser extends AbstractLangiumParser {
         if (isDataTypeNode(obj)) {
             return this.converter.convert(obj.value, obj.$cstNode);
         } else {
-            this.assignMandatoryProperties(obj);
+            assignMandatoryAstProperties(this.astReflection, obj);
         }
         return obj;
     }
 
-    private assignMandatoryProperties(obj: any): void {
-        const typeMetaData = this.astReflection.getTypeMetaData(obj.$type);
-        for (const mandatoryProperty of typeMetaData.mandatory) {
-            const value = obj[mandatoryProperty.name];
-            if (mandatoryProperty.type === 'array' && !Array.isArray(value)) {
-                obj[mandatoryProperty.name] = [];
-            } else if (mandatoryProperty.type === 'boolean' && value === undefined) {
-                obj[mandatoryProperty.name] = false;
-            }
-        }
-    }
-
     private getAssignment(feature: AbstractElement): AssignmentElement {
         if (!this.assignmentMap.has(feature)) {
             const assignment = getContainerOfType(feature, isAssignment);
diff --git a/packages/langium/src/utils/ast-util.ts b/packages/langium/src/utils/ast-util.ts
@@ -5,7 +5,7 @@
  ******************************************************************************/
 
 import type { Range } from 'vscode-languageserver';
-import type { AstNode, CstNode, GenericAstNode, Reference, ReferenceInfo } from '../syntax-tree.js';
+import type { AstNode, AstReflection, CstNode, GenericAstNode, Reference, ReferenceInfo } from '../syntax-tree.js';
 import type { Stream, TreeStream } from '../utils/stream.js';
 import type { LangiumDocument } from '../workspace/documents.js';
 import { isAstNode, isReference } from '../syntax-tree.js';
@@ -232,6 +232,25 @@ export function findLocalReferences(targetNode: AstNode, lookup = getDocument(ta
     return stream(refs);
 }
 
+/**
+ * Assigns all mandatory AST properties to the specified node.
+ *
+ * @param reflection Reflection object used to gather mandatory properties for the node.
+ * @param node Specified node is modified in place and properties are directly assigned.
+ */
+export function assignMandatoryAstProperties(reflection: AstReflection, node: AstNode): void {
+    const typeMetaData = reflection.getTypeMetaData(node.$type);
+    const genericNode = node as GenericAstNode;
+    for (const mandatoryProperty of typeMetaData.mandatory) {
+        const value = genericNode[mandatoryProperty.name];
+        if (mandatoryProperty.type === 'array' && !Array.isArray(value)) {
+            genericNode[mandatoryProperty.name] = [];
+        } else if (mandatoryProperty.type === 'boolean' && value === undefined) {
+            genericNode[mandatoryProperty.name] = false;
+        }
+    }
+}
+
 /**
  * Creates a deep copy of the specified AST node.
  * The resulting copy will only contain semantically relevant information, such as the `$type` property and AST properties.
