Repository: elm-tooling/elm-language-server
Pull request #381 (2020-11-30 06:48:03)

Fix import code actions
Adds a new system for code actions using a registration model that should be nicer and easy to handle "fix all" code actions. Moves the import code action out of the `elmMakeDiagnostics` and into a shared code action registration that is used for type inference diagnostics as well. In the future I will work to consolidate other code actions to this new way. Also started adding tests for code actions. I assume it will need some cleanup, but its most of the way there. 

Another thing I changed was the diagnostics behavior between the compiler and our analyzer. Now, when you save, we clear our diagnostics and get them only from the compiler. Then if you make a change, we clear the compiler ones and only get ours. We can talk about this more if we need to.


Patch:
diff --git a/src/elmWorkspace.ts b/src/elmWorkspace.ts
index ea8631b..88c2623 100644
--- a/src/elmWorkspace.ts
+++ b/src/elmWorkspace.ts
@@ -4,10 +4,11 @@ import os from "os";
 import path from "path";
 import { container } from "tsyringe";
 import util from "util";
-import { Connection } from "vscode-languageserver";
+import { Connection, Diagnostic } from "vscode-languageserver";
 import { URI } from "vscode-uri";
 import Parser, { Tree } from "web-tree-sitter";
-import { Forest, IForest } from "./forest";
+import { ICancellationToken } from "./cancellation";
+import { Forest, IForest, ITreeContainer } from "./forest";
 import * as utils from "./util/elmUtils";
 import {
   IPossibleImportsCache,
@@ -83,6 +84,14 @@ export interface IElmWorkspace {
   markAsDirty(): void;
   getPossibleImportsCache(): IPossibleImportsCache;
   getOperatorsCache(): Map<string, DefinitionResult>;
+  getDiagnostics(
+    sourceFile: ITreeContainer,
+    cancellationToken?: ICancellationToken,
+  ): Diagnostic[];
+  getDiagnosticsAsync(
+    sourceFile: ITreeContainer,
+    cancellationToken?: ICancellationToken,
+  ): Promise<Diagnostic[]>;
 }
 
 export interface IRootFolder {
@@ -101,6 +110,7 @@ export class ElmWorkspace implements IElmWorkspace {
   private dirty = true;
   private possibleImportsCache: IPossibleImportsCache;
   private operatorsCache: Map<string, DefinitionResult>;
+  private diagnosticsCache: Map<string, Diagnostic[]>;
 
   constructor(private rootPath: URI) {
     this.settings = container.resolve("Settings");
@@ -113,6 +123,7 @@ export class ElmWorkspace implements IElmWorkspace {
     this.typeCache = new TypeCache();
     this.possibleImportsCache = new PossibleImportsCache();
     this.operatorsCache = new Map<string, DefinitionResult>();
+    this.diagnosticsCache = new Map<string, Diagnostic[]>();
   }
 
   public async init(
@@ -165,6 +176,7 @@ export class ElmWorkspace implements IElmWorkspace {
     if (!this.dirty) {
       this.dirty = true;
       this.typeChecker = undefined;
+      this.diagnosticsCache.clear();
     }
   }
 
@@ -176,6 +188,44 @@ export class ElmWorkspace implements IElmWorkspace {
     return this.operatorsCache;
   }
 
+  public getDiagnostics(
+    sourceFile: ITreeContainer,
+    cancellationToken?: ICancellationToken,
+  ): Diagnostic[] {
+    const cached = this.diagnosticsCache.get(sourceFile.uri);
+
+    if (cached) {
+      return cached;
+    }
+
+    const diagnostics = this.getTypeChecker().getDiagnostics(
+      sourceFile,
+      cancellationToken,
+    );
+
+    this.diagnosticsCache.set(sourceFile.uri, diagnostics);
+    return diagnostics;
+  }
+
+  public async getDiagnosticsAsync(
+    sourceFile: ITreeContainer,
+    cancellationToken?: ICancellationToken,
+  ): Promise<Diagnostic[]> {
+    const cached = this.diagnosticsCache.get(sourceFile.uri);
+
+    if (cached) {
+      return Promise.resolve(cached);
+    }
+
+    const diagnostics = await this.getTypeChecker().getDiagnosticsAsync(
+      sourceFile,
+      cancellationToken,
+    );
+
+    this.diagnosticsCache.set(sourceFile.uri, diagnostics);
+    return diagnostics;
+  }
+
   private async initWorkspace(
     progressCallback: (percent: number) => void,
   ): Promise<void> {
diff --git a/src/providers/codeAction/importCodeAction.ts b/src/providers/codeAction/importCodeAction.ts
new file mode 100644
index 0000000..c7ae7b9
--- /dev/null
+++ b/src/providers/codeAction/importCodeAction.ts
@@ -0,0 +1,122 @@
+/* eslint-disable @typescript-eslint/no-use-before-define */
+import { CodeAction, TextEdit } from "vscode-languageserver";
+import { Range } from "vscode-languageserver-textdocument";
+import { SyntaxNode } from "web-tree-sitter";
+import { ITreeContainer } from "../../forest";
+import { ImportUtils, IPossibleImport } from "../../util/importUtils";
+import { RefactorEditUtils } from "../../util/refactorEditUtils";
+import { TreeUtils } from "../../util/treeUtils";
+import { Diagnostics } from "../../util/types/diagnostics";
+import { CodeActionProvider, ICodeActionParams } from "../codeActionProvider";
+
+const errorCodes = [Diagnostics.MissingValue.code];
+
+CodeActionProvider.registerCodeAction({
+  errorCodes,
+  getCodeActions: (params: ICodeActionParams): CodeAction[] | undefined => {
+    const valueNode = TreeUtils.getNamedDescendantForRange(
+      params.sourceFile,
+      params.range,
+    );
+
+    return getPossibleImports(params, params.range).map((possibleImport) => {
+      const edit = getEditFromPossibleImport(
+        params.sourceFile,
+        params.range,
+        possibleImport,
+      );
+
+      const valueToImport = getValueToImport(valueNode, possibleImport);
+
+      return CodeActionProvider.getCodeAction(
+        params,
+        valueToImport
+          ? `Import '${valueToImport}' from module "${possibleImport.module}"`
+          : `Import module "${possibleImport.module}"`,
+        edit ? [edit] : [],
+      );
+    });
+  },
+  getFixAllCodeAction: (params: ICodeActionParams) => {
+    return CodeActionProvider.getFixAllCodeAction(
+      "Add all missing imports",
+      params,
+      errorCodes,
+      (edits, diagnostic) => {
+        const firstPossibleImport = getPossibleImports(
+          params,
+          diagnostic.range,
+        )[0];
+
+        if (firstPossibleImport) {
+          const edit = getEditFromPossibleImport(
+            params.sourceFile,
+            diagnostic.range,
+            firstPossibleImport,
+          );
+
+          if (edit) {
+            edits.push(edit);
+          }
+        }
+      },
+    );
+  },
+});
+
+function getPossibleImports(
+  params: ICodeActionParams,
+  range: Range,
+): IPossibleImport[] {
+  const valueNode = TreeUtils.getNamedDescendantForRange(
+    params.sourceFile,
+    range,
+  );
+
+  const possibleImports = ImportUtils.getPossibleImports(
+    params.program.getForest(),
+    params.sourceFile.uri,
+  );
+
+  // Add import quick fixes
+  if (valueNode) {
+    return possibleImports.filter(
+      (exposed) =>
+        exposed.value === valueNode.text ||
+        ((valueNode.type === "upper_case_qid" ||
+          valueNode.type === "value_qid") &&
+          exposed.value ===
+            valueNode.namedChildren[valueNode.namedChildren.length - 1].text &&
+          exposed.module ===
+            valueNode.namedChildren
+              .slice(0, valueNode.namedChildren.length - 2) // Dots are also namedNodes
+              .map((a) => a.text)
+              .join("")),
+    );
+  }
+
+  return [];
+}
+
+function getEditFromPossibleImport(
+  sourceFile: ITreeContainer,
+  range: Range,
+  possibleImport: IPossibleImport,
+): TextEdit | undefined {
+  const valueNode = TreeUtils.getNamedDescendantForRange(sourceFile, range);
+
+  return RefactorEditUtils.addImport(
+    sourceFile.tree,
+    possibleImport.module,
+    getValueToImport(valueNode, possibleImport),
+  );
+}
+
+function getValueToImport(
+  valueNode: SyntaxNode,
+  possibleImport: IPossibleImport,
+): string | undefined {
+  return valueNode.type !== "upper_case_qid" && valueNode.type !== "value_qid"
+    ? possibleImport.valueToImport ?? possibleImport.value
+    : undefined;
+}
diff --git a/src/providers/codeAction/index.ts b/src/providers/codeAction/index.ts
new file mode 100644
index 0000000..2bc047a
--- /dev/null
+++ b/src/providers/codeAction/index.ts
@@ -0,0 +1,1 @@
+import "./importCodeAction";
diff --git a/src/providers/codeActionProvider.ts b/src/providers/codeActionProvider.ts
index f1f2751..b6354aa 100644
--- a/src/providers/codeActionProvider.ts
+++ b/src/providers/codeActionProvider.ts
@@ -10,21 +10,37 @@ import {
 import { URI } from "vscode-uri";
 import { SyntaxNode, Tree } from "web-tree-sitter";
 import { IElmWorkspace } from "../elmWorkspace";
-import { ElmWorkspaceMatcher } from "../util/elmWorkspaceMatcher";
+import { ElmWorkspaceMatcher, IParams } from "../util/elmWorkspaceMatcher";
+import { MultiMap } from "../util/multiMap";
 import { RefactorEditUtils } from "../util/refactorEditUtils";
 import { Settings } from "../util/settings";
-import { TreeUtils } from "../util/treeUtils";
+import { flatMap, TreeUtils } from "../util/treeUtils";
+import { Diagnostics } from "../util/types/diagnostics";
 import { ElmLsDiagnostics } from "./diagnostics/elmLsDiagnostics";
 import { ElmMakeDiagnostics } from "./diagnostics/elmMakeDiagnostics";
+import { diagnosticsEquals } from "./diagnostics/fileDiagnostics";
 import { ExposeUnexposeHandler } from "./handlers/exposeUnexposeHandler";
 import { MoveRefactoringHandler } from "./handlers/moveRefactoringHandler";
 
+export type ICodeActionParams = CodeActionParams & IParams;
+
+export interface ICodeActionRegistration {
+  errorCodes: string[];
+  getCodeActions(params: ICodeActionParams): CodeAction[] | undefined;
+  getFixAllCodeAction(params: ICodeActionParams): CodeAction | undefined;
+}
+
 export class CodeActionProvider {
   private connection: Connection;
   private settings: Settings;
   private elmMake: ElmMakeDiagnostics;
   private elmDiagnostics: ElmLsDiagnostics;
 
+  private static errorCodeToRegistrationMap = new MultiMap<
+    string,
+    ICodeActionRegistration
+  >();
+
   constructor() {
     this.settings = container.resolve("Settings");
     this.elmMake = container.resolve(ElmMakeDiagnostics);
@@ -35,7 +51,7 @@ export class CodeActionProvider {
     this.connection.onCodeAction(
       new ElmWorkspaceMatcher((param: CodeActionParams) =>
         URI.parse(param.textDocument.uri),
-      ).handlerForWorkspace(this.onCodeAction.bind(this)),
+      ).handle(this.onCodeAction.bind(this)),
     );
 
     if (this.settings.extendedCapabilities?.moveFunctionRefactoringSupport) {
@@ -45,37 +61,144 @@ export class CodeActionProvider {
     new ExposeUnexposeHandler();
   }
 
-  private onCodeAction(
-    params: CodeActionParams,
-    elmWorkspace: IElmWorkspace,
-  ): CodeAction[] {
+  public static registerCodeAction(
+    registration: ICodeActionRegistration,
+  ): void {
+    registration.errorCodes.forEach((code) => {
+      CodeActionProvider.errorCodeToRegistrationMap.set(code, registration);
+    });
+  }
+
+  private static forEachDiagnostic(
+    params: ICodeActionParams,
+    errorCodes: string[],
+    callback: (diagnostic: Diagnostic) => void,
+  ): void {
+    params.program.getDiagnostics(params.sourceFile).forEach((diagnostic) => {
+      if (
+        typeof diagnostic.code === "string" &&
+        errorCodes.includes(diagnostic.code)
+      ) {
+        callback(diagnostic);
+      }
+    });
+  }
+
+  public static getCodeAction(
+    params: ICodeActionParams,
+    title: string,
+    edits: TextEdit[],
+  ): CodeAction {
+    const changes = { [params.sourceFile.uri]: edits };
+    return {
+      title,
+      kind: CodeActionKind.QuickFix,
+      edit: { changes },
+      isPreferred: true,
+    };
+  }
+
+  public static getFixAllCodeAction(
+    title: string,
+    params: ICodeActionParams,
+    errorCodes: string[],
+    callback: (edits: TextEdit[], diagnostic: Diagnostic) => void,
+  ): CodeAction {
+    const edits: TextEdit[] = [];
+    const changes = {
+      [params.sourceFile.uri]: edits,
+    };
+
+    const diagnostics: Diagnostic[] = [];
+    CodeActionProvider.forEachDiagnostic(params, errorCodes, (diagnostic) => {
+      diagnostics.push(diagnostic);
+      callback(edits, diagnostic);
+    });
+
+    return {
+      title,
+      kind: CodeActionKind.SourceFixAll,
+      diagnostics,
+      edit: { changes },
+    };
+  }
+
+  protected onCodeAction(params: ICodeActionParams): CodeAction[] | undefined {
     this.connection.console.info("A code action was requested");
     const make = this.elmMake.onCodeAction(params);
     const elmDiagnostics = this.elmDiagnostics.onCodeAction(params);
+
+    const results: CodeAction[] = [];
+
+    // For each diagnostic in the context, get the code action registration that
+    // handles the diagnostic error code and ask for the code actions for that error
+    // and the fix all code action for that error if there are other diagnostics with
+    // the same error code
+    params.context.diagnostics.forEach((diagnostic) => {
+      if (typeof diagnostic?.code === "string") {
+        const registrations = CodeActionProvider.errorCodeToRegistrationMap.getAll(
+          diagnostic.code,
+        );
+
+        // Set the params range to the diagnostic range so we get the correct nodes
+        params.range = diagnostic.range;
+
+        results.push(
+          ...flatMap(registrations, (reg) => {
+            const codeActions =
+              reg
+                .getCodeActions(params)
+                ?.map((codeAction) =>
+                  this.addDiagnosticToCodeAction(codeAction, diagnostic),
+                ) ?? [];
+
+            if (
+              codeActions.length > 0 &&
+              params.program
+                .getDiagnostics(params.sourceFile)
+                .some(
+                  (diag) =>
+                    !diagnosticsEquals(diag, diagnostic) &&
+                    diag.code === diagnostic.code,
+                )
+            ) {
+              const fixAllCodeAction = reg.getFixAllCodeAction(params);
+
+              if (fixAllCodeAction) {
+                codeActions?.push(fixAllCodeAction);
+              }
+            }
+
+            return codeActions;
+          }),
+        );
+      }
+    });
+
     return [
+      ...results,
       ...this.convertDiagnosticsToCodeActions(
         params.context.diagnostics,
-        elmWorkspace,
+        params.program,
         params.textDocument.uri,
       ),
-      ...this.getRefactorCodeActions(params, elmWorkspace),
-      ...this.getTypeAnnotationCodeActions(params, elmWorkspace),
+      ...this.getRefactorCodeActions(params),
+      ...this.getTypeAnnotationCodeActions(params),
       ...make,
       ...elmDiagnostics,
     ];
   }
 
   private getTypeAnnotationCodeActions(
-    params: CodeActionParams,
-    elmWorkspace: IElmWorkspace,
+    params: ICodeActionParams,
   ): CodeAction[] {
     // Top level annotation are handled by diagnostics
     const codeActions: CodeAction[] = [];
 
-    const forest = elmWorkspace.getForest();
+    const forest = params.program.getForest();
     const treeContainer = forest.getByUri(params.textDocument.uri);
     const tree = treeContainer?.tree;
-    const checker = elmWorkspace.getTypeChecker();
+    const checker = params.program.getTypeChecker();
 
     if (tree) {
       const nodeAtPosition = TreeUtils.getNamedDescendantForPosition(
@@ -119,13 +242,10 @@ export class CodeActionProvider {
     return codeActions;
   }
 
-  private getRefactorCodeActions(
-    params: CodeActionParams,
-    elmWorkspace: IElmWorkspace,
-  ): CodeAction[] {
+  private getRefactorCodeActions(params: ICodeActionParams): CodeAction[] {
     const codeActions: CodeAction[] = [];
 
-    const forest = elmWorkspace.getForest();
+    const forest = params.program.getForest();
     const tree = forest.getTree(params.textDocument.uri);
 
     if (tree) {
@@ -137,11 +257,7 @@ export class CodeActionProvider {
       codeActions.push(
         ...this.getFunctionCodeActions(params, tree, nodeAtPosition),
         ...this.getTypeAliasCodeActions(params, tree, nodeAtPosition),
-        ...this.getMakeDeclarationFromUsageCodeActions(
-          params,
-          elmWorkspace,
-          nodeAtPosition,
-        ),
+        ...this.getMakeDeclarationFromUsageCodeActions(params, nodeAtPosition),
       );
     }
 
@@ -168,7 +284,11 @@ export class CodeActionProvider {
           command: {
             title: "Refactor",
             command: "elm.refactor",
-            arguments: ["moveFunction", params, functionName],
+            arguments: [
+              "moveFunction",
+              { textDocument: params.textDocument, range: params.range },
+              functionName,
+            ],
           },
           kind: CodeActionKind.RefactorRewrite,
         });
@@ -265,8 +385,7 @@ export class CodeActionProvider {
   }
 
   private getMakeDeclarationFromUsageCodeActions(
-    params: CodeActionParams,
-    elmWorkspace: IElmWorkspace,
+    params: ICodeActionParams,
     nodeAtPosition: SyntaxNode,
   ): CodeAction[] {
     const codeActions: CodeAction[] = [];
@@ -279,14 +398,10 @@ export class CodeActionProvider {
     ) {
       const funcName = nodeAtPosition.text;
 
-      const treeContainer = elmWorkspace
-        .getForest()
-        .getByUri(params.textDocument.uri);
-      const tree = treeContainer?.tree;
-      const checker = elmWorkspace.getTypeChecker();
+      const tree = params.sourceFile.tree;
+      const checker = params.program.getTypeChecker();
 
       if (
-        tree &&
         !TreeUtils.findAllTopLevelFunctionDeclarations(tree)?.some(
           (a) =>
             a.firstChild?.text == funcName ||
@@ -299,7 +414,7 @@ export class CodeActionProvider {
 
         const typeString: string = checker.typeToString(
           checker.findType(nodeAtPosition),
-          treeContainer,
+          params.sourceFile,
         );
 
         const edit = RefactorEditUtils.createTopLevelFunction(
@@ -340,7 +455,7 @@ export class CodeActionProvider {
     if (treeContainer) {
       diagnostics.forEach((diagnostic) => {
         switch (diagnostic.code) {
-          case "missing_type_annotation":
+          case Diagnostics.MissingTypeAnnotation.code:
             {
               const nodeAtPosition = TreeUtils.getNamedDescendantForPosition(
                 treeContainer.tree.rootNode,
@@ -391,4 +506,12 @@ export class CodeActionProvider {
       title,
     };
   }
+
+  private addDiagnosticToCodeAction(
+    codeAction: CodeAction,
+    diagnostic: Diagnostic,
+  ): CodeAction {
+    codeAction.diagnostics = [diagnostic];
+    return codeAction;
+  }
 }
diff --git a/src/providers/diagnostics/diagnosticsProvider.ts b/src/providers/diagnostics/diagnosticsProvider.ts
index 6a41418..a7afb90 100644
--- a/src/providers/diagnostics/diagnosticsProvider.ts
+++ b/src/providers/diagnostics/diagnosticsProvider.ts
@@ -164,6 +164,23 @@ export class DiagnosticsProvider {
     });
   }
 
+  public interuptDiagnostics<T>(f: () => T): T {
+    if (!this.pendingRequest) {
+      return f();
+    }
+
+    this.pendingRequest.cancel();
+    this.pendingRequest = undefined;
+    const result = f();
+
+    this.triggerDiagnostics();
+    return result;
+  }
+
+  public getCurrentDiagnostics(uri: string): Diagnostic[] {
+    return this.currentDiagnostics.get(uri)?.get() ?? [];
+  }
+
   private requestDiagnostics(uri: string): void {
     this.pendingDiagnostics.set(uri, Date.now());
     this.triggerDiagnostics();
@@ -180,20 +197,6 @@ export class DiagnosticsProvider {
 
     this.triggerDiagnostics();
   }
-
-  public interuptDiagnostics<T>(f: () => T): T {
-    if (!this.pendingRequest) {
-      return f();
-    }
-
-    this.pendingRequest.cancel();
-    this.pendingRequest = undefined;
-    const result = f();
-
-    this.triggerDiagnostics();
-    return result;
-  }
-
   private triggerDiagnostics(delay = 200): void {
     const sendPendingDiagnostics = (): void => {
       const orderedFiles = this.pendingDiagnostics.getOrderedFiles();
@@ -298,7 +301,7 @@ export class DiagnosticsProvider {
             }
 
             const uri = files[index];
-            const workspace = this.elmWorkspaceMatcher.getElmWorkspaceFor(
+            const workspace = this.elmWorkspaceMatcher.getProgramFor(
               URI.parse(uri),
             );
 
@@ -320,6 +323,7 @@ export class DiagnosticsProvider {
                 return;
               }
 
+              this.updateDiagnostics(uri, DiagnosticKind.ElmMake, []);
               this.updateDiagnostics(
                 uri,
                 DiagnosticKind.TypeInference,
@@ -349,7 +353,7 @@ export class DiagnosticsProvider {
     );
   }
 
-  public async getElmMakeDiagnostics(uri: string): Promise<void> {
+  private async getElmMakeDiagnostics(uri: string): Promise<void> {
     const elmMakeDiagnostics = await this.elmMakeDiagnostics.createDiagnostics(
       URI.parse(uri),
     );
@@ -357,6 +361,7 @@ export class DiagnosticsProvider {
     this.resetDiagnostics(elmMakeDiagnostics, DiagnosticKind.ElmMake);
 
     elmMakeDiagnostics.forEach((diagnostics, diagnosticsUri) => {
+      this.updateDiagnostics(diagnosticsUri, DiagnosticKind.TypeInference, []);
       this.updateDiagnostics(
         diagnosticsUri,
         DiagnosticKind.ElmMake,
diff --git a/src/providers/diagnostics/elmDiagnosticsHelper.ts b/src/providers/diagnostics/elmDiagnosticsHelper.ts
index 9c904b4..a746937 100644
--- a/src/providers/diagnostics/elmDiagnosticsHelper.ts
+++ b/src/providers/diagnostics/elmDiagnosticsHelper.ts
@@ -1,7 +1,9 @@
 import path from "path";
 import { Diagnostic, DiagnosticSeverity, Range } from "vscode-languageserver";
 import { URI } from "vscode-uri";
+import { Diagnostics } from "../../util/types/diagnostics";
 import { IElmIssue } from "./diagnosticsProvider";
+import { NAMING_ERROR } from "./elmMakeDiagnostics";
 
 export class ElmDiagnosticsHelper {
   public static issuesToDiagnosticMap(
@@ -50,11 +52,17 @@ export class ElmDiagnosticsHelper {
 
     const messagePrefix = issue.overview ? `${issue.overview} - ` : "";
 
+    let code = undefined;
+
+    if (issue.overview.startsWith(NAMING_ERROR)) {
+      code = Diagnostics.MissingValue.code;
+    }
+
     return Diagnostic.create(
       lineRange,
       `${messagePrefix}${issue.details.replace(/\[\d+m/g, "")}`,
       this.severityStringToDiagnosticSeverity(issue.type),
-      undefined,
+      code,
       "Elm",
     );
   }
diff --git a/src/providers/diagnostics/elmLsDiagnostics.ts b/src/providers/diagnostics/elmLsDiagnostics.ts
index 735da86..a0e72b4 100644
--- a/src/providers/diagnostics/elmLsDiagnostics.ts
+++ b/src/providers/diagnostics/elmLsDiagnostics.ts
@@ -373,9 +373,7 @@ export class ElmLsDiagnostics {
   ): CodeAction[] {
     const result: CodeAction[] = [];
 
-    const elmWorkspace = this.elmWorkspaceMatcher.getElmWorkspaceFor(
-      URI.parse(uri),
-    );
+    const elmWorkspace = this.elmWorkspaceMatcher.getProgramFor(URI.parse(uri));
 
     const forest = elmWorkspace.getForest();
 
diff --git a/src/providers/diagnostics/elmMakeDiagnostics.ts b/src/providers/diagnostics/elmMakeDiagnostics.ts
index 4ed7e5c..67774d3 100644
--- a/src/providers/diagnostics/elmMakeDiagnostics.ts
+++ b/src/providers/diagnostics/elmMakeDiagnostics.ts
@@ -28,7 +28,7 @@ import execa = require("execa");
 import { IElmWorkspace } from "../../elmWorkspace";
 
 const ELM_MAKE = "Elm";
-const NAMING_ERROR = "NAMING ERROR";
+export const NAMING_ERROR = "NAMING ERROR";
 const RANDOM_ID = randomBytes(16).toString("hex");
 export const CODE_ACTION_ELM_MAKE = `elmLS.elmMakeFixer-${RANDOM_ID}`;
 const readFile = util.promisify(fs.readFile);
@@ -124,7 +124,7 @@ export class ElmMakeDiagnostics {
     filePath: URI,
   ): Promise<Map<string, Diagnostic[]>> => {
     const workspaceRootPath = this.elmWorkspaceMatcher
-      .getElmWorkspaceFor(filePath)
+      .getProgramFor(filePath)
       .getRootPath();
     const diagnostics = await this.checkForErrors(
       workspaceRootPath.fsPath,
@@ -136,9 +136,7 @@ export class ElmMakeDiagnostics {
     });
 
     // Handle import all
-    const forest = this.elmWorkspaceMatcher
-      .getElmWorkspaceFor(filePath)
-      .getForest();
+    const forest = this.elmWorkspaceMatcher.getProgramFor(filePath).getForest();
 
     const exposedValues = ImportUtils.getPossibleImports(
       forest,
@@ -215,95 +213,12 @@ export class ElmMakeDiagnostics {
   ): CodeAction[] {
     const result: CodeAction[] = [];
 
-    const elmWorkspace = this.elmWorkspaceMatcher.getElmWorkspaceFor(
-      URI.parse(uri),
-    );
-
+    const elmWorkspace = this.elmWorkspaceMatcher.getProgramFor(URI.parse(uri));
     const forest = elmWorkspace.getForest();
 
-    const exposedValues = ImportUtils.getPossibleImports(forest, uri);
-
     const sourceTree = forest.getByUri(uri);
 
     diagnostics.forEach((diagnostic) => {
-      if (diagnostic.message.startsWith(NAMING_ERROR)) {
-        const valueNode = sourceTree?.tree.rootNode.namedDescendantForPosition(
-          {
-            column: diagnostic.range.start.character,
-            row: diagnostic.range.start.line,
-          },
-          {
-            column: diagnostic.range.end.character,
-            row: diagnostic.range.end.line,
-          },
-        );
-
-        let hasImportFix = false;
-
-        // Add import quick fixes
-        if (valueNode) {
-          exposedValues
-            .filter(
-              (exposed) =>
-                exposed.value === valueNode.text ||
-                ((valueNode.type === "upper_case_qid" ||
-                  valueNode.type === "value_qid") &&
-                  exposed.value ===
-                    valueNode.namedChildren[valueNode.namedChildren.length - 1]
-                      .text &&
-                  exposed.module ===
-                    valueNode.namedChildren
-                      .slice(0, valueNode.namedChildren.length - 2) // Dots are also namedNodes
-                      .map((a) => a.text)
-                      .join("")),
-            )
-            .forEach((exposed) => {
-              hasImportFix = true;
-              result.push(
-                this.createImportQuickFix(
-                  uri,
-                  diagnostic,
-                  exposed.module,
-                  valueNode.type !== "upper_case_qid" &&
-                    valueNode.type !== "value_qid"
-                    ? exposed.valueToImport
-                      ? exposed.valueToImport
-                      : exposed.value
-                    : undefined,
-                ),
-              );
-            });
-        }
-
-        // Add import all quick fix
-        const filteredImports =
-          this.neededImports
-            .get(uri)
-            ?.filter(
-              (data, i, array) =>
-                array.findIndex(
-                  (d) =>
-                    data.moduleName === d.moduleName &&
-                    data.valueName === d.valueName,
-                ) === i,
-            ) ?? [];
-
-        if (hasImportFix && filteredImports.length > 1) {
-          // Sort so that the first diagnostic is this one
-          this.neededImports
-            .get(uri)
-            ?.sort((a, b) =>
-              a.diagnostic.message === diagnostic.message
-                ? -1
-                : b.diagnostic.message === diagnostic.message
-                ? 1
-                : 0,
-            );
-
-          result.push(this.createImportAllQuickFix(uri));
-        }
-      }
-
       if (
         diagnostic.message.startsWith(NAMING_ERROR) ||
         diagnostic.message.startsWith("BAD IMPORT") ||
@@ -531,7 +446,7 @@ export class ElmMakeDiagnostics {
     }
 
     const tree = this.elmWorkspaceMatcher
-      .getElmWorkspaceFor(URI.parse(uri))
+      .getProgramFor(URI.parse(uri))
       .getForest()
       .getTree(uri);
 
@@ -554,37 +469,6 @@ export class ElmMakeDiagnostics {
     };
   }
 
-  private createImportAllQuickFix(uri: string): CodeAction {
-    const changes: {
-      [uri: string]: TextEdit[];
-    } = {};
-    if (!changes[uri]) {
-      changes[uri] = [];
-    }
-
-    const tree = this.elmWorkspaceMatcher
-      .getElmWorkspaceFor(URI.parse(uri))
-      .getForest()
-      .getTree(uri);
-
-    const imports = this.neededImports.get(uri);
-
-    if (tree && imports) {
-      const edit = RefactorEditUtils.addImports(tree, imports);
-
-      if (edit) {
-        changes[uri].push(edit);
-      }
-    }
-
-    return {
-      diagnostics: imports?.map((data) => data.diagnostic),
-      edit: { changes },
-      kind: CodeActionKind.QuickFix,
-      title: `Add all missing imports`,
-    };
-  }
-
   private filterElmMakeDiagnostics(diagnostics: Diagnostic[]): Diagnostic[] {
     return diagnostics.filter((diagnostic) => diagnostic.source === ELM_MAKE);
   }
diff --git a/src/providers/diagnostics/typeInferenceDiagnostics.ts b/src/providers/diagnostics/typeInferenceDiagnostics.ts
index 149b869..36cede4 100644
--- a/src/providers/diagnostics/typeInferenceDiagnostics.ts
+++ b/src/providers/diagnostics/typeInferenceDiagnostics.ts
@@ -14,10 +14,10 @@ export class TypeInferenceDiagnostics {
 
   public getDiagnosticsForFileAsync(
     treeContainer: ITreeContainer,
-    elmWorkspace: IElmWorkspace,
+    program: IElmWorkspace,
     cancellationToken: CancellationToken,
   ): Promise<Diagnostic[]> {
-    const checker = elmWorkspace.getTypeChecker();
+    const checker = program.getTypeChecker();
 
     const diagnostics: Diagnostic[] = [];
 
@@ -27,7 +27,7 @@ export class TypeInferenceDiagnostics {
       TreeUtils.findAllTopLevelFunctionDeclarations(treeContainer.tree) ?? [];
 
     return new Promise((resolve, reject) => {
-      checker
+      program
         .getDiagnosticsAsync(
           treeContainer,
           new ServerCancellationToken(cancellationToken),
diff --git a/src/providers/index.ts b/src/providers/index.ts
index 765c8ca..d333f7f 100644
--- a/src/providers/index.ts
+++ b/src/providers/index.ts
@@ -14,3 +14,4 @@ export * from "./hoverProvider";
 export * from "./referencesProvider";
 export * from "./renameProvider";
 export * from "./workspaceSymbolProvider";
+import "./codeAction";
diff --git a/src/util/elmWorkspaceMatcher.ts b/src/util/elmWorkspaceMatcher.ts
index a4ff9a6..1b018d3 100644
--- a/src/util/elmWorkspaceMatcher.ts
+++ b/src/util/elmWorkspaceMatcher.ts
@@ -2,12 +2,19 @@ import { container } from "tsyringe";
 import { CancellationToken } from "vscode-languageserver";
 import { URI } from "vscode-uri";
 import { IElmWorkspace } from "../elmWorkspace";
+import { ITreeContainer } from "../forest";
 import { NoWorkspaceContainsError } from "./noWorkspaceContainsError";
 
+export interface IParams {
+  program: IElmWorkspace;
+  sourceFile: ITreeContainer;
+}
+
 /**
  * Identifies the relevant ElmWorkspace for a given ParamType, either directly
  * (getElmWorkspaceFor) or when an event handler receives a ParamType
- * (handlerForWorkspace).
+ * (handle) it returns a params object with a combined type ParamType and IParams
+ * which has the program and sourceFile.
  */
 export class ElmWorkspaceMatcher<ParamType> {
   private elmWorkspaces: IElmWorkspace[];
@@ -16,19 +23,41 @@ export class ElmWorkspaceMatcher<ParamType> {
     this.elmWorkspaces = container.resolve("ElmWorkspaces");
   }
 
+  public handle<ResultType>(
+    handler: (
+      param: ParamType & IParams,
+      token?: CancellationToken,
+    ) => ResultType,
+  ): (param: ParamType, token?: CancellationToken) => ResultType {
+    return (param: ParamType, token?: CancellationToken): ResultType => {
+      const program = this.getProgramFor(param);
+      return handler(
+        {
+          ...param,
+          program,
+          sourceFile: this.getSourceFileFor(param, program),
+        },
+        token,
+      );
+    };
+  }
+
+  /**
+   * @deprecated Use handle() instead, which returns a params with the program and source file in it
+   */
   public handlerForWorkspace<ResultType>(
     handler: (
       param: ParamType,
-      elmWorkspace: IElmWorkspace,
+      program: IElmWorkspace,
       token?: CancellationToken,
     ) => ResultType,
   ): (param: ParamType, token?: CancellationToken) => ResultType {
     return (param: ParamType, token?: CancellationToken): ResultType => {
-      return handler(param, this.getElmWorkspaceFor(param), token);
+      return handler(param, this.getProgramFor(param), token);
     };
   }
 
-  public getElmWorkspaceFor(param: ParamType): IElmWorkspace {
+  public getProgramFor(param: ParamType): IElmWorkspace {
     const uri = this.getUriFor(param);
     const workspace =
       // first look for a workspace where the file has been parsed to a tree
@@ -42,4 +71,19 @@ export class ElmWorkspaceMatcher<ParamType> {
 
     return workspace;
   }
+
+  public getSourceFileFor(
+    param: ParamType,
+    program: IElmWorkspace,
+  ): ITreeContainer {
+    const uri = this.getUriFor(param).toString();
+
+    const sourceFile = program.getForest().getByUri(uri);
+
+    if (!sourceFile) {
+      throw new NoWorkspaceContainsError(this.getUriFor(param));
+    }
+
+    return sourceFile;
+  }
 }
diff --git a/src/util/multiMap.ts b/src/util/multiMap.ts
index bb5a501..1cda19e 100644
--- a/src/util/multiMap.ts
+++ b/src/util/multiMap.ts
@@ -22,6 +22,16 @@ export class MultiMap<K, V> extends Map<K, V | V[]> {
     }
   }
 
+  public getAll(key: K): V[] | undefined {
+    const found = super.get(key);
+
+    if (Array.isArray(found)) {
+      return found;
+    } else if (found) {
+      return [found];
+    }
+  }
+
   public set(key: K, val: V): this {
     if (super.has(key)) {
       const existing = super.get(key);
diff --git a/src/util/treeUtils.ts b/src/util/treeUtils.ts
index 9451599..641055e 100644
--- a/src/util/treeUtils.ts
+++ b/src/util/treeUtils.ts
@@ -8,6 +8,7 @@ import {
   EFunctionCallExpr,
   mapSyntaxNodeToExpression,
 } from "./types/expressionTree";
+import { Range } from "vscode-languageserver-textdocument";
 
 export type NodeType =
   | "Function"
@@ -39,14 +40,16 @@ export interface IExposed {
 export type IExposing = Map<string, IExposed>;
 
 export function flatMap<T, U>(
-  array: T[],
-  callback: (value: T, index: number, array: T[]) => U[],
+  array: T[] | undefined,
+  callback: (value: T, index: number, array: T[]) => U[] | undefined,
 ): U[] {
   const flattend: U[] = [];
-  for (let i = 0; i < array.length; i++) {
-    const elementArray = callback(array[i], i, array);
-    for (const el of elementArray) {
-      flattend.push(el);
+  if (array) {
+    for (let i = 0; i < array.length; i++) {
+      const elementArray = callback(array[i], i, array);
+      if (elementArray) {
+        flattend.push(...elementArray);
+      }
     }
   }
   return flattend;
@@ -806,6 +809,22 @@ export class TreeUtils {
     }
   }
 
+  public static getNamedDescendantForRange(
+    sourceFile: ITreeContainer,
+    range: Range,
+  ): SyntaxNode {
+    return sourceFile.tree.rootNode.namedDescendantForPosition(
+      {
+        column: range.start.character,
+        row: range.start.line,
+      },
+      {
+        column: range.end.character,
+        row: range.end.line,
+      },
+    );
+  }
+
   public static findPreviousNode(
     node: SyntaxNode,
     position: Position,
diff --git a/src/util/types/typeInference.ts b/src/util/types/typeInference.ts
index b0b5424..aaf6518 100644
--- a/src/util/types/typeInference.ts
+++ b/src/util/types/typeInference.ts
@@ -449,7 +449,7 @@ function argumentCountError(
       node,
       Diagnostics.ArgumentCount,
       endNode,
-      name ? `\`${name}\`` : "This value",
+      name ?? "This value",
       actual,
     );
   } else {
diff --git a/src/util/utils.ts b/src/util/utils.ts
index 2c1b966..c038c72 100644
--- a/src/util/utils.ts
+++ b/src/util/utils.ts
@@ -31,4 +31,23 @@ export class Utils {
       a.end.line === b.end.line
     );
   }
+
+  public static rangeOverlaps(a: Range, b: Range): boolean {
+    if (b.start.line < a.start.line || b.end.line < a.start.line) {
+      return false;
+    }
+    if (b.start.line > a.end.line || b.end.line > a.end.line) {
+      return false;
+    }
+    if (
+      b.start.line === a.start.line &&
+      b.start.character < a.start.character
+    ) {
+      return false;
+    }
+    if (b.end.line === a.end.line && b.end.character > a.end.character) {
+      return false;
+    }
+    return true;
+  }
 }


Test patch:
diff --git a/test/codeActionProvider.test.ts b/test/codeActionProvider.test.ts
new file mode 100644
index 0000000..1572c0d
--- /dev/null
+++ b/test/codeActionProvider.test.ts
@@ -0,0 +1,183 @@
+import { container } from "tsyringe";
+import { CodeAction } from "vscode-languageserver";
+import { URI } from "vscode-uri";
+import { IElmWorkspace } from "../src/elmWorkspace";
+import { CodeActionProvider, ICodeActionParams } from "../src/providers";
+import { Utils } from "../src/util/utils";
+import { baseUri } from "./utils/mockElmWorkspace";
+import { getTargetPositionFromSource } from "./utils/sourceParser";
+import { SourceTreeParser } from "./utils/sourceTreeParser";
+
+function codeActionEquals(a: CodeAction, b: CodeAction): boolean {
+  return a.title === b.title;
+}
+
+const basicsSources = `
+--@ Basics.elm
+module Basics exposing ((+), (|>), (==), Int, Float, Bool(..), Order(..))
+
+infix left  0 (|>) = apR
+infix non   4 (==) = eq
+infix left  6 (+)  = add
+
+type Int = Int
+
+type Float = Float
+
+type Bool = True | False
+
+add : number -> number -> number
+add =
+  Elm.Kernel.Basics.add
+
+apR : a -> (a -> b) -> b
+apR x f =
+  f x
+
+eq : a -> a -> Bool
+eq =
+  Elm.Kernel.Utils.equal
+
+type Order = LT | EQ | GT
+`;
+
+class MockCodeActionsProvider extends CodeActionProvider {
+  public handleCodeAction(params: ICodeActionParams): CodeAction[] | undefined {
+    return this.onCodeAction(params);
+  }
+}
+
+describe("test codeActionProvider", () => {
+  const treeParser = new SourceTreeParser();
+  let codeActionProvider: MockCodeActionsProvider;
+
+  const debug = process.argv.find((arg) => arg === "--debug");
+
+  async function testCodeAction(
+    source: string,
+    expectedCodeActions: CodeAction[],
+  ) {
+    await treeParser.init();
+
+    if (!codeActionProvider) {
+      codeActionProvider = new MockCodeActionsProvider();
+    }
+
+    const result = getTargetPositionFromSource(source);
+
+    if (!result) {
+      throw new Error("Getting sources failed");
+    }
+
+    const testUri = URI.file(baseUri + "Test.elm").toString();
+
+    const program = treeParser.getWorkspace(result.sources);
+    const sourceFile = program.getForest().getByUri(testUri);
+
+    if (!sourceFile) throw new Error("Getting tree failed");
+
+    const workspaces = container.resolve<IElmWorkspace[]>("ElmWorkspaces");
+    workspaces.splice(0, workspaces.length);
+    workspaces.push(program);
+
+    const range = { start: result.position, end: result.position };
+    const codeActions =
+      codeActionProvider.handleCodeAction({
+        program,
+        sourceFile,
+        range,
+        textDocument: { uri: testUri },
+        context: {
+          diagnostics: program
+            .getDiagnostics(sourceFile)
+            .filter((diag) => Utils.rangeOverlaps(diag.range, range)),
+        },
+      }) ?? [];
+
+    const codeActionsExist = expectedCodeActions.every((codeAction) =>
+      codeActions.find((c) => codeActionEquals(codeAction, c)),
+    );
+
+    if (debug && !codeActionsExist) {
+      console.log(
+        `Expecting ${JSON.stringify(expectedCodeActions)}, got ${JSON.stringify(
+          codeActions,
+        )}`,
+      );
+    }
+
+    expect(codeActionsExist).toBeTruthy();
+  }
+
+  test("add import of value", async () => {
+    const source = `
+--@ Test.elm
+module Test exposing (..)
+
+func = foo
+      --^
+
+--@ App.elm
+module App exposing (foo)
+
+foo = ""
+`;
+    await testCodeAction(basicsSources + source, [
+      { title: `Import 'foo' from module "App"` },
+    ]);
+  });
+
+  test("add import of qualified value", async () => {
+    const source = `
+--@ Test.elm
+module Test exposing (..)
+
+func = App.foo
+          --^
+
+--@ App.elm
+module App exposing (foo)
+
+foo = ""
+`;
+    await testCodeAction(basicsSources + source, [
+      { title: `Import module "App"` },
+    ]);
+
+    const source2 = `
+--@ Test.elm
+module Test exposing (..)
+
+func = App.foo
+      --^
+
+--@ App.elm
+module App exposing (foo)
+
+foo = ""
+`;
+    await testCodeAction(basicsSources + source2, [
+      { title: `Import module "App"` },
+    ]);
+  });
+
+  test("add all missing imports", async () => {
+    const source = `
+--@ Test.elm
+module Test exposing (..)
+
+func = foo + bar
+      --^
+
+--@ App.elm
+module App exposing (foo, bar)
+
+foo = ""
+
+bar = ""
+`;
+    await testCodeAction(basicsSources + source, [
+      { title: `Add all missing imports` },
+    ]);
+  });
+});
diff --git a/test/diagnosticTests/elmDiagnostics.test.ts b/test/diagnosticTests/elmDiagnostics.test.ts
index aff2d5d..8d2cd2d 100644
--- a/test/diagnosticTests/elmDiagnostics.test.ts
+++ b/test/diagnosticTests/elmDiagnostics.test.ts
@@ -69,25 +69,23 @@ describe("test elm diagnostics", () => {
 
     const testUri = URI.file(baseUri + "Test.elm").toString();
 
-    const workspace = treeParser.getWorkspace(result.sources);
-    const treeContainer = workspace.getForest().getByUri(testUri);
+    const program = treeParser.getWorkspace(result.sources);
+    const treeContainer = program.getForest().getByUri(testUri);
 
     if (!treeContainer) throw new Error("Getting tree failed");
 
-    const checker = workspace.getTypeChecker();
     const diagnostics: Diagnostic[] = [];
 
-    workspace.getForest().treeMap.forEach((treeContainer) => {
+    program.getForest().treeMap.forEach((treeContainer) => {
       if (!treeContainer.uri.includes("Basic")) {
-        diagnostics.push(...checker.getDiagnostics(treeContainer));
+        diagnostics.push(...program.getDiagnostics(treeContainer));
       }
     });
 
     let nodeAtPosition: SyntaxNode;
 
     if ("position" in result) {
-      const rootNode = workspace.getForest().treeMap.get(testUri)!.tree
-        .rootNode;
+      const rootNode = program.getForest().treeMap.get(testUri)!.tree.rootNode;
       nodeAtPosition = TreeUtils.getNamedDescendantForPosition(
         rootNode,
         result.position,
diff --git a/test/utils/mockElmWorkspace.ts b/test/utils/mockElmWorkspace.ts
index 1276e28..09b2d71 100644
--- a/test/utils/mockElmWorkspace.ts
+++ b/test/utils/mockElmWorkspace.ts
@@ -2,7 +2,7 @@ import path, * as Path from "path";
 import { URI } from "vscode-uri";
 import Parser, { Tree } from "web-tree-sitter";
 import { IElmWorkspace } from "../../src/elmWorkspace";
-import { Forest, IForest } from "../../src/forest";
+import { Forest, IForest, ITreeContainer } from "../../src/forest";
 import { Imports } from "../../src/imports";
 import { container } from "tsyringe";
 import { TypeCache } from "../../src/util/types/typeCache";
@@ -16,6 +16,8 @@ import {
   IPossibleImportsCache,
   PossibleImportsCache,
 } from "../../src/util/possibleImportsCache";
+import { Diagnostic } from "vscode-languageserver";
+import { ICancellationToken } from "../../src/cancellation";
 
 export const baseUri = Path.join(__dirname, "../sources/src/");
 
@@ -25,6 +27,7 @@ export class MockElmWorkspace implements IElmWorkspace {
   private typeCache = new TypeCache();
   private possibleImportsCache = new PossibleImportsCache();
   private operatorsCache = new Map<string, DefinitionResult>();
+  private diagnosticsCache = new Map<string, Diagnostic[]>();
 
   constructor(sources: { [K: string]: string }) {
     this.parser = container.resolve("Parser");
@@ -47,19 +50,16 @@ export class MockElmWorkspace implements IElmWorkspace {
     return;
   }
 
-  // eslint-disable-next-line @typescript-eslint/no-unused-vars
-  hasDocument(uri: URI): boolean {
-    return false;
+  public hasDocument(uri: URI): boolean {
+    return !!this.forest.getTree(uri.toString());
   }
 
-  // eslint-disable-next-line @typescript-eslint/no-unused-vars
-  hasPath(uri: URI): boolean {
-    return false;
+  public hasPath(uri: URI): boolean {
+    return !!this.getPath(uri);
   }
 
-  // eslint-disable-next-line @typescript-eslint/no-unused-vars
-  getPath(uri: URI): string | undefined {
-    return;
+  public getPath(uri: URI): string | undefined {
+    return baseUri;
   }
 
   getForest(): IForest {
@@ -90,6 +90,44 @@ export class MockElmWorkspace implements IElmWorkspace {
     return this.operatorsCache;
   }
 
+  public getDiagnostics(
+    sourceFile: ITreeContainer,
+    cancellationToken?: ICancellationToken,
+  ): Diagnostic[] {
+    const cached = this.diagnosticsCache.get(sourceFile.uri);
+
+    if (cached) {
+      return cached;
+    }
+
+    const diagnostics = this.getTypeChecker().getDiagnostics(
+      sourceFile,
+      cancellationToken,
+    );
+
+    this.diagnosticsCache.set(sourceFile.uri, diagnostics);
+    return diagnostics;
+  }
+
+  public async getDiagnosticsAsync(
+    sourceFile: ITreeContainer,
+    cancellationToken?: ICancellationToken,
+  ): Promise<Diagnostic[]> {
+    const cached = this.diagnosticsCache.get(sourceFile.uri);
+
+    if (cached) {
+      return Promise.resolve(cached);
+    }
+
+    const diagnostics = await this.getTypeChecker().getDiagnosticsAsync(
+      sourceFile,
+      cancellationToken,
+    );
+
+    this.diagnosticsCache.set(sourceFile.uri, diagnostics);
+    return diagnostics;
+  }
+
   private parseAndAddToForest(fileName: string, source: string): void {
     const tree: Tree | undefined = this.parser.parse(source);
     this.forest.setTree(