Add quickfix to add a struct field

5 files changed, 134 insertions, 5 deletions

diff --git a/crates/ra_hir/src/semantics.rs b/crates/ra_hir/src/semantics.rs
index a232a5856..6a49c424a 100644
--- a/crates/ra_hir/src/semantics.rs
+++ b/crates/ra_hir/src/semantics.rs
@@ -6,9 +6,9 @@ use std::{cell::RefCell, fmt, iter::successors};
 
 use hir_def::{
     resolver::{self, HasResolver, Resolver},
-    AsMacroCall, TraitId,
+    AsMacroCall, TraitId, VariantId,
 };
-use hir_expand::{hygiene::Hygiene, ExpansionInfo};
+use hir_expand::{diagnostics::AstDiagnostic, hygiene::Hygiene, ExpansionInfo};
 use hir_ty::associated_type_shorthand_candidates;
 use itertools::Itertools;
 use ra_db::{FileId, FileRange};
@@ -104,6 +104,13 @@ impl<'db, DB: HirDatabase> Semantics<'db, DB> {
         tree
     }
 
+    pub fn ast<T: AstDiagnostic + Diagnostic>(&self, d: &T) -> <T as AstDiagnostic>::AST {
+        let file_id = d.source().file_id;
+        let root = self.db.parse_or_expand(file_id).unwrap();
+        self.cache(root, file_id);
+        d.ast(self.db)
+    }
+
     pub fn expand(&self, macro_call: &ast::MacroCall) -> Option<SyntaxNode> {
         let macro_call = self.find_file(macro_call.syntax().clone()).with_value(macro_call);
         let sa = self.analyze2(macro_call.map(|it| it.syntax()), None);
@@ -247,6 +254,10 @@ impl<'db, DB: HirDatabase> Semantics<'db, DB> {
         self.analyze(path.syntax()).resolve_path(self.db, path)
     }
 
+    pub fn resolve_variant(&self, record_lit: ast::RecordLit) -> Option<VariantId> {
+        self.analyze(record_lit.syntax()).resolve_variant(self.db, record_lit)
+    }
+
     pub fn lower_path(&self, path: &ast::Path) -> Option<Path> {
         let src = self.find_file(path.syntax().clone());
         Path::from_src(path.clone(), &Hygiene::new(self.db.upcast(), src.file_id.into()))
diff --git a/crates/ra_hir/src/source_analyzer.rs b/crates/ra_hir/src/source_analyzer.rs
index 7c6bbea13..757d1e397 100644
--- a/crates/ra_hir/src/source_analyzer.rs
+++ b/crates/ra_hir/src/source_analyzer.rs
@@ -313,6 +313,16 @@ impl SourceAnalyzer {
         })?;
         Some(macro_call_id.as_file())
     }
+
+    pub(crate) fn resolve_variant(
+        &self,
+        db: &dyn HirDatabase,
+        record_lit: ast::RecordLit,
+    ) -> Option<VariantId> {
+        let infer = self.infer.as_ref()?;
+        let expr_id = self.expr_id(db, &record_lit.into())?;
+        infer.variant_resolution_for_expr(expr_id)
+    }
 }
 
 fn scope_for(
diff --git a/crates/ra_hir_ty/src/diagnostics.rs b/crates/ra_hir_ty/src/diagnostics.rs
index 2c7298714..ebd9cb08f 100644
--- a/crates/ra_hir_ty/src/diagnostics.rs
+++ b/crates/ra_hir_ty/src/diagnostics.rs
@@ -6,7 +6,7 @@ use hir_expand::{db::AstDatabase, name::Name, HirFileId, InFile};
 use ra_syntax::{ast, AstNode, AstPtr, SyntaxNodePtr};
 use stdx::format_to;
 
-pub use hir_def::{diagnostics::UnresolvedModule, expr::MatchArm};
+pub use hir_def::{diagnostics::UnresolvedModule, expr::MatchArm, path::Path};
 pub use hir_expand::diagnostics::{AstDiagnostic, Diagnostic, DiagnosticSink};
 
 #[derive(Debug)]
@@ -29,6 +29,16 @@ impl Diagnostic for NoSuchField {
     }
 }
 
