5 files changed, 185 insertions, 101 deletions

diff --git a/editors/code/src/utils/diagnostics/SuggestedFix.ts b/editors/code/src/utils/diagnostics/SuggestedFix.ts
new file mode 100644
index 000000000..b1be2a225
--- /dev/null
+++ b/editors/code/src/utils/diagnostics/SuggestedFix.ts
@@ -0,0 +1,67 @@
+import * as vscode from 'vscode';
+
+import { SuggestionApplicability } from './rust';
+
+/**
+ * Model object for text replacements suggested by the Rust compiler
+ *
+ * This is an intermediate form between the raw `rustc` JSON and a
+ * `vscode.CodeAction`. It's optimised for the use-cases of
+ * `SuggestedFixCollection`.
+ */
+export default class SuggestedFix {
+    public readonly title: string;
+    public readonly location: vscode.Location;
+    public readonly replacement: string;
+    public readonly applicability: SuggestionApplicability;
+
+    /**
+     * Diagnostics this suggested fix could resolve
+     */
+    public diagnostics: vscode.Diagnostic[];
+
+    constructor(
+        title: string,
+        location: vscode.Location,
+        replacement: string,
+        applicability: SuggestionApplicability = SuggestionApplicability.Unspecified
+    ) {
+        this.title = title;
+        this.location = location;
+        this.replacement = replacement;
+        this.applicability = applicability;
+        this.diagnostics = [];
+    }
+
+    /**
+     * Determines if this suggested fix is equivalent to another instance
+     */
+    public isEqual(other: SuggestedFix): boolean {
+        return (
+            this.title === other.title &&
+            this.location.range.isEqual(other.location.range) &&
+            this.replacement === other.replacement &&
+            this.applicability === other.applicability
+        );
+    }
+
+    /**
+     * Converts this suggested fix to a VS Code Quick Fix code action
+     */
+    public toCodeAction(): vscode.CodeAction {
+        const codeAction = new vscode.CodeAction(
+            this.title,
+            vscode.CodeActionKind.QuickFix
+        );
+
+        const edit = new vscode.WorkspaceEdit();
+        edit.replace(this.location.uri, this.location.range, this.replacement);
+        codeAction.edit = edit;
+
+        codeAction.isPreferred =
+            this.applicability === SuggestionApplicability.MachineApplicable;
+
+        codeAction.diagnostics = [...this.diagnostics];
+        return codeAction;
+    }
+}
diff --git a/editors/code/src/utils/diagnostics/SuggestedFixCollection.ts b/editors/code/src/utils/diagnostics/SuggestedFixCollection.ts
new file mode 100644
index 000000000..132ce12f8
--- /dev/null
+++ b/editors/code/src/utils/diagnostics/SuggestedFixCollection.ts
@@ -0,0 +1,77 @@
+import * as vscode from 'vscode';
+import SuggestedFix from './SuggestedFix';
+
+/**
+ * Collection of suggested fixes across multiple documents
+ *
+ * This stores `SuggestedFix` model objects and returns them via the
+ * `vscode.CodeActionProvider` interface.
+ */
+export default class SuggestedFixCollection
+    implements vscode.CodeActionProvider {
+    public static PROVIDED_CODE_ACTION_KINDS = [vscode.CodeActionKind.QuickFix];
+
+    /**
+     * Map of document URI strings to suggested fixes
+     */
+    private suggestedFixes: Map<string, SuggestedFix[]>;
+
+    constructor() {
+        this.suggestedFixes = new Map();
+    }
+
+    /**
+     * Clears all suggested fixes across all documents
+     */
+    public clear(): void {
+        this.suggestedFixes = new Map();
+    }
+
+    /**
+     * Adds a suggested fix for the given diagnostic
+     *
+     * Some suggested fixes will appear in multiple diagnostics. For example,
+     * forgetting a `mut` on a variable will suggest changing the delaration on
+     * every mutable usage site. If the suggested fix has already been added
+     * this method will instead associate the existing fix with the new
+     * diagnostic.
+     */
+    public addSuggestedFixForDiagnostic(
+        suggestedFix: SuggestedFix,
+        diagnostic: vscode.Diagnostic
+    ): void {
+        const fileUriString = suggestedFix.location.uri.toString();
+        const fileSuggestions = this.suggestedFixes.get(fileUriString) || [];
+
+        const existingSuggestion = fileSuggestions.find(s =>
+            s.isEqual(suggestedFix)
+        );
+
+        if (existingSuggestion) {
+            // The existing suggestion also applies to this new diagnostic
+            existingSuggestion.diagnostics.push(diagnostic);
+        } else {
+            // We haven't seen this suggestion before
+            suggestedFix.diagnostics.push(diagnostic);
+            fileSuggestions.push(suggestedFix);
+        }
+
+        this.suggestedFixes.set(fileUriString, fileSuggestions);
+    }
+
+    /**
+     * Filters suggested fixes by their document and range and converts them to
+     * code actions
+     */
+    public provideCodeActions(
+        document: vscode.TextDocument,
+        range: vscode.Range
+    ): vscode.CodeAction[] {
+        const documentUriString = document.uri.toString();
+
+        const suggestedFixes = this.suggestedFixes.get(documentUriString);
+        return (suggestedFixes || [])
+            .filter(({ location }) => location.range.intersection(range))
+            .map(suggestedEdit => suggestedEdit.toCodeAction());
+    }
+}
diff --git a/editors/code/src/utils/rust_diagnostics.ts b/editors/code/src/utils/diagnostics/rust.ts
index 3c524cb37..d16576eb1 100644
--- a/editors/code/src/utils/rust_diagnostics.ts
+++ b/editors/code/src/utils/diagnostics/rust.ts
@@ -1,6 +1,15 @@
 import * as path from 'path';
 import * as vscode from 'vscode';
 
