Add a fix that inserts a missing method

Author: MatrixFrog

crates/hir/src/has_source.rs

@@ -12,7 +12,7 @@ use syntax::ast;
 use tt::TextRange;
 
 use crate::{
-    Adt, Callee, Const, Enum, ExternCrateDecl, Field, FieldSource, Function, Impl,
+    Adt, AssocItem, Callee, Const, Enum, ExternCrateDecl, Field, FieldSource, Function, Impl,
     InlineAsmOperand, Label, LifetimeParam, LocalSource, Macro, Module, Param, SelfParam, Static,
     Struct, Trait, TraitAlias, TypeAlias, TypeOrConstParam, Union, Variant, VariantDef,
     db::HirDatabase,
@@ -100,6 +100,17 @@ impl HasSource for Field {
         Some(field_source)
     }
 }
+impl HasSource for AssocItem {
+    type Ast = ast::AssocItem;
+
+    fn source(self, db: &dyn HirDatabase) -> Option<InFile<Self::Ast>> {
+        match self {
+            AssocItem::Const(c) => Some(c.source(db)?.map(ast::AssocItem::Const)),
+            AssocItem::Function(f) => Some(f.source(db)?.map(ast::AssocItem::Fn)),
+            AssocItem::TypeAlias(t) => Some(t.source(db)?.map(ast::AssocItem::TypeAlias)),
+        }
+    }
+}
 impl HasSource for Adt {
     type Ast = ast::Adt;
     fn source(self, db: &dyn HirDatabase) -> Option<InFile<Self::Ast>> {

crates/ide-diagnostics/src/handlers/unresolved_method.rs

@@ -1,3 +1,4 @@
+use hir::HasSource;
 use hir::{FileRange, HirDisplay, InFile, db::ExpandDatabase};
 use ide_db::text_edit::TextEdit;
 use ide_db::{
@@ -11,7 +12,7 @@ use syntax::{
     format_smolstr,
 };
 
-use crate::{Diagnostic, DiagnosticCode, DiagnosticsContext, adjusted_display_range};
+use crate::{Diagnostic, DiagnosticCode, DiagnosticsContext, adjusted_display_range, fix};
 
 // Diagnostic: unresolved-method
 //
@@ -67,9 +68,44 @@ fn fixes(ctx: &DiagnosticsContext<'_>, d: &hir::UnresolvedMethodCall) -> Option<
         fixes.push(assoc_func_fix);
     }
 
+    if let Some(method_fix) = add_method_fix(ctx, d) {
+        fixes.push(method_fix);
+    }
+
     if fixes.is_empty() { None } else { Some(fixes) }
 }
 
+/// Fix to add the missing method.
+fn add_method_fix(ctx: &DiagnosticsContext<'_>, d: &hir::UnresolvedMethodCall) -> Option<Assist> {
+    let root = ctx.sema.db.parse_or_expand(d.expr.file_id);
+    let expr = d.expr.value.to_node(&root).left()?;
+
+    let db = ctx.sema.db;
+    let ty = d.receiver.clone();
+
+    let impl_block =
+        hir::Impl::all_for_type(db, ty).into_iter().find(|block| block.trait_(db).is_none())?;
+    let items = impl_block.items(db);
+    let last_item = items.last()?;
+    let source = last_item.source(db)?;
+    let file_id = match source.file_id {
+        hir::HirFileId::FileId(file_id) => file_id,
+        hir::HirFileId::MacroFile(_) => return None,
+    };
+    let end_of_last_item = source.node_file_range().file_range()?.range.end();
+
+    let text_to_insert = format!("\n  fn {}(&self) {{}}", d.name.as_str());
+    Some(fix(
+        "add-missing-method",
+        "Add missing method",
+        SourceChange::from_text_edit(
+            file_id.file_id(db),
+            TextEdit::insert(end_of_last_item, text_to_insert),
+        ),
+        ctx.sema.original_range(expr.syntax()).range,
+    ))
+}
+
 fn field_fix(
     ctx: &DiagnosticsContext<'_>,
     d: &hir::UnresolvedMethodCall,
@@ -286,6 +322,35 @@ fn main() {
         );
     }
 
+    #[test]
+    fn test_add_method_fix() {
+        check_fix(
+            r#"
+struct Tiger;
+
+impl Tiger {
+  fn sleep(&self) {}
+}
+
+fn main() {
+  let t = Tiger;
+  t.roar$0();
+}"#,
+            r#"
+struct Tiger;
+
+impl Tiger {
+  fn sleep(&self) {}
+  fn roar(&self) {}
+}
+
+fn main() {
+  let t = Tiger;
+  t.roar();
+}"#,
+        );
+    }
+
     #[test]
     fn smoke_test_in_macro_def_site() {
         check_diagnostics(