Merge #9804

9804: Generate method from call r=matklad a=mahdi-frms

Needs a bit of refactoring. Tests also should be added.

Co-authored-by: mahdi-frms <mahdif1380@outlook.com>

crates/ide_assists/src/handlers/generate_function.rs

@@ -1,4 +1,4 @@
-use hir::{HirDisplay, TypeInfo};
+use hir::{HasSource, HirDisplay, InFile, Module, TypeInfo};
 use ide_db::{base_db::FileId, helpers::SnippetCap};
 use rustc_hash::{FxHashMap, FxHashSet};
 use stdx::to_lower_snake_case;
@@ -13,7 +13,7 @@ use syntax::{
 
 use crate::{
     utils::useless_type_special_case,
-    utils::{render_snippet, Cursor},
+    utils::{find_struct_impl, render_snippet, Cursor},
     AssistContext, AssistId, AssistKind, Assists,
 };
 
@@ -43,6 +43,31 @@ use crate::{
 //
 // ```
 pub(crate) fn generate_function(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
+    gen_fn(acc, ctx).or_else(|| gen_method(acc, ctx))
+}
+
+enum FuncExpr {
+    Func(ast::CallExpr),
+    Method(ast::MethodCallExpr),
+}
+
+impl FuncExpr {
+    fn arg_list(&self) -> Option<ast::ArgList> {
+        match self {
+            FuncExpr::Func(fn_call) => fn_call.arg_list(),
+            FuncExpr::Method(m_call) => m_call.arg_list(),
+        }
+    }
+
+    fn syntax(&self) -> &SyntaxNode {
+        match self {
+            FuncExpr::Func(fn_call) => fn_call.syntax(),
+            FuncExpr::Method(m_call) => m_call.syntax(),
+        }
+    }
+}
+
+fn gen_fn(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
     let path_expr: ast::PathExpr = ctx.find_node_at_offset()?;
     let call = path_expr.syntax().parent().and_then(ast::CallExpr::cast)?;
 
@@ -79,6 +104,69 @@ pub(crate) fn generate_function(acc: &mut Assists, ctx: &AssistContext) -> Optio
     )
 }
 
+fn gen_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
+    let call: ast::MethodCallExpr = ctx.find_node_at_offset()?;
+    let fn_name: ast::NameRef = ast::NameRef::cast(
+        call.syntax().children().find(|child| child.kind() == SyntaxKind::NAME_REF)?,
+    )?;
+    let ty = ctx.sema.type_of_expr(&call.receiver()?)?.original().strip_references().as_adt()?;
+
+    let current_module =
+        ctx.sema.scope(ctx.find_node_at_offset::<ast::MethodCallExpr>()?.syntax()).module()?;
+    let target_module = ty.module(ctx.sema.db);
+
+    if current_module.krate() != target_module.krate() {
+        return None;
+    }
+
+    let (impl_, file) = match ty {
+        hir::Adt::Struct(strukt) => get_impl(strukt.source(ctx.sema.db)?.syntax(), &fn_name, ctx),
+        hir::Adt::Enum(en) => get_impl(en.source(ctx.sema.db)?.syntax(), &fn_name, ctx),
+        hir::Adt::Union(union) => get_impl(union.source(ctx.sema.db)?.syntax(), &fn_name, ctx),
+    }?;
+
+    let function_builder = FunctionBuilder::from_method_call(
+        ctx,
+        &call,
+        &fn_name,
+        &impl_,
+        file,
+        target_module,
+        current_module,
+    )?;
+    let target = call.syntax().text_range();
+
+    acc.add(
+        AssistId("generate_function", AssistKind::Generate),
+        format!("Generate `{}` function", function_builder.fn_name),
+        target,
+        |builder| {
+            let function_template = function_builder.render();
+            builder.edit_file(function_template.file);
+            let mut new_fn = function_template.to_string(ctx.config.snippet_cap);
+            if impl_.is_none() {
+                new_fn = format!("\nimpl {} {{\n{}\n}}", ty.name(ctx.sema.db), new_fn,);
+            }
+            match ctx.config.snippet_cap {
+                Some(cap) => builder.insert_snippet(cap, function_template.insert_offset, new_fn),
+                None => builder.insert(function_template.insert_offset, new_fn),
+            }
+        },
+    )
+}
+
+fn get_impl(
+    adt: InFile<&SyntaxNode>,
+    fn_name: &ast::NameRef,
+    ctx: &AssistContext,
+) -> Option<(Option<ast::Impl>, FileId)> {
+    let file = adt.file_id.original_file(ctx.sema.db);
+    let adt = adt.value;
+    let adt = ast::Adt::cast(adt.clone())?;
+    let r = find_struct_impl(ctx, &adt, fn_name.text().as_str())?;
+    Some((r, file))
+}
+
 struct FunctionTemplate {
     insert_offset: TextSize,
     leading_ws: String,
@@ -130,12 +218,12 @@ impl FunctionBuilder {
                 file = in_file;
                 target
             }
-            None => next_space_for_fn_after_call_site(call)?,
+            None => next_space_for_fn_after_call_site(FuncExpr::Func(call.clone()))?,
         };
         let needs_pub = target_module.is_some();
         let target_module = target_module.or_else(|| ctx.sema.scope(target.syntax()).module())?;
         let fn_name = fn_name(path)?;
