3 files changed, 211 insertions, 154 deletions

diff --git a/crates/ra_ide/src/diagnostics.rs b/crates/ra_ide/src/diagnostics.rs
index 73c0b8275..1046d7ab3 100644
--- a/crates/ra_ide/src/diagnostics.rs
+++ b/crates/ra_ide/src/diagnostics.rs
@@ -6,22 +6,21 @@
 
 use std::cell::RefCell;
 
-use hir::{
-    diagnostics::{AstDiagnostic, Diagnostic as _, DiagnosticSinkBuilder},
-    HasSource, HirDisplay, Semantics, VariantDef,
-};
+use hir::{diagnostics::DiagnosticSinkBuilder, Semantics};
 use itertools::Itertools;
 use ra_db::SourceDatabase;
 use ra_ide_db::RootDatabase;
 use ra_prof::profile;
 use ra_syntax::{
-    algo,
-    ast::{self, edit::IndentLevel, make, AstNode},
+    ast::{self, AstNode},
     SyntaxNode, TextRange, T,
 };
 use ra_text_edit::{TextEdit, TextEditBuilder};
 
-use crate::{Diagnostic, FileId, FileSystemEdit, Fix, SourceFileEdit};
+use crate::{Diagnostic, FileId, Fix, SourceFileEdit};
+
+mod diagnostics_with_fix;
+use diagnostics_with_fix::DiagnosticWithFix;
 
 #[derive(Debug, Copy, Clone)]
 pub enum Severity {
@@ -54,73 +53,16 @@ pub(crate) fn diagnostics(
     let res = RefCell::new(res);
     let mut sink = DiagnosticSinkBuilder::new()
         .on::<hir::diagnostics::UnresolvedModule, _>(|d| {
-            let original_file = d.source().file_id.original_file(db);
-            let fix = Fix::new(
-                "Create module",
-                FileSystemEdit::CreateFile { anchor: original_file, dst: d.candidate.clone() }
-                    .into(),
-            );
-            res.borrow_mut().push(Diagnostic {
-                range: sema.diagnostics_range(d).range,
-                message: d.message(),
-                severity: Severity::Error,
-                fix: Some(fix),
-            })
+            res.borrow_mut().push(diagnostic_with_fix(d, &sema));
         })
         .on::<hir::diagnostics::MissingFields, _>(|d| {
-            // Note that although we could add a diagnostics to
-            // fill the missing tuple field, e.g :
-            // `struct A(usize);`
-            // `let a = A { 0: () }`
-            // but it is uncommon usage and it should not be encouraged.
-            let fix = if d.missed_fields.iter().any(|it| it.as_tuple_index().is_some()) {
-                None
-            } else {
-                let mut field_list = d.ast(db);
-                for f in d.missed_fields.iter() {
-                    let field = make::record_expr_field(
-                        make::name_ref(&f.to_string()),
-                        Some(make::expr_unit()),
-                    );
-                    field_list = field_list.append_field(&field);
-                }
-
-                let edit = {
-                    let mut builder = TextEditBuilder::default();
-                    algo::diff(&d.ast(db).syntax(), &field_list.syntax())
-                        .into_text_edit(&mut builder);
-                    builder.finish()
-                };
-                Some(Fix::new("Fill struct fields", SourceFileEdit { file_id, edit }.into()))
-            };
-
-            res.borrow_mut().push(Diagnostic {
-                range: sema.diagnostics_range(d).range,
-                message: d.message(),
-                severity: Severity::Error,
-                fix,
-            })
+            res.borrow_mut().push(diagnostic_with_fix(d, &sema));
         })
         .on::<hir::diagnostics::MissingOkInTailExpr, _>(|d| {
-            let node = d.ast(db);
-            let replacement = format!("Ok({})", node.syntax());
-            let edit = TextEdit::replace(node.syntax().text_range(), replacement);
-            let source_change = SourceFileEdit { file_id, edit }.into();
-            let fix = Fix::new("Wrap with ok", source_change);
-            res.borrow_mut().push(Diagnostic {
-                range: sema.diagnostics_range(d).range,
-                message: d.message(),
-                severity: Severity::Error,
-                fix: Some(fix),
-            })
+            res.borrow_mut().push(diagnostic_with_fix(d, &sema));
         })
         .on::<hir::diagnostics::NoSuchField, _>(|d| {
-            res.borrow_mut().push(Diagnostic {
-                range: sema.diagnostics_range(d).range,
-                message: d.message(),
-                severity: Severity::Error,
-                fix: missing_struct_field_fix(&sema, file_id, d),
-            })
+            res.borrow_mut().push(diagnostic_with_fix(d, &sema));
         })
         // Only collect experimental diagnostics when they're enabled.
         .filter(|diag| !diag.is_experimental() || enable_experimental)
@@ -128,7 +70,7 @@ pub(crate) fn diagnostics(
         .build(|d| {
             res.borrow_mut().push(Diagnostic {
                 message: d.message(),
-                range: sema.diagnostics_range(d).range,
+                range: sema.diagnostics_display_range(d).range,
                 severity: Severity::Error,
                 fix: None,
             })
@@ -141,77 +83,12 @@ pub(crate) fn diagnostics(
     res.into_inner()
 }
 
-fn missing_struct_field_fix(
-    sema: &Semantics<RootDatabase>,
-    usage_file_id: FileId,
-    d: &hir::diagnostics::NoSuchField,
-) -> Option<Fix> {
-    let record_expr = sema.ast(d);
-
-    let record_lit = ast::RecordExpr::cast(record_expr.syntax().parent()?.parent()?)?;
-    let def_id = sema.resolve_variant(record_lit)?;
-    let module;
-    let def_file_id;
-    let record_fields = match VariantDef::from(def_id) {
-        VariantDef::Struct(s) => {
-            module = s.module(sema.db);
-            let source = s.source(sema.db);
-            def_file_id = source.file_id;
-            let fields = source.value.field_list()?;
-            record_field_list(fields)?
-        }
-        VariantDef::Union(u) => {
-            module = u.module(sema.db);
-            let source = u.source(sema.db);
-            def_file_id = source.file_id;
-            source.value.record_field_list()?
-        }
-        VariantDef::EnumVariant(e) => {
-            module = e.module(sema.db);
-            let source = e.source(sema.db);
-            def_file_id = source.file_id;
-            let fields = source.value.field_list()?;
-            record_field_list(fields)?
-        }
-    };
-    let def_file_id = def_file_id.original_file(sema.db);
-
-    let new_field_type = sema.type_of_expr(&record_expr.expr()?)?;
-    if new_field_type.is_unknown() {
-        return None;
-    }
-    let new_field = make::record_field(
-        record_expr.field_name()?,
-        make::ty(&new_field_type.display_source_code(sema.db, module.into()).ok()?),
-    );
-
-    let last_field = record_fields.fields().last()?;
-    let last_field_syntax = last_field.syntax();
-    let indent = IndentLevel::from_node(last_field_syntax);
-
-    let mut new_field = new_field.to_string();
-    if usage_file_id != def_file_id {
-        new_field = format!("pub(crate) {}", new_field);
-    }
-    new_field = format!("\n{}{}", indent, new_field);
-
-    let needs_comma = !last_field_syntax.to_string().ends_with(',');
-    if needs_comma {
-        new_field = format!(",{}", new_field);
-    }
-
-    let source_change = SourceFileEdit {
-        file_id: def_file_id,
-        edit: TextEdit::insert(last_field_syntax.text_range().end(), new_field),
-    };
-    let fix = Fix::new("Create field", source_change.into());
-    return Some(fix);
-
-    fn record_field_list(field_def_list: ast::FieldList) -> Option<ast::RecordFieldList> {
-        match field_def_list {
-            ast::FieldList::RecordFieldList(it) => Some(it),
-            ast::FieldList::TupleFieldList(_) => None,
-        }
+fn diagnostic_with_fix<D: DiagnosticWithFix>(d: &D, sema: &Semantics<RootDatabase>) -> Diagnostic {
+    Diagnostic {
+        range: sema.diagnostics_display_range(d).range,
+        message: d.message(),
+        severity: Severity::Error,
+        fix: d.fix(&sema),
     }
 }
 
@@ -222,24 +99,25 @@ fn check_unnecessary_braces_in_use_statement(
 ) -> Option<()> {
     let use_tree_list = ast::UseTreeList::cast(node.clone())?;
     if let Some((single_use_tree,)) = use_tree_list.use_trees().collect_tuple() {
-        let range = use_tree_list.syntax().text_range();
+        let use_range = use_tree_list.syntax().text_range();
         let edit =
             text_edit_for_remove_unnecessary_braces_with_self_in_use_statement(&single_use_tree)
                 .unwrap_or_else(|| {
                     let to_replace = single_use_tree.syntax().text().to_string();
                     let mut edit_builder = TextEditBuilder::default();
-                    edit_builder.delete(range);
-                    edit_builder.insert(range.start(), to_replace);
+                    edit_builder.delete(use_range);
+                    edit_builder.insert(use_range.start(), to_replace);
                     edit_builder.finish()
                 });
 
         acc.push(Diagnostic {
-            range,
+            range: use_range,
             message: "Unnecessary braces in use statement".to_string(),
             severity: Severity::WeakWarning,
             fix: Some(Fix::new(
                 "Remove unnecessary braces",
                 SourceFileEdit { file_id, edit }.into(),
+                use_range,
             )),
         });
     }
@@ -254,8 +132,7 @@ fn text_edit_for_remove_unnecessary_braces_with_self_in_use_statement(
     if single_use_tree.path()?.segment()?.syntax().first_child_or_token()?.kind() == T![self] {
         let start = use_tree_list_node.prev_sibling_or_token()?.text_range().start();
         let end = use_tree_list_node.text_range().end();
-        let range = TextRange::new(start, end);
-        return Some(TextEdit::delete(range));
+        return Some(TextEdit::delete(TextRange::new(start, end)));
     }
     None
 }
@@ -278,13 +155,15 @@ fn check_struct_shorthand_initialization(
                 edit_builder.insert(record_field.syntax().text_range().start(), field_name);
                 let edit = edit_builder.finish();
 
+                let field_range = record_field.syntax().text_range();
                 acc.push(Diagnostic {
-                    range: record_field.syntax().text_range(),
+                    range: field_range,
                     message: "Shorthand struct initialization".to_string(),
                     severity: Severity::WeakWarning,
                     fix: Some(Fix::new(
                         "Use struct shorthand initialization",
                         SourceFileEdit { file_id, edit }.into(),
+                        field_range,
                     )),
                 });
             }
@@ -304,7 +183,7 @@ mod tests {
     /// Takes a multi-file input fixture with annotated cursor positions,
     /// and checks that:
     ///  * a diagnostic is produced
-    ///  * this diagnostic touches the input cursor position
+    ///  * this diagnostic fix trigger range touches the input cursor position
     ///  * that the contents of the file containing the cursor match `after` after the diagnostic fix is applied
     fn check_fix(ra_fixture_before: &str, ra_fixture_after: &str) {
         let after = trim_indent(ra_fixture_after);
@@ -322,10 +201,10 @@ mod tests {
 
         assert_eq_text!(&after, &actual);
         assert!(
-            diagnostic.range.start() <= file_position.offset
-                && diagnostic.range.end() >= file_position.offset,
-            "diagnostic range {:?} does not touch cursor position {:?}",
-            diagnostic.range,
+            fix.fix_trigger_range.start() <= file_position.offset
+                && fix.fix_trigger_range.end() >= file_position.offset,
+            "diagnostic fix range {:?} does not touch cursor position {:?}",
+            fix.fix_trigger_range,
             file_position.offset
         );
     }
@@ -642,6 +521,7 @@ fn test_fn() {
                                     ],
                                     is_snippet: false,
                                 },
+                                fix_trigger_range: 0..8,
                             },
                         ),
                     },
diff --git a/crates/ra_ide/src/diagnostics/diagnostics_with_fix.rs b/crates/ra_ide/src/diagnostics/diagnostics_with_fix.rs
new file mode 100644
index 000000000..f7c73773f
--- /dev/null
+++ b/crates/ra_ide/src/diagnostics/diagnostics_with_fix.rs
@@ -0,0 +1,171 @@
+//! Provides a way to attach fixes to the diagnostics.
+//! The same module also has all curret custom fixes for the diagnostics implemented.
+use crate::Fix;
+use ast::{edit::IndentLevel, make};
+use hir::{
+    db::AstDatabase,
+    diagnostics::{Diagnostic, MissingFields, MissingOkInTailExpr, NoSuchField, UnresolvedModule},
+    HasSource, HirDisplay, Semantics, VariantDef,
+};
+use ra_db::FileId;
+use ra_ide_db::{
+    source_change::{FileSystemEdit, SourceFileEdit},
+    RootDatabase,
+};
+use ra_syntax::{algo, ast, AstNode};
+use ra_text_edit::{TextEdit, TextEditBuilder};
+
+/// A [Diagnostic] that potentially has a fix available.
+///
+/// [Diagnostic]: hir::diagnostics::Diagnostic
+pub trait DiagnosticWithFix: Diagnostic {
+    fn fix(&self, sema: &Semantics<RootDatabase>) -> Option<Fix>;
+}
+
+impl DiagnosticWithFix for UnresolvedModule {
+    fn fix(&self, sema: &Semantics<RootDatabase>) -> Option<Fix> {
+        let root = sema.db.parse_or_expand(self.file)?;
+        let unresolved_module = self.decl.to_node(&root);
+        Some(Fix::new(
+            "Create module",
+            FileSystemEdit::CreateFile {
+                anchor: self.file.original_file(sema.db),
+                dst: self.candidate.clone(),
+            }
+            .into(),
+            unresolved_module.syntax().text_range(),
+        ))
+    }
+}
+
+impl DiagnosticWithFix for NoSuchField {
+    fn fix(&self, sema: &Semantics<RootDatabase>) -> Option<Fix> {
+        let root = sema.db.parse_or_expand(self.file)?;
+        missing_record_expr_field_fix(
+            &sema,
+            self.file.original_file(sema.db),
+            &self.field.to_node(&root),
+        )
+    }
+}
+
+impl DiagnosticWithFix for MissingFields {
+    fn fix(&self, sema: &Semantics<RootDatabase>) -> Option<Fix> {
+        // Note that although we could add a diagnostics to
+        // fill the missing tuple field, e.g :
+        // `struct A(usize);`
+        // `let a = A { 0: () }`
+        // but it is uncommon usage and it should not be encouraged.
+        if self.missed_fields.iter().any(|it| it.as_tuple_index().is_some()) {
+            return None;
+        }
+
+        let root = sema.db.parse_or_expand(self.file)?;
+        let old_field_list = self.field_list_parent.to_node(&root).record_expr_field_list()?;
+        let mut new_field_list = old_field_list.clone();
+        for f in self.missed_fields.iter() {
+            let field =
+                make::record_expr_field(make::name_ref(&f.to_string()), Some(make::expr_unit()));
+            new_field_list = new_field_list.append_field(&field);
+        }
+
+        let edit = {
+            let mut builder = TextEditBuilder::default();
+            algo::diff(&old_field_list.syntax(), &new_field_list.syntax())
+                .into_text_edit(&mut builder);
+            builder.finish()
+        };
+        Some(Fix::new(
+            "Fill struct fields",
+            SourceFileEdit { file_id: self.file.original_file(sema.db), edit }.into(),
+            sema.original_range(&old_field_list.syntax()).range,
+        ))
+    }
+}
+
+impl DiagnosticWithFix for MissingOkInTailExpr {
+    fn fix(&self, sema: &Semantics<RootDatabase>) -> Option<Fix> {
+        let root = sema.db.parse_or_expand(self.file)?;
+        let tail_expr = self.expr.to_node(&root);
+        let tail_expr_range = tail_expr.syntax().text_range();
+        let edit = TextEdit::replace(tail_expr_range, format!("Ok({})", tail_expr.syntax()));
+        let source_change =
+            SourceFileEdit { file_id: self.file.original_file(sema.db), edit }.into();
+        Some(Fix::new("Wrap with ok", source_change, tail_expr_range))
+    }
+}
+
+fn missing_record_expr_field_fix(
+    sema: &Semantics<RootDatabase>,
+    usage_file_id: FileId,
+    record_expr_field: &ast::RecordExprField,
+) -> Option<Fix> {
+    let record_lit = ast::RecordExpr::cast(record_expr_field.syntax().parent()?.parent()?)?;
+    let def_id = sema.resolve_variant(record_lit)?;
+    let module;
+    let def_file_id;
+    let record_fields = match VariantDef::from(def_id) {
+        VariantDef::Struct(s) => {
+            module = s.module(sema.db);
+            let source = s.source(sema.db);
+            def_file_id = source.file_id;
+            let fields = source.value.field_list()?;
+            record_field_list(fields)?
+        }
+        VariantDef::Union(u) => {
+            module = u.module(sema.db);
+            let source = u.source(sema.db);
+            def_file_id = source.file_id;
+            source.value.record_field_list()?
+        }
+        VariantDef::EnumVariant(e) => {
+            module = e.module(sema.db);
+            let source = e.source(sema.db);
+            def_file_id = source.file_id;
+            let fields = source.value.field_list()?;
+            record_field_list(fields)?
+        }
+    };
+    let def_file_id = def_file_id.original_file(sema.db);
+
+    let new_field_type = sema.type_of_expr(&record_expr_field.expr()?)?;
+    if new_field_type.is_unknown() {
+        return None;
+    }
+    let new_field = make::record_field(
+        record_expr_field.field_name()?,
+        make::ty(&new_field_type.display_source_code(sema.db, module.into()).ok()?),
+    );
+
+    let last_field = record_fields.fields().last()?;
+    let last_field_syntax = last_field.syntax();
+    let indent = IndentLevel::from_node(last_field_syntax);
+
+    let mut new_field = new_field.to_string();
+    if usage_file_id != def_file_id {
+        new_field = format!("pub(crate) {}", new_field);
+    }
+    new_field = format!("\n{}{}", indent, new_field);
+
+    let needs_comma = !last_field_syntax.to_string().ends_with(',');
+    if needs_comma {
+        new_field = format!(",{}", new_field);
+    }
+
+    let source_change = SourceFileEdit {
+        file_id: def_file_id,
+        edit: TextEdit::insert(last_field_syntax.text_range().end(), new_field),
+    };
+    return Some(Fix::new(
+        "Create field",
+        source_change.into(),
+        record_expr_field.syntax().text_range(),
+    ));
+
+    fn record_field_list(field_def_list: ast::FieldList) -> Option<ast::RecordFieldList> {
+        match field_def_list {
+            ast::FieldList::RecordFieldList(it) => Some(it),
+            ast::FieldList::TupleFieldList(_) => None,
+        }
+    }
+}
diff --git a/crates/ra_ide/src/lib.rs b/crates/ra_ide/src/lib.rs
index 0fede0d87..89fcb6f17 100644
--- a/crates/ra_ide/src/lib.rs
+++ b/crates/ra_ide/src/lib.rs
@@ -112,13 +112,19 @@ pub struct Diagnostic {
 pub struct Fix {
     pub label: String,
     pub source_change: SourceChange,
+    /// Allows to trigger the fix only when the caret is in the range given
+    pub fix_trigger_range: TextRange,
 }
 
 impl Fix {
-    pub fn new(label: impl Into<String>, source_change: SourceChange) -> Self {
+    pub fn new(
+        label: impl Into<String>,
+        source_change: SourceChange,
+        fix_trigger_range: TextRange,
+    ) -> Self {
         let label = label.into();
         assert!(label.starts_with(char::is_uppercase) && !label.ends_with('.'));
-        Self { label, source_change }
+        Self { label, source_change, fix_trigger_range }
     }
 }