+import SuggestedFix from './SuggestedFix';
+
+export enum SuggestionApplicability {
+    MachineApplicable = 'MachineApplicable',
+    HasPlaceholders = 'HasPlaceholders',
+    MaybeIncorrect = 'MaybeIncorrect',
+    Unspecified = 'Unspecified'
+}
+
 // Reference:
 // https://github.com/rust-lang/rust/blob/master/src/libsyntax/json.rs
 export interface RustDiagnosticSpan {
@@ -12,11 +21,7 @@ export interface RustDiagnosticSpan {
     file_name: string;
     label?: string;
     suggested_replacement?: string;
-    suggestion_applicability?:
-        | 'MachineApplicable'
-        | 'HasPlaceholders'
-        | 'MaybeIncorrect'
-        | 'Unspecified';
+    suggestion_applicability?: SuggestionApplicability;
 }
 
 export interface RustDiagnostic {
@@ -33,12 +38,12 @@ export interface RustDiagnostic {
 export interface MappedRustDiagnostic {
     location: vscode.Location;
     diagnostic: vscode.Diagnostic;
-    codeActions: vscode.CodeAction[];
+    suggestedFixes: SuggestedFix[];
 }
 
 interface MappedRustChildDiagnostic {
     related?: vscode.DiagnosticRelatedInformation;
-    codeAction?: vscode.CodeAction;
+    suggestedFix?: SuggestedFix;
     messageLine?: string;
 }
 
@@ -130,24 +135,19 @@ function mapRustChildDiagnostic(rd: RustDiagnostic): MappedRustChildDiagnostic {
 
     // We need to distinguish `null` from an empty string
     if (span && typeof span.suggested_replacement === 'string') {
-        const edit = new vscode.WorkspaceEdit();
-        edit.replace(location.uri, location.range, span.suggested_replacement);
-
-        // Include our replacement in the label unless it's empty
+        // Include our replacement in the title unless it's empty
         const title = span.suggested_replacement
             ? `${rd.message}: \`${span.suggested_replacement}\``
             : rd.message;
 
-        const codeAction = new vscode.CodeAction(
-            title,
-            vscode.CodeActionKind.QuickFix
-        );
-
-        codeAction.edit = edit;
-        codeAction.isPreferred =
-            span.suggestion_applicability === 'MachineApplicable';
-
-        return { codeAction };
+        return {
+            suggestedFix: new SuggestedFix(
+                title,
+                location,
+                span.suggested_replacement,
+                span.suggestion_applicability
+            )
+        };
     } else {
         const related = new vscode.DiagnosticRelatedInformation(
             location,
@@ -165,7 +165,7 @@ function mapRustChildDiagnostic(rd: RustDiagnostic): MappedRustChildDiagnostic {
  *
  * 1. Creating a `vscode.Diagnostic` with the root message and primary span.
  * 2. Adding any labelled secondary spans to `relatedInformation`
- * 3. Categorising child diagnostics as either Quick Fix actions,
+ * 3. Categorising child diagnostics as either `SuggestedFix`es,
  *    `relatedInformation` or additional message lines.
  *
  * If the diagnostic has no primary span this will return `undefined`
@@ -173,8 +173,6 @@ function mapRustChildDiagnostic(rd: RustDiagnostic): MappedRustChildDiagnostic {
 export function mapRustDiagnosticToVsCode(
     rd: RustDiagnostic
 ): MappedRustDiagnostic | undefined {
-    const codeActions = [];
-
     const primarySpan = rd.spans.find(s => s.is_primary);
     if (!primarySpan) {
         return;
@@ -208,16 +206,17 @@ export function mapRustDiagnosticToVsCode(
         }
     }
 
+    const suggestedFixes = [];
     for (const child of rd.children) {
-        const { related, codeAction, messageLine } = mapRustChildDiagnostic(
+        const { related, suggestedFix, messageLine } = mapRustChildDiagnostic(
             child
         );
 
         if (related) {
             vd.relatedInformation.push(related);
         }
-        if (codeAction) {
-            codeActions.push(codeAction);
+        if (suggestedFix) {
+            suggestedFixes.push(suggestedFix);
         }
         if (messageLine) {
             vd.message += `\n${messageLine}`;
@@ -231,6 +230,6 @@ export function mapRustDiagnosticToVsCode(
     return {
         location,
         diagnostic: vd,
-        codeActions
+        suggestedFixes
     };
 }
diff --git a/editors/code/src/utils/diagnostics/vscode.ts b/editors/code/src/utils/diagnostics/vscode.ts
new file mode 100644
index 000000000..d8b85b720
--- /dev/null
+++ b/editors/code/src/utils/diagnostics/vscode.ts
@@ -0,0 +1,14 @@
+import * as vscode from 'vscode';
+
+/** Compares two `vscode.Diagnostic`s for equality */
+export function areDiagnosticsEqual(
+    left: vscode.Diagnostic,
+    right: vscode.Diagnostic
+): boolean {
+    return (
+        left.source === right.source &&
+        left.severity === right.severity &&
+        left.range.isEqual(right.range) &&
+        left.message === right.message
+    );
+}
diff --git a/editors/code/src/utils/vscode_diagnostics.ts b/editors/code/src/utils/vscode_diagnostics.ts
deleted file mode 100644
index 9d763c8d6..000000000
--- a/editors/code/src/utils/vscode_diagnostics.ts
+++ /dev/null
@@ -1,73 +0,0 @@
-import * as vscode from 'vscode';
-
-/** Compares two `vscode.Diagnostic`s for equality */
-export function areDiagnosticsEqual(
-    left: vscode.Diagnostic,
-    right: vscode.Diagnostic
-): boolean {
-    return (
-        left.source === right.source &&
-        left.severity === right.severity &&
-        left.range.isEqual(right.range) &&
-        left.message === right.message
-    );
-}
-
-/** Compares two `vscode.TextEdit`s for equality */
-function areTextEditsEqual(
-    left: vscode.TextEdit,
-    right: vscode.TextEdit
-): boolean {
-    if (!left.range.isEqual(right.range)) {
-        return false;
-    }
-
-    if (left.newText !== right.newText) {
-        return false;
-    }
-
-    return true;
-}
-
-/** Compares two `vscode.CodeAction`s for equality */
-export function areCodeActionsEqual(
-    left: vscode.CodeAction,
-    right: vscode.CodeAction
-): boolean {
-    if (
-        left.kind !== right.kind ||
-        left.title !== right.title ||
-        !left.edit ||
-        !right.edit
-    ) {
-        return false;
-    }
-
-    const leftEditEntries = left.edit.entries();
-    const rightEditEntries = right.edit.entries();
-
-    if (leftEditEntries.length !== rightEditEntries.length) {
-        return false;
-    }
-
-    for (let i = 0; i < leftEditEntries.length; i++) {
-        const [leftUri, leftEdits] = leftEditEntries[i];
-        const [rightUri, rightEdits] = rightEditEntries[i];
-
-        if (leftUri.toString() !== rightUri.toString()) {
-            return false;
-        }
-
-        if (leftEdits.length !== rightEdits.length) {
-            return false;
-        }
-
-        for (let j = 0; j < leftEdits.length; j++) {
-            if (!areTextEditsEqual(leftEdits[j], rightEdits[j])) {
-                return false;
-            }
-        }
-    }
-
-    return true;
-}