Merge #7898

7898: generate_function assist: infer return type r=JoshMcguigan a=JoshMcguigan

This PR makes two changes to the generate function assist:

1. Attempt to infer an appropriate return type for the generated function
2. If a return type is inferred, and that return type is not unit, don't render the snippet

```rust
fn main() {
    let x: u32 = foo$0();
    //              ^^^ trigger the assist to generate this function
}

// BEFORE
fn foo() ${0:-> ()} {
    todo!()
}

// AFTER (only change 1)
fn foo() ${0:-> u32} {
    todo!()
}

// AFTER (change  1 and 2, note the lack of snippet around the return type)
fn foo() -> u32 {
    todo!()
}
```

These changes are made as two commits, in case we want to omit change 2. I personally feel like it is a nice change, but I could understand there being some opposition.

#### Pros of change 2
If we are able to infer a return type, and especially if that return type is not the unit type, the return type is almost as likely to be correct as the argument names/types. I think this becomes even more true as people learn how this feature works.

#### Cons of change 2

We could never be as confident about the return type as we are about the function argument types, so it is more likely a user will want to change that. Plus it is a confusing UX to sometimes have the cursor highlight the return type after triggering this assist and sometimes not have that happen.

#### Why omit unit type?

The assumption is that if we infer the return type as unit, it is likely just because of the current structure of the code rather than that actually being the desired return type. However, this is obviously just a heuristic and will sometimes be wrong. But being wrong here just means falling back to the exact behavior that existed before this PR.



Co-authored-by: Josh Mcguigan <[email protected]>

1 files changed, 67 insertions, 6 deletions

diff --git a/crates/ide_assists/src/handlers/generate_function.rs b/crates/ide_assists/src/handlers/generate_function.rs
index 3870b7e75..6f95b1a07 100644
--- a/crates/ide_assists/src/handlers/generate_function.rs
+++ b/crates/ide_assists/src/handlers/generate_function.rs
@@ -83,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)
     }
@@ -104,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,
 }
@@ -132,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 {
@@ -145,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;
@@ -171,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,
@@ -546,7 +586,7 @@ impl Baz {
     }
 }
 
-fn bar(baz: Baz) ${0:-> ()} {
+fn bar(baz: Baz) -> Baz {
     todo!()
 }
 ",
@@ -1060,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,