1 files changed, 119 insertions, 10 deletions

diff --git a/crates/ide_assists/src/handlers/generate_function.rs b/crates/ide_assists/src/handlers/generate_function.rs
index 959824981..6f95b1a07 100644
--- a/crates/ide_assists/src/handlers/generate_function.rs
+++ b/crates/ide_assists/src/handlers/generate_function.rs
@@ -1,6 +1,7 @@
 use hir::HirDisplay;
 use ide_db::{base_db::FileId, helpers::SnippetCap};
 use rustc_hash::{FxHashMap, FxHashSet};
+use stdx::to_lower_snake_case;
 use syntax::{
     ast::{
         self,
@@ -82,17 +83,18 @@ struct FunctionTemplate {
     leading_ws: String,
     fn_def: ast::Fn,
     ret_type: ast::RetType,
+    should_render_snippet: bool,
     trailing_ws: String,
     file: FileId,
 }
 
 impl FunctionTemplate {
     fn to_string(&self, cap: Option<SnippetCap>) -> String {
-        let f = match cap {
-            Some(cap) => {
+        let f = match (cap, self.should_render_snippet) {
+            (Some(cap), true) => {
                 render_snippet(cap, self.fn_def.syntax(), Cursor::Replace(self.ret_type.syntax()))
             }
-            None => self.fn_def.to_string(),
+            _ => self.fn_def.to_string(),
         };
         format!("{}{}{}", self.leading_ws, f, self.trailing_ws)
     }
@@ -103,6 +105,8 @@ struct FunctionBuilder {
     fn_name: ast::Name,
     type_params: Option<ast::GenericParamList>,
     params: ast::ParamList,
+    ret_type: ast::RetType,
+    should_render_snippet: bool,
     file: FileId,
     needs_pub: bool,
 }
@@ -131,7 +135,43 @@ impl FunctionBuilder {
         let fn_name = fn_name(&path)?;
         let (type_params, params) = fn_args(ctx, target_module, &call)?;
 
-        Some(Self { target, fn_name, type_params, params, file, needs_pub })
+        // 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::CallExpr(call.clone())) {
+                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,
+        })
     }
 
     fn render(self) -> FunctionTemplate {
@@ -144,7 +184,7 @@ impl FunctionBuilder {
             self.type_params,
             self.params,
             fn_body,
-            Some(make::ret_type(make::ty_unit())),
+            Some(self.ret_type),
         );
         let leading_ws;
         let trailing_ws;
@@ -170,6 +210,7 @@ impl FunctionBuilder {
             insert_offset,
             leading_ws,
             ret_type: fn_def.ret_type().unwrap(),
+            should_render_snippet: self.should_render_snippet,
             fn_def,
             trailing_ws,
             file: self.file,
@@ -257,14 +298,15 @@ fn deduplicate_arg_names(arg_names: &mut Vec<String>) {
 fn fn_arg_name(fn_arg: &ast::Expr) -> Option<String> {
     match fn_arg {
         ast::Expr::CastExpr(cast_expr) => fn_arg_name(&cast_expr.expr()?),
-        _ => Some(
-            fn_arg
+        _ => {
+            let s = fn_arg
                 .syntax()
                 .descendants()
                 .filter(|d| ast::NameRef::can_cast(d.kind()))
                 .last()?
-                .to_string(),
-        ),
+                .to_string();
+            Some(to_lower_snake_case(&s))
+        }
     }
 }
 
@@ -448,6 +490,52 @@ mod baz {
     }
 
     #[test]
+    fn add_function_with_upper_camel_case_arg() {
+        check_assist(
+            generate_function,
+            r"
+struct BazBaz;
+fn foo() {
+    bar$0(BazBaz);
+}
+",
+            r"
+struct BazBaz;
+fn foo() {
+    bar(BazBaz);
+}
+
+fn bar(baz_baz: BazBaz) ${0:-> ()} {
+    todo!()
+}
+",
+        );
+    }
+
+    #[test]
+    fn add_function_with_upper_camel_case_arg_as_cast() {
+        check_assist(
+            generate_function,
+            r"
+struct BazBaz;
+fn foo() {
+    bar$0(&BazBaz as *const BazBaz);
+}
+",
+            r"
+struct BazBaz;
+fn foo() {
+    bar(&BazBaz as *const BazBaz);
+}
+
+fn bar(baz_baz: *const BazBaz) ${0:-> ()} {
+    todo!()
+}
+",
+        );
+    }
+
+    #[test]
     fn add_function_with_function_call_arg() {
         check_assist(
             generate_function,
@@ -498,7 +586,7 @@ impl Baz {
     }
 }
 
-fn bar(baz: Baz) ${0:-> ()} {
+fn bar(baz: Baz) -> Baz {
     todo!()
 }
 ",
@@ -1012,6 +1100,27 @@ pub(crate) fn bar() ${0:-> ()} {
     }
 
     #[test]
+    fn add_function_with_return_type() {
+        check_assist(
+            generate_function,
+            r"
+fn main() {
+    let x: u32 = foo$0();
+}
+",
+            r"
+fn main() {
+    let x: u32 = foo();
+}
+
+fn foo() -> u32 {
+    todo!()
+}
+",
+        )
+    }
+
+    #[test]
     fn add_function_not_applicable_if_function_already_exists() {
         check_assist_not_applicable(
             generate_function,