Improve source range display on definition request (#2063)

Author: msujew

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

@@ -24,7 +24,7 @@ import type { MarkupContent } from 'vscode-languageserver';
 import { CompletionItemKind, CompletionList, Position } from 'vscode-languageserver';
 import * as ast from '../../languages/generated/ast.js';
 import { assignMandatoryProperties, getContainerOfType } from '../../utils/ast-utils.js';
-import { findDeclarationNodeAtOffset, findLeafNodeBeforeOffset } from '../../utils/cst-utils.js';
+import { findDeclarationNodeAtOffset, findLeafNodeBeforeOffset, getDatatypeNode } from '../../utils/cst-utils.js';
 import { getEntryRule, getExplicitRuleType } from '../../utils/grammar-utils.js';
 import { stream, type Stream } from '../../utils/stream.js';
 import { findFirstFeatures, findNextFeatures } from './follow-element-computation.js';
@@ -326,18 +326,14 @@ export class DefaultCompletionProvider implements CompletionProvider {
     }
 
     protected findDataTypeRuleStart(cst: CstNode, offset: number): [number, number] | undefined {
-        let containerNode: CstNode | undefined = findDeclarationNodeAtOffset(cst, offset, this.grammarConfig.nameRegexp);
+        const containerNode = findDeclarationNodeAtOffset(cst, offset, this.grammarConfig.nameRegexp);
+        if (!containerNode) {
+            return undefined;
+        }
         // Identify whether the element was parsed as part of a data type rule
-        let isDataTypeNode = Boolean(getContainerOfType(containerNode?.grammarSource, ast.isParserRule)?.dataType);
-        if (isDataTypeNode) {
-            while (isDataTypeNode) {
-                // Use the container to find the correct parent element
-                containerNode = containerNode?.container;
-                isDataTypeNode = Boolean(getContainerOfType(containerNode?.grammarSource, ast.isParserRule)?.dataType);
-            }
-            if (containerNode) {
-                return [containerNode.offset, containerNode.end];
-            }
+        const fullNode = getDatatypeNode(containerNode);
+        if (fullNode) {
+            return [fullNode.offset, fullNode.end];
         }
         return undefined;
     }

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

@@ -15,7 +15,7 @@ import type { MaybePromise } from '../utils/promise-utils.js';
 import type { LangiumDocument } from '../workspace/documents.js';
 import { LocationLink } from 'vscode-languageserver';
 import { getDocument } from '../utils/ast-utils.js';
-import { findDeclarationNodeAtOffset } from '../utils/cst-utils.js';
+import { findDeclarationNodeAtOffset, getDatatypeNode } from '../utils/cst-utils.js';
 
 /**
  * Language-specific service for handling go to definition requests.
@@ -79,12 +79,17 @@ export class DefaultDefinitionProvider implements DefinitionProvider {
     }
 
     protected findLinks(source: CstNode): GoToLink[] {
+        const datatypeSourceNode = getDatatypeNode(source) ?? source;
         const targets = this.references.findDeclarationNodes(source);
         const links: GoToLink[] = [];
         for (const target of targets) {
             const targetDocument = getDocument(target.astNode);
             if (targets && targetDocument) {
-                links.push({ source, target, targetDocument });
+                links.push({
+                    source: datatypeSourceNode,
+                    target,
+                    targetDocument
+                });
             }
         }
         return links;

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

@@ -11,6 +11,34 @@ import type { DocumentSegment } from '../workspace/documents.js';
 import type { Stream, TreeStream } from './stream.js';
 import { isCompositeCstNode, isLeafCstNode, isRootCstNode } from '../syntax-tree.js';
 import { TreeStreamImpl } from './stream.js';
+import { getContainerOfType } from './ast-utils.js';
+import { isParserRule } from '../languages/generated/ast.js';
+
+/**
+ * Attempts to find the CST node that belongs to the datatype element that contains the given CST node.
+ *
+ * @param cstNode The CST node for which to find the datatype node.
+ * @returns The CST node corresponding to the datatype element, or the undefined if no such element exists.
+ */
+export function getDatatypeNode(cstNode: CstNode): CstNode | undefined {
+    let current: CstNode | undefined = cstNode;
+    let found = false;
+    while (current) {
+        const definingRule = getContainerOfType(current.grammarSource, isParserRule);
+        if (definingRule && definingRule.dataType) {
+            // Go up the chain. This element might be part of a larger datatype rule
+            current = current.container;
+            found = true;
+        } else if (found) {
+            // The last datatype node is the one we are looking for
+            return current;
+        } else {
+            // We haven't found any datatype node yet and we've reached a non-datatype rule
+            return undefined;
+        }
+    }
+    return undefined;
+}
 
 /**
  * Create a stream of all CST nodes that are directly and indirectly contained in the given root node,

packages/langium/test/lsp/goto-definition.test.ts

@@ -4,10 +4,12 @@
  * terms of the MIT License, which is available in the project root.
  ******************************************************************************/
 
-import { describe, test } from 'vitest';
-import { EmptyFileSystem } from 'langium';
+import { describe, expect, test } from 'vitest';
+import { EmptyFileSystem, URI } from 'langium';
 import { createLangiumGrammarServices, createServicesForGrammar } from 'langium/grammar';
 import { expectGoToDefinition } from 'langium/test';
+import { expandToString } from 'langium/generate';
+import type { Range } from 'vscode-languageserver';
 
 /**
  * Represents a grammar file
@@ -113,6 +115,57 @@ describe('Definition Provider', () => {
             });
         });
     });
+
+    test('Should highlight full datatype rule node', async () => {
+        const grammar = `
+        grammar Test
+        entry Model: (elements+=Element)*;
+        Element: Source | Target;
+        Source: 'source' name=FQN;
+        Target: 'target' ref=[Source];
+        FQN returns string: ID ('.' ID)*;
+        terminal ID: /\\w+/;
+        hidden terminal WS: /\\s+/;
+        `;
+        const services = await createServicesForGrammar({ grammar });
+        const text = expandToString`
+            target a.b.c;
+            source a.b.c;
+        `;
+        const workspace = services.shared.workspace;
+        const document = workspace.LangiumDocumentFactory.fromString(text, URI.file('test.txt'));
+        workspace.LangiumDocuments.addDocument(document);
+        await workspace.DocumentBuilder.build([document]);
+        const targetTextRange: Range = {
+            start: document.textDocument.positionAt(text.indexOf('a.b.c')),
+            end: document.textDocument.positionAt(text.indexOf('a.b.c') + 'a.b.c'.length)
+        };
+        const sourceTextRange: Range = {
+            start: document.textDocument.positionAt(text.lastIndexOf('a.b.c')),
+            end: document.textDocument.positionAt(text.lastIndexOf('a.b.c') + 'a.b.c'.length)
+        };
+        const provider = services.lsp.DefinitionProvider!;
+        // Go to definition from target to source
+        const defFromTarget = await provider.getDefinition(document, {
+            textDocument: { uri: document.uri.toString() },
+            position: targetTextRange.start,
+        });
+        expect(defFromTarget).toBeDefined();
+        expect(defFromTarget).toHaveLength(1);
+        const targetSourceRange = defFromTarget![0].originSelectionRange!;
+        expect(targetSourceRange).toBeDefined();
+        expect(targetSourceRange).toEqual(targetTextRange);
+        // Go to definition from target to itself
+        const defFromSource = await provider.getDefinition(document, {
+            textDocument: { uri: document.uri.toString() },
+            position: sourceTextRange.start,
+        });
+        expect(defFromSource).toBeDefined();
+        expect(defFromSource).toHaveLength(1);
+        const sourceRange = defFromSource![0].originSelectionRange!;
+        expect(sourceRange).toBeDefined();
+        expect(sourceRange).toEqual(sourceTextRange);
+    });
 });
 
 describe('Definition Provider with Infix Operators', async () => {