2 files changed, 352 insertions, 54 deletions

diff --git a/editors/code/src/commands/cargo_watch.ts b/editors/code/src/commands/cargo_watch.ts
index 13adf4c10..a662f7cc8 100644
--- a/editors/code/src/commands/cargo_watch.ts
+++ b/editors/code/src/commands/cargo_watch.ts
@@ -4,6 +4,10 @@ import * as path from 'path';
 import * as vscode from 'vscode';
 import { Server } from '../server';
 import { terminate } from '../utils/processes';
+import {
+    mapRustDiagnosticToVsCode,
+    RustDiagnostic
+} from '../utils/rust_diagnostics';
 import { LineBuffer } from './line_buffer';
 import { StatusDisplay } from './watch_status';
 
@@ -33,10 +37,17 @@ export function registerCargoWatchProvider(
     return provider;
 }
 
-export class CargoWatchProvider implements vscode.Disposable {
+export class CargoWatchProvider
+    implements vscode.Disposable, vscode.CodeActionProvider {
     private readonly diagnosticCollection: vscode.DiagnosticCollection;
     private readonly statusDisplay: StatusDisplay;
     private readonly outputChannel: vscode.OutputChannel;
+
+    private codeActions: {
+        [fileUri: string]: vscode.CodeAction[];
+    };
+    private readonly codeActionDispose: vscode.Disposable;
+
     private cargoProcess?: child_process.ChildProcess;
 
     constructor() {
@@ -49,6 +60,16 @@ export class CargoWatchProvider implements vscode.Disposable {
         this.outputChannel = vscode.window.createOutputChannel(
             'Cargo Watch Trace'
         );
+
+        // Register code actions for rustc's suggested fixes
+        this.codeActions = {};
+        this.codeActionDispose = vscode.languages.registerCodeActionsProvider(
+            [{ scheme: 'file', language: 'rust' }],
+            this,
+            {
+                providedCodeActionKinds: [vscode.CodeActionKind.QuickFix]
+            }
+        );
     }
 
     public start() {
@@ -127,6 +148,14 @@ export class CargoWatchProvider implements vscode.Disposable {
         this.diagnosticCollection.dispose();
         this.outputChannel.dispose();
         this.statusDisplay.dispose();
+        this.codeActionDispose.dispose();
+    }
+
+    public provideCodeActions(
+        document: vscode.TextDocument
+    ): vscode.ProviderResult<Array<vscode.Command | vscode.CodeAction>> {
+        const documentActions = this.codeActions[document.uri.toString()];
+        return documentActions || [];
     }
 
     private logInfo(line: string) {
@@ -147,6 +176,7 @@ export class CargoWatchProvider implements vscode.Disposable {
     private parseLine(line: string) {
         if (line.startsWith('[Running')) {
             this.diagnosticCollection.clear();
+            this.codeActions = {};
             this.statusDisplay.show();
         }
 
@@ -154,34 +184,65 @@ export class CargoWatchProvider implements vscode.Disposable {
             this.statusDisplay.hide();
         }
 
-        function getLevel(s: string): vscode.DiagnosticSeverity {
-            if (s === 'error') {
-                return vscode.DiagnosticSeverity.Error;
+        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
+            );
+        }
+
+        function areCodeActionsEqual(
+            left: vscode.CodeAction,
+            right: vscode.CodeAction
+        ): boolean {
+            if (
+                left.kind !== right.kind ||
+                left.title !== right.title ||
+                !left.edit ||
+                !right.edit
+            ) {
+                return false;
             }
-            if (s.startsWith('warn')) {
-                return vscode.DiagnosticSeverity.Warning;
+
+            const leftEditEntries = left.edit.entries();
+            const rightEditEntries = right.edit.entries();
+
+            if (leftEditEntries.length !== leftEditEntries.length) {
+                return false;
             }
-            return vscode.DiagnosticSeverity.Information;
-        }
 
-        // Reference:
-        // https://github.com/rust-lang/rust/blob/master/src/libsyntax/json.rs
-        interface RustDiagnosticSpan {
-            line_start: number;
-            line_end: number;
-            column_start: number;
-            column_end: number;
-            is_primary: boolean;
-            file_name: string;
-        }
+            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;
+                }
 
-        interface RustDiagnostic {
-            spans: RustDiagnosticSpan[];
-            rendered: string;
-            level: string;
-            code?: {
-                code: string;
-            };
+                if (leftEdits.length !== rightEdits.length) {
+                    return false;
+                }
+
+                for (let j = 0; j < leftEdits.length; j++) {
+                    const leftEdit = leftEdits[j];
+                    const rightEdit = rightEdits[j];
+
+                    if (!leftEdit.range.isEqual(rightEdit.range)) {
+                        return false;
+                    }
+
+                    if (leftEdit.newText !== rightEdit.newText) {
+                        return false;
+                    }
+                }
+            }
+
+            return true;
         }
 
         interface CargoArtifact {
@@ -215,41 +276,58 @@ export class CargoWatchProvider implements vscode.Disposable {
         } else if (data.reason === 'compiler-message') {
             const msg = data.message as RustDiagnostic;
 
-            const spans = msg.spans.filter(o => o.is_primary);
-
-            // We only handle primary span right now.
-            if (spans.length > 0) {
-                const o = spans[0];
+            const mapResult = mapRustDiagnosticToVsCode(msg);
+            if (!mapResult) {
+                return;
+            }
 
-                const rendered = msg.rendered;
-                const level = getLevel(msg.level);
-                const range = new vscode.Range(
-                    new vscode.Position(o.line_start - 1, o.column_start - 1),
-                    new vscode.Position(o.line_end - 1, o.column_end - 1)
-                );
+            const { location, diagnostic, codeActions } = mapResult;
+            const fileUri = location.uri;
 
-                const fileName = path.join(
-                    vscode.workspace.rootPath!,
-                    o.file_name
-                );
-                const diagnostic = new vscode.Diagnostic(
-                    range,
-                    rendered,
-                    level
-                );
+            const diagnostics: vscode.Diagnostic[] = [
+                ...(this.diagnosticCollection!.get(fileUri) || [])
+            ];
 
-                diagnostic.source = 'rustc';
-                diagnostic.code = msg.code ? msg.code.code : undefined;
-                diagnostic.relatedInformation = [];
+            // If we're building multiple targets it's possible we've already seen this diagnostic
+            const isDuplicate = diagnostics.some(d =>
+                areDiagnosticsEqual(d, diagnostic)
+            );
 
-                const fileUrl = vscode.Uri.file(fileName!);
+            if (isDuplicate) {
+                return;
+            }
 
-                const diagnostics: vscode.Diagnostic[] = [
-                    ...(this.diagnosticCollection!.get(fileUrl) || [])
-                ];
-                diagnostics.push(diagnostic);
+            diagnostics.push(diagnostic);
+            this.diagnosticCollection!.set(fileUri, diagnostics);
+
+            if (codeActions.length) {
+                const fileUriString = fileUri.toString();
+                const existingActions = this.codeActions[fileUriString] || [];
+
+                for (const newAction of codeActions) {
+                    const existingAction = existingActions.find(existing =>
+                        areCodeActionsEqual(existing, newAction)
+                    );
+
+                    if (existingAction) {
+                        if (!existingAction.diagnostics) {
+                            existingAction.diagnostics = [];
+                        }
+                        // This action also applies to this diagnostic
+                        existingAction.diagnostics.push(diagnostic);
+                    } else {
+                        newAction.diagnostics = [diagnostic];
+                        existingActions.push(newAction);
+                    }
+                }
 
-                this.diagnosticCollection!.set(fileUrl, diagnostics);
+                // Have VsCode query us for the code actions
+                this.codeActions[fileUriString] = existingActions;
+                vscode.commands.executeCommand(
+                    'vscode.executeCodeActionProvider',
+                    fileUri,
+                    diagnostic.range
+                );
             }
         }
     }
diff --git a/editors/code/src/utils/rust_diagnostics.ts b/editors/code/src/utils/rust_diagnostics.ts
new file mode 100644
index 000000000..7d8cc0e0b
--- /dev/null
+++ b/editors/code/src/utils/rust_diagnostics.ts
@@ -0,0 +1,220 @@
+import * as path from 'path';
+import * as vscode from 'vscode';
+
+// Reference:
+// https://github.com/rust-lang/rust/blob/master/src/libsyntax/json.rs
+export interface RustDiagnosticSpan {
+    line_start: number;
+    line_end: number;
+    column_start: number;
+    column_end: number;
+    is_primary: boolean;
+    file_name: string;
+    label?: string;
+    suggested_replacement?: string;
+    suggestion_applicability?:
+        | 'MachineApplicable'
+        | 'HasPlaceholders'
+        | 'MaybeIncorrect'
+        | 'Unspecified';
+}
+
+export interface RustDiagnostic {
+    spans: RustDiagnosticSpan[];
+    rendered: string;
+    message: string;
+    level: string;
+    code?: {
+        code: string;
+    };
+    children: RustDiagnostic[];
+}
+
+export interface MappedRustDiagnostic {
+    location: vscode.Location;
+    diagnostic: vscode.Diagnostic;
+    codeActions: vscode.CodeAction[];
+}
+
+interface MappedRustChildDiagnostic {
+    related?: vscode.DiagnosticRelatedInformation;
+    codeAction?: vscode.CodeAction;
+    messageLine?: string;
+}
+
+/**
+ * Converts a Rust level string to a VsCode severity
+ */
+function mapLevelToSeverity(s: string): vscode.DiagnosticSeverity {
+    if (s === 'error') {
+        return vscode.DiagnosticSeverity.Error;
+    }
+    if (s.startsWith('warn')) {
+        return vscode.DiagnosticSeverity.Warning;
+    }
+    return vscode.DiagnosticSeverity.Information;
+}
+
+/**
+ * Converts a Rust span to a VsCode location
+ */
+function mapSpanToLocation(span: RustDiagnosticSpan): vscode.Location {
+    const fileName = path.join(vscode.workspace.rootPath!, span.file_name);
+    const fileUri = vscode.Uri.file(fileName);
+
+    const range = new vscode.Range(
+        new vscode.Position(span.line_start - 1, span.column_start - 1),
+        new vscode.Position(span.line_end - 1, span.column_end - 1)
+    );
+
+    return new vscode.Location(fileUri, range);
+}
+
+/**
+ * Converts a secondary Rust span to a VsCode related information
+ *
+ * If the span is unlabelled this will return `undefined`.
+ */
+function mapSecondarySpanToRelated(
+    span: RustDiagnosticSpan
+): vscode.DiagnosticRelatedInformation | undefined {
+    if (!span.label) {
+        // Nothing to label this with
+        return;
+    }
+
+    const location = mapSpanToLocation(span);
+    return new vscode.DiagnosticRelatedInformation(location, span.label);
+}
+
+/**
+ * Determines if diagnostic is related to unused code
+ */
+function isUnusedOrUnnecessary(rd: RustDiagnostic): boolean {
+    if (!rd.code) {
+        return false;
+    }
+
+    const { code } = rd.code;
+    return code.startsWith('unused_') || code === 'dead_code';
+}
+
+/**
+ * Converts a Rust child diagnostic to a VsCode related information
+ *
+ * This can have three outcomes:
+ *
+ * 1. If this is no primary span this will return a `noteLine`
+ * 2. If there is a primary span with a suggested replacement it will return a
+ *    `codeAction`.
+ * 3. If there is a primary span without a suggested replacement it will return
+ *    a `related`.
+ */
+function mapRustChildDiagnostic(rd: RustDiagnostic): MappedRustChildDiagnostic {
+    const span = rd.spans.find(s => s.is_primary);
+
+    if (!span) {
+        // `rustc` uses these spanless children as a way to print multi-line
+        // messages
+        return { messageLine: rd.message };
+    }
+
+    // If we have a primary span use its location, otherwise use the parent
+    const location = mapSpanToLocation(span);
+
+    // 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
+        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 };
+    } else {
+        const related = new vscode.DiagnosticRelatedInformation(
+            location,
+            rd.message
+        );
+
+        return { related };
+    }
+}
+
+/**
+ * Converts a Rust root diagnostic to VsCode form
+ *
+ * This flattens the Rust diagnostic by:
+ *
+ * 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,
+ *    `relatedInformation` or additional message lines.
+ *
+ * If the diagnostic has no primary span this will return `undefined`
+ */
+export function mapRustDiagnosticToVsCode(
+    rd: RustDiagnostic
+): MappedRustDiagnostic | undefined {
+    const codeActions = [];
+
+    const primarySpan = rd.spans.find(s => s.is_primary);
+    if (!primarySpan) {
+        return;
+    }
+
+    const location = mapSpanToLocation(primarySpan);
+    const secondarySpans = rd.spans.filter(s => !s.is_primary);
+
+    const severity = mapLevelToSeverity(rd.level);
+
+    const vd = new vscode.Diagnostic(location.range, rd.message, severity);
+
+    vd.source = 'rustc';
+    vd.code = rd.code ? rd.code.code : undefined;
+    vd.relatedInformation = [];
+
+    for (const secondarySpan of secondarySpans) {
+        const related = mapSecondarySpanToRelated(secondarySpan);
+        if (related) {
+            vd.relatedInformation.push(related);
+        }
+    }
+
+    for (const child of rd.children) {
+        const { related, codeAction, messageLine } = mapRustChildDiagnostic(
+            child
+        );
+
+        if (related) {
+            vd.relatedInformation.push(related);
+        }
+        if (codeAction) {
+            codeActions.push(codeAction);
+        }
+        if (messageLine) {
+            vd.message += `\n${messageLine}`;
+        }
+    }
+
+    if (isUnusedOrUnnecessary(rd)) {
+        vd.tags = [vscode.DiagnosticTag.Unnecessary];
+    }
+
+    return {
+        location,
+        diagnostic: vd,
+        codeActions
+    };
+}