Merge pull request #574 from jvalue/import-specific-element

Import Specific Element

libs/language-server/src/grammar/main.langium

@@ -36,4 +36,7 @@ ExportDefinition:
   'publish' element=[ExportableElement] ('as' alias=ID)? ';';
 
 ImportDefinition:
-  'use' '*' 'from' path=STRING ';';
+  'use' (
+    useAll?='*' 
+    | '{' (usedElements+=ID) (',' usedElements+=ID)* '}'
+  ) 'from' path=STRING ';';

libs/language-server/src/lib/ast/model-util.ts

@@ -2,13 +2,22 @@
 //
 // SPDX-License-Identifier: AGPL-3.0-only
 
+// eslint-disable-next-line unicorn/prefer-node-protocol
+import { strict as assert } from 'assert';
+
 import { type AstNode, AstUtils, type LangiumDocuments } from 'langium';
 
 import {
   type BuiltinBlockTypeDefinition,
   type BuiltinConstrainttypeDefinition,
+  type ExportDefinition,
+  type ExportableElement,
+  type JayveeModel,
   isBuiltinBlockTypeDefinition,
   isBuiltinConstrainttypeDefinition,
+  isExportDefinition,
+  isExportableElement,
+  isExportableElementDefinition,
   isJayveeModel,
 } from './generated/ast';
 import {
@@ -121,3 +130,89 @@ export function getAllBuiltinConstraintTypes(
     });
   return allBuiltinConstraintTypes;
 }
+
+export interface ExportDetails {
+  /**
+   * The exported element
+   */
+  element: ExportableElement;
+
+  /**
+   * The name which the exported element is available under.
+   */
+  alias: string;
+}
+
+/**
+ * Gets all exported elements from a document.
+ * This logic cannot reside in a {@link ScopeComputationProvider} but should be handled here:
+ * https://github.com/eclipse-langium/langium/discussions/1508#discussioncomment-9524544
+ */
+export function getExportedElements(model: JayveeModel): ExportDetails[] {
+  const exportedElements: ExportDetails[] = [];
+
+  for (const node of AstUtils.streamAllContents(model)) {
+    if (isExportableElementDefinition(node) && node.isPublished) {
+      assert(
+        isExportableElement(node),
+        'Exported node is not an ExportableElement',
+      );
+      exportedElements.push({
+        element: node,
+        alias: node.name,
+      });
+    }
+
+    if (isExportDefinition(node)) {
+      const originalDefinition = followExportDefinitionChain(node);
+      if (originalDefinition !== undefined) {
+        const exportName = node.alias ?? originalDefinition.name;
+        exportedElements.push({
+          element: originalDefinition,
+          alias: exportName,
+        });
+      }
+    }
+  }
+  return exportedElements;
+}
+
+/**
+ * Follow an export statement to its original definition.
+ */
+export function followExportDefinitionChain(
+  exportDefinition: ExportDefinition,
+): ExportableElement | undefined {
+  const referenced = exportDefinition.element.ref;
+
+  if (referenced === undefined) {
+    return undefined; // Cannot follow reference to original definition
+  }
+
+  if (!isElementExported(referenced)) {
+    return undefined;
+  }
+
+  return referenced; // Reached original definition
+}
+
+/**
+ * Checks whether an exportable @param element is exported (either in definition or via an delayed export definition).
+ */
+export function isElementExported(element: ExportableElement): boolean {
+  if (isExportableElementDefinition(element) && element.isPublished) {
+    return true;
+  }
+
+  const model = AstUtils.getContainerOfType(element, isJayveeModel);
+  assert(
+    model !== undefined,
+    // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
+    `Could not get container of exportable element ${element.name ?? ''}`,
+  );
+
+  const isExported = model.exports.some(
+    (exportDefinition) => exportDefinition.element.ref === element,
+  );
+  return isExported;
+}

libs/language-server/src/lib/lsp/jayvee-completion-provider.ts