+impl AstDiagnostic for NoSuchField {
+    type AST = ast::RecordField;
+
+    fn ast(&self, db: &impl AstDatabase) -> Self::AST {
+        let root = db.parse_or_expand(self.source().file_id).unwrap();
+        let node = self.source().value.to_node(&root);
+        ast::RecordField::cast(node).unwrap()
+    }
+}
+
 #[derive(Debug)]
 pub struct MissingFields {
     pub file: HirFileId,
diff --git a/crates/ra_ide/src/diagnostics.rs b/crates/ra_ide/src/diagnostics.rs
index e1bfd72f9..365d52168 100644
--- a/crates/ra_ide/src/diagnostics.rs
+++ b/crates/ra_ide/src/diagnostics.rs
@@ -8,7 +8,7 @@ use std::cell::RefCell;
 
 use hir::{
     diagnostics::{AstDiagnostic, Diagnostic as _, DiagnosticSink},
-    Semantics,
+    HasSource, HirDisplay, Semantics, VariantDef,
 };
 use itertools::Itertools;
 use ra_db::{RelativePath, SourceDatabase, SourceDatabaseExt};
@@ -16,7 +16,7 @@ use ra_ide_db::RootDatabase;
 use ra_prof::profile;
 use ra_syntax::{
     algo,
-    ast::{self, make, AstNode},
+    ast::{self, edit::IndentLevel, make, AstNode},
     SyntaxNode, TextRange, T,
 };
 use ra_text_edit::{TextEdit, TextEditBuilder};
@@ -123,7 +123,16 @@ pub(crate) fn diagnostics(db: &RootDatabase, file_id: FileId) -> Vec<Diagnostic>
             severity: Severity::Error,
             fix: Some(fix),
         })
+    })
+    .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),
+        })
     });
+
     if let Some(m) = sema.to_module_def(file_id) {
         m.diagnostics(db, &mut sink);
     };
@@ -131,6 +140,68 @@ pub(crate) fn diagnostics(db: &RootDatabase, file_id: FileId) -> Vec<Diagnostic>
     res.into_inner()
 }
 
+fn missing_struct_field_fix(
+    sema: &Semantics<RootDatabase>,
+    file_id: FileId,
+    d: &hir::diagnostics::NoSuchField,
+) -> Option<Fix> {
+    let record_expr = sema.ast(d);
+
+    let record_lit = ast::RecordLit::cast(record_expr.syntax().parent()?.parent()?)?;
+    let def_id = sema.resolve_variant(record_lit)?;
+    let module;
+    let record_fields = match VariantDef::from(def_id) {
+        VariantDef::Struct(s) => {
+            module = s.module(sema.db);
+            let source = s.source(sema.db);
+            let fields = source.value.field_def_list()?;
+            record_field_def_list(fields)?
+        }
+        VariantDef::Union(u) => {
+            module = u.module(sema.db);
+            let source = u.source(sema.db);
+            source.value.record_field_def_list()?
+        }
+        VariantDef::EnumVariant(e) => {
+            module = e.module(sema.db);
+            let source = e.source(sema.db);
+            let fields = source.value.field_def_list()?;
+            record_field_def_list(fields)?
+        }
+    };
+
+    let new_field_type = sema.type_of_expr(&record_expr.expr()?)?;
+    let new_field = make::record_field_def(
+        record_expr.field_name()?,
+        make::type_ref(&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 = 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,
+        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_def_list(field_def_list: ast::FieldDefList) -> Option<ast::RecordFieldDefList> {
+        match field_def_list {
+            ast::FieldDefList::RecordFieldDefList(it) => Some(it),
+            ast::FieldDefList::TupleFieldDefList(_) => None,
+        }
+    }
+}
+
 fn check_unnecessary_braces_in_use_statement(
     acc: &mut Vec<Diagnostic>,
     file_id: FileId,
@@ -795,4 +866,27 @@ fn main() {
             check_struct_shorthand_initialization,
         );
     }
+
+    #[test]
+    fn test_add_field_from_usage() {
+        check_apply_diagnostic_fix(
+            r"
+            fn main() {
+                Foo { bar: 3, baz: false};
+            }
+            struct Foo {
+                bar: i32
+            }
+            ",
+            r"
+            fn main() {
+                Foo { bar: 3, baz: false};
+            }
+            struct Foo {
+                bar: i32,
+                baz: bool
+            }
+            ",
+        )
+    }
 }
diff --git a/crates/ra_syntax/src/ast/make.rs b/crates/ra_syntax/src/ast/make.rs
index da0eb0926..192c610f1 100644
--- a/crates/ra_syntax/src/ast/make.rs
+++ b/crates/ra_syntax/src/ast/make.rs
@@ -75,6 +75,10 @@ pub fn record_field(name: ast::NameRef, expr: Option<ast::Expr>) -> ast::RecordF
     }
 }
 
+pub fn record_field_def(name: ast::NameRef, ty: ast::TypeRef) -> ast::RecordFieldDef {
+    ast_from_text(&format!("struct S {{ {}: {}, }}", name, ty))
+}
+
 pub fn block_expr(
     stmts: impl IntoIterator<Item = ast::Stmt>,
     tail_expr: Option<ast::Expr>,