Add code action to import unimported element

libs/language-server/src/lib/jayvee-module.ts

@@ -27,6 +27,7 @@ import { WrapperFactoryProvider } from './ast/wrappers/wrapper-factory-provider'
 import { JayveeWorkspaceManager } from './builtin-library/jayvee-workspace-manager';
 import { JayveeValueConverter } from './jayvee-value-converter';
 import {
+  JayveeCodeActionProvider,
   JayveeCompletionProvider,
   JayveeDefinitionProvider,
   JayveeFormatter,
@@ -86,6 +87,7 @@ export const JayveeModule: Module<
       new JayveeHoverProvider(services),
     Formatter: () => new JayveeFormatter(),
     DefinitionProvider: (services) => new JayveeDefinitionProvider(services),
+    CodeActionProvider: (services) => new JayveeCodeActionProvider(services),
   },
   references: {
     ScopeProvider: (services) => new JayveeScopeProvider(services),

libs/language-server/src/lib/lsp/index.ts

@@ -8,3 +8,4 @@ export * from './jayvee-hover-provider';
 export * from './jayvee-scope-computation';
 export * from './jayvee-scope-provider';
 export * from './jayvee-definition-provider';
+export * from './jayvee-code-action-provider';

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

@@ -0,0 +1,179 @@
+// eslint-disable-next-line unicorn/prefer-node-protocol
+import { strict as assert } from 'assert';
+
+import {
+  type AstReflection,
+  DocumentValidator,
+  type IndexManager,
+  type LangiumDocument,
+  type LinkingErrorData,
+  type MaybePromise,
+  type Reference,
+  type ReferenceInfo,
+  type URI,
+  UriUtils,
+} from 'langium';
+import { type CodeActionProvider } from 'langium/lsp';
+import {
+  type CodeAction,
+  CodeActionKind,
+  type CodeActionParams,
+  type Command,
+  type Diagnostic,
+  type Position,
+} from 'vscode-languageserver-protocol';
+
+import { type JayveeModel } from '../ast';
+import { type JayveeServices } from '../jayvee-module';
+
+export class JayveeCodeActionProvider implements CodeActionProvider {
+  protected readonly reflection: AstReflection;
+  protected readonly indexManager: IndexManager;
+
+  constructor(services: JayveeServices) {
+    this.reflection = services.shared.AstReflection;
+    this.indexManager = services.shared.workspace.IndexManager;
+  }
+
+  getCodeActions(
+    document: LangiumDocument,
+    params: CodeActionParams,
+  ): MaybePromise<Array<Command | CodeAction>> {
+    const actions: CodeAction[] = [];
+
+    for (const diagnostic of params.context.diagnostics) {
+      const diagnosticActions = this.getCodeActionsForDiagnostic(
+        diagnostic,
+        document,
+      );
+      actions.push(...diagnosticActions);
+    }
+    return actions;
+  }
+
+  protected getCodeActionsForDiagnostic(
+    diagnostic: Diagnostic,
+    document: LangiumDocument,
+  ): CodeAction[] {
+    const actions: CodeAction[] = [];
+
+    const diagnosticData = diagnostic.data as unknown;
+    const diagnosticCode = (diagnosticData as { code?: string } | undefined)
+      ?.code;
+    if (diagnosticData === undefined || diagnosticCode === undefined) {
+      return actions;
+    }
+
+    switch (diagnosticCode) {
+      case DocumentValidator.LinkingError: {
+        const linkingData = diagnosticData as LinkingErrorData;
+        actions.push(
+          ...this.getCodeActionsForLinkingError(
+            diagnostic,
+            linkingData,
+            document,
+          ),
+        );
+      }
+    }
+
+    return actions;
+  }
+
+  protected getCodeActionsForLinkingError(
+    diagnostic: Diagnostic,
+    linkingData: LinkingErrorData,
+    document: LangiumDocument,
+  ): CodeAction[] {
+    const refInfo: ReferenceInfo = {
+      container: {
+        $type: linkingData.containerType,
+      },
+      property: linkingData.property,
+      reference: {
+        $refText: linkingData.refText,
+      } as Reference,
+    };
+    const refType = this.reflection.getReferenceType(refInfo);
+    const importCandidates = this.indexManager
+      .allElements(refType)
+      .filter((e) => e.name === linkingData.refText);
+
+    const actions: CodeAction[] = [];
+    for (const importCandidate of importCandidates) {
+      const isInCurrentFile = UriUtils.equals(
+        importCandidate.documentUri,
+        document.uri,
+      );
+      if (isInCurrentFile) {
+        continue;
+      }
+
+      const importPath = this.getRelativeImportPath(
+        document.uri,
+        importCandidate.documentUri,
+      );
+
+      const importPosition = this.getImportLinePosition(
+        document.parseResult.value as JayveeModel,
+      );
+      if (importPosition === undefined) {
+        continue;
+      }
+
+      actions.push({
+        title: `Use from '${importPath}'`,
+        kind: CodeActionKind.QuickFix,
+        diagnostics: [diagnostic],
+        isPreferred: false,
+        edit: {
+          changes: {
+            [document.textDocument.uri]: [
+              {
+                range: {
+                  start: importPosition,
+                  end: importPosition,
+                },
+                newText: `use * from "${importPath}";\n`,
+              },
+            ],
+          },
+        },
+      });
+    }
+
+    return actions;
+  }
+
+  protected getImportLinePosition(
+    javeeModel: JayveeModel,
+  ): Position | undefined {
+    const currentModelImports = javeeModel.imports;
+
+    // Put the new import after the last import
+    if (currentModelImports.length > 0) {
+      const lastImportEnd =
+        currentModelImports[currentModelImports.length - 1]?.$cstNode?.range
+          .end;
+      assert(
+        lastImportEnd !== undefined,
+        'Could not find end of last import statement.',
+      );
+      return { line: lastImportEnd.line + 1, character: 0 };
+    }
+
+    // For now, we just add it in the first row if there is no import yet
+    return { line: 0, character: 0 };
+  }
+
+  private getRelativeImportPath(source: URI, target: URI): string {
+    const sourceDir = UriUtils.dirname(source);
+    const relativePath = UriUtils.relative(sourceDir, target);
+
+    if (!relativePath.startsWith('./') && !relativePath.startsWith('../')) {
+      return `./${relativePath}`;
+    }
+
+    return relativePath;
+  }
+}