@@ -41,20 +41,24 @@ import {
 import {
   getAllBuiltinBlockTypes,
   getAllBuiltinConstraintTypes,
+  getExportedElements,
 } from '../ast/model-util';
 import { LspDocGenerator } from '../docs/lsp-doc-generator';
 import { type JayveeServices } from '../jayvee-module';
+import { type JayveeImportResolver } from '../services';
 
 const RIGHT_ARROW_SYMBOL = '\u{2192}';
 
 export class JayveeCompletionProvider extends DefaultCompletionProvider {
   protected langiumDocuments: LangiumDocuments;
   protected readonly wrapperFactories: WrapperFactoryProvider;
+  protected readonly importResolver: JayveeImportResolver;
 
   constructor(services: JayveeServices) {
     super(services);
     this.langiumDocuments = services.shared.workspace.LangiumDocuments;
     this.wrapperFactories = services.WrapperFactories;
+    this.importResolver = services.ImportResolver;
   }
 
   override completionFor(
@@ -94,6 +98,12 @@ export class JayveeCompletionProvider extends DefaultCompletionProvider {
       if (isImportPathCompletion) {
         return this.completionForImportPath(astNode, context, acceptor);
       }
+
+      const isImportElementCompletion =
+        isImportDefinition(astNode) && next.property === 'usedElements';
+      if (isImportElementCompletion) {
+        return this.completionForImportElement(astNode, context, acceptor);
+      }
     }
     return super.completionFor(context, next, acceptor);
   }
@@ -261,6 +271,47 @@ export class JayveeCompletionProvider extends DefaultCompletionProvider {
     }
   }
 
+  private completionForImportElement(
+    importDefinition: ImportDefinition,
+    context: CompletionContext,
+    acceptor: CompletionAcceptor,
+  ) {
+    const resolvedModel = this.importResolver.resolveImport(importDefinition);
+    if (resolvedModel === undefined) {
+      return;
+    }
+
+    const documentText = context.textDocument.getText();
+    const existingElementName = documentText.substring(
+      context.tokenOffset,
+      context.offset,
+    );
+
+    const exportedElementNames = getExportedElements(resolvedModel).map(
+      (x) => x.alias,
+    );
+
+    const suggestedElementNames = exportedElementNames.filter((x) =>
+      x.startsWith(existingElementName),
+    );
+
+    const insertRange: Range = {
+      start: context.textDocument.positionAt(context.tokenOffset),
+      end: context.textDocument.positionAt(context.tokenEndOffset),
+    };
+
+    for (const elementName of suggestedElementNames) {
+      acceptor(context, {
+        label: elementName,
+        textEdit: {
+          newText: elementName,
+          range: insertRange,
+        },
+        kind: CompletionItemKind.Reference,
+      });
+    }
+  }
+
   /**
    * Gets all paths to available documents, formatted as relative paths.
    * Does not include path to stdlib files as they don't need to be imported.

libs/language-server/src/lib/lsp/jayvee-definition-provider.ts

@@ -2,7 +2,12 @@
 //
 // SPDX-License-Identifier: AGPL-3.0-only
 
+// eslint-disable-next-line unicorn/prefer-node-protocol
+import { strict as assert } from 'assert';
+
 import {
+  type AstNode,
+  AstUtils,
   GrammarUtils,
   type LangiumDocuments,
   type LeafCstNode,
@@ -15,7 +20,7 @@ import {
   Range,
 } from 'vscode-languageserver-protocol';
 
-import { isImportDefinition } from '../ast';
+import { getExportedElements, isImportDefinition, isJayveeModel } from '../ast';
 import { type JayveeServices } from '../jayvee-module';
 import { type JayveeImportResolver } from '../services/import-resolver';
 
@@ -34,6 +39,7 @@ export class JayveeDefinitionProvider extends DefaultDefinitionProvider {
     params: DefinitionParams,
   ): MaybePromise<LocationLink[] | undefined> {
     const sourceAstNode = sourceCstNode.astNode;
+
     if (
       isImportDefinition(sourceAstNode) &&
       GrammarUtils.findAssignment(sourceCstNode)?.feature === 'path'
@@ -44,23 +50,76 @@ export class JayveeDefinitionProvider extends DefaultDefinitionProvider {
         return undefined;
       }
 
-      const jumpTarget = importedModel;
-
-      const selectionRange =
-        this.nameProvider.getNameNode(jumpTarget)?.range ??
-        Range.create(0, 0, 0, 0);
-      const previewRange =
-        jumpTarget.$cstNode?.range ?? Range.create(0, 0, 0, 0);
-
-      return [
-        LocationLink.create(
-          importedModel.$document.uri.toString(),
-          previewRange,
-          selectionRange,
-          sourceCstNode.range,
-        ),
-      ];
+      return this.getLocationLink(sourceCstNode, importedModel);
+    }
+
+    if (
+      isImportDefinition(sourceAstNode) &&
+      GrammarUtils.findAssignment(sourceCstNode)?.feature === 'usedElements'
+    ) {
+      const clickedIndex =
+        GrammarUtils.findAssignment(sourceCstNode)?.$container?.$containerIndex;
+      assert(
+        clickedIndex !== undefined,
+        'Could not read index of selected element',
+      );
+      const indexOfElement = clickedIndex - 1;
+      assert(
+        indexOfElement < sourceAstNode.usedElements.length,
+        'Index of selected element is not correctly computed',
+      );
+      const refString = sourceAstNode.usedElements[indexOfElement];
+      assert(
+        refString !== undefined,
+        'Could not read reference text to imported element',
+      );
+
+      const importedModel = this.importResolver.resolveImport(sourceAstNode);
+
+      if (importedModel?.$document === undefined) {
+        return undefined;
+      }
+
+      const allExportDefinitions = getExportedElements(importedModel);
+
+      const referencedExport = allExportDefinitions.find((x) => {
+        return x.alias === refString;
+      });
+      if (referencedExport === undefined) {
+        return;
+      }
+
+      return this.getLocationLink(sourceCstNode, referencedExport.element);
     }
     return super.collectLocationLinks(sourceCstNode, params);
   }
+
+  protected getLocationLink(
+    sourceCstNode: LeafCstNode,
+    jumpTarget: AstNode,
+  ): LocationLink[] | undefined {
+    // need to go over model as jumpTarget might not have $document associated
+    const containingDocument = AstUtils.getContainerOfType(
+      jumpTarget,
+      isJayveeModel,
+    )?.$document;
+
+    if (containingDocument === undefined) {
+      return undefined;
+    }
+
+    const selectionRange =
+      this.nameProvider.getNameNode(jumpTarget)?.range ??
+      Range.create(0, 0, 0, 0);
+    const previewRange = jumpTarget.$cstNode?.range ?? Range.create(0, 0, 0, 0);
+
+    return [
+      LocationLink.create(
+        containingDocument.uri.toString(),
+        previewRange,
+        selectionRange,
+        sourceCstNode.range,
+      ),
+    ];
+  }
 }