-        let (type_params, params) = fn_args(ctx, target_module, call)?;
+        let (type_params, params) = fn_args(ctx, target_module, FuncExpr::Func(call.clone()))?;
 
         let await_expr = call.syntax().parent().and_then(ast::AwaitExpr::cast);
         let is_async = await_expr.is_some();
@@ -181,6 +269,79 @@ impl FunctionBuilder {
         })
     }
 
+    fn from_method_call(
+        ctx: &AssistContext,
+        call: &ast::MethodCallExpr,
+        name: &ast::NameRef,
+        impl_: &Option<ast::Impl>,
+        file: FileId,
+        target_module: Module,
+        current_module: Module,
+    ) -> Option<Self> {
+        // let mut file = ctx.frange.file_id;
+        // let target_module = ctx.sema.scope(call.syntax()).module()?;
+        let target = match impl_ {
+            Some(impl_) => next_space_for_fn_in_impl(&impl_)?,
+            None => {
+                next_space_for_fn_in_module(
+                    ctx.sema.db,
+                    &target_module.definition_source(ctx.sema.db),
+                )?
+                .1
+            }
+        };
+        let needs_pub = !module_is_descendant(&current_module, &target_module, ctx);
+
+        let fn_name = make::name(&name.text());
+        let (type_params, params) = fn_args(ctx, target_module, FuncExpr::Method(call.clone()))?;
+
+        let await_expr = call.syntax().parent().and_then(ast::AwaitExpr::cast);
+        let is_async = await_expr.is_some();
+
+        // should_render_snippet intends to express a rough level of confidence about
+        // the correctness of the return type.
+        //
+        // If we are able to infer some return type, and that return type is not unit, we
+        // don't want to render the snippet. The assumption here is in this situation the
+        // return type is just as likely to be correct as any other part of the generated
+        // function.
+        //
+        // In the case where the return type is inferred as unit it is likely that the
+        // user does in fact intend for this generated function to return some non unit
+        // type, but that the current state of their code doesn't allow that return type
+        // to be accurately inferred.
+        let (ret_ty, should_render_snippet) = {
+            match ctx
+                .sema
+                .type_of_expr(&ast::Expr::MethodCallExpr(call.clone()))
+                .map(TypeInfo::original)
+            {
+                Some(ty) if ty.is_unknown() || ty.is_unit() => (make::ty_unit(), true),
+                Some(ty) => {
+                    let rendered = ty.display_source_code(ctx.db(), target_module.into());
+                    match rendered {
+                        Ok(rendered) => (make::ty(&rendered), false),
+                        Err(_) => (make::ty_unit(), true),
+                    }
+                }
+                None => (make::ty_unit(), true),
+            }
+        };
+        let ret_type = make::ret_type(ret_ty);
+
+        Some(Self {
+            target,
+            fn_name,
+            type_params,
+            params,
+            ret_type,
+            should_render_snippet,
+            file,
+            needs_pub,
+            is_async,
+        })
+    }
+
     fn render(self) -> FunctionTemplate {
         let placeholder_expr = make::ext::expr_todo();
         let fn_body = make::block_expr(vec![], Some(placeholder_expr));
@@ -249,7 +410,7 @@ fn fn_name(call: &ast::Path) -> Option<ast::Name> {
 fn fn_args(
     ctx: &AssistContext,
     target_module: hir::Module,
-    call: &ast::CallExpr,
+    call: FuncExpr,
 ) -> Option<(Option<ast::GenericParamList>, ast::ParamList)> {
     let mut arg_names = Vec::new();
     let mut arg_types = Vec::new();
@@ -277,7 +438,17 @@ fn fn_args(
     let params = arg_names.into_iter().zip(arg_types).map(|(name, ty)| {
         make::param(make::ext::simple_ident_pat(make::name(&name)).into(), make::ty(&ty))
     });
-    Some((None, make::param_list(None, params)))
+
+    Some((
+        None,
+        make::param_list(
+            match call {
+                FuncExpr::Func(_) => None,
+                FuncExpr::Method(_) => Some(make::self_param()),
+            },
+            params,
+        ),
+    ))
 }
 
 /// Makes duplicate argument names unique by appending incrementing numbers.
@@ -348,7 +519,7 @@ fn fn_arg_type(
 /// directly after the current block
 /// We want to write the generated function directly after
 /// fns, impls or macro calls, but inside mods
-fn next_space_for_fn_after_call_site(expr: &ast::CallExpr) -> Option<GeneratedFunctionTarget> {
+fn next_space_for_fn_after_call_site(expr: FuncExpr) -> Option<GeneratedFunctionTarget> {
     let mut ancestors = expr.syntax().ancestors().peekable();
     let mut last_ancestor: Option<SyntaxNode> = None;
     while let Some(next_ancestor) = ancestors.next() {
@@ -401,6 +572,26 @@ fn next_space_for_fn_in_module(
     Some((file, assist_item))
 }
 
+fn next_space_for_fn_in_impl(impl_: &ast::Impl) -> Option<GeneratedFunctionTarget> {
+    if let Some(last_item) = impl_.assoc_item_list().and_then(|it| it.assoc_items().last()) {
+        Some(GeneratedFunctionTarget::BehindItem(last_item.syntax().clone()))
+    } else {
+        Some(GeneratedFunctionTarget::InEmptyItemList(impl_.assoc_item_list()?.syntax().clone()))
+    }
+}
+
+fn module_is_descendant(module: &hir::Module, ans: &hir::Module, ctx: &AssistContext) -> bool {
+    if module == ans {
+        return true;
+    }
+    for c in ans.children(ctx.sema.db) {
+        if module_is_descendant(module, &c, ctx) {
+            return true;
+        }
+    }
+    false
+}
+
 #[cfg(test)]
 mod tests {
     use crate::tests::{check_assist, check_assist_not_applicable};
@@ -1164,8 +1355,7 @@ fn bar(baz: ()) {}
 
     #[test]
     fn create_method_with_no_args() {
-        // FIXME: This is wrong, this should just work.
-        check_assist_not_applicable(
+        check_assist(
             generate_function,
             r#"
 struct Foo;
@@ -1174,6 +1364,18 @@ impl Foo {
         self.bar()$0;
     }
 }
+"#,
+            r#"
+struct Foo;
+impl Foo {
+    fn foo(&self) {
+        self.bar();
+    }
+
+    fn bar(&self) ${0:-> ()} {
+        todo!()
+    }
+}
 "#,
         )
     }
@@ -1195,6 +1397,131 @@ fn foo() {
 async fn bar(arg: i32) ${0:-> ()} {
     todo!()
 }
+",
+        )
+    }
+
+    #[test]
+    fn create_method() {
+        check_assist(
+            generate_function,
+            r"
+struct S;
+fn foo() {S.bar$0();}
+",
+            r"
+struct S;
+fn foo() {S.bar();}
+impl S {
+
+
+fn bar(&self) ${0:-> ()} {
+    todo!()
+}
+}
+",
+        )
+    }
+
+    #[test]
+    fn create_method_within_an_impl() {
+        check_assist(
+            generate_function,
+            r"
+struct S;
+fn foo() {S.bar$0();}
+impl S {}
+
+",
+            r"
+struct S;
+fn foo() {S.bar();}
+impl S {
+    fn bar(&self) ${0:-> ()} {
+        todo!()
+    }
+}
+
+",
+        )
+    }
+
+    #[test]
+    fn create_method_from_different_module() {
+        check_assist(
+            generate_function,
+            r"
+mod s {
+    pub struct S;
+}
+fn foo() {s::S.bar$0();}
+",
+            r"
+mod s {
+    pub struct S;
+impl S {
+
+
+    pub(crate) fn bar(&self) ${0:-> ()} {
+        todo!()
+    }
+}
+}
+fn foo() {s::S.bar();}
+",
+        )
+    }
+
+    #[test]
+    fn create_method_from_descendant_module() {
+        check_assist(
+            generate_function,
+            r"
+struct S;
+mod s {
+    fn foo() {
+        super::S.bar$0();
+    }
+}
+
+",
+            r"
+struct S;
+mod s {
+    fn foo() {
+        super::S.bar();
+    }
+}
+impl S {
+
+
+fn bar(&self) ${0:-> ()} {
+    todo!()
+}
+}
+
+",
+        )
+    }
+
+    #[test]
+    fn create_method_with_cursor_anywhere_on_call_expresion() {
+        check_assist(
+            generate_function,
+            r"
+struct S;
+fn foo() {$0S.bar();}
+",
+            r"
+struct S;
+fn foo() {S.bar();}
+impl S {
+
+
+fn bar(&self) ${0:-> ()} {
+    todo!()
+}
+}
 ",
         )
     }

crates/syntax/src/ast/make.rs

@@ -531,6 +531,10 @@ pub fn param(pat: ast::Pat, ty: ast::Type) -> ast::Param {
     ast_from_text(&format!("fn f({}: {}) {{ }}", pat, ty))
 }
 
+pub fn self_param() -> ast::SelfParam {
+    ast_from_text(&format!("fn f(&self) {{ }}"))
+}
+
 pub fn ret_type(ty: ast::Type) -> ast::RetType {
     ast_from_text(&format!("fn f() -> {} {{ }}", ty))
 }