aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--crates/ide_assists/src/handlers/generate_function.rs46
1 files changed, 43 insertions, 3 deletions
diff --git a/crates/ide_assists/src/handlers/generate_function.rs b/crates/ide_assists/src/handlers/generate_function.rs
index 3870b7e75..fd4f2fbed 100644
--- a/crates/ide_assists/src/handlers/generate_function.rs
+++ b/crates/ide_assists/src/handlers/generate_function.rs
@@ -104,6 +104,7 @@ struct FunctionBuilder {
104 fn_name: ast::Name, 104 fn_name: ast::Name,
105 type_params: Option<ast::GenericParamList>, 105 type_params: Option<ast::GenericParamList>,
106 params: ast::ParamList, 106 params: ast::ParamList,
107 ret_type: Option<ast::RetType>,
107 file: FileId, 108 file: FileId,
108 needs_pub: bool, 109 needs_pub: bool,
109} 110}
@@ -131,8 +132,9 @@ impl FunctionBuilder {
131 let target_module = target_module.or_else(|| ctx.sema.scope(target.syntax()).module())?; 132 let target_module = target_module.or_else(|| ctx.sema.scope(target.syntax()).module())?;
132 let fn_name = fn_name(&path)?; 133 let fn_name = fn_name(&path)?;
133 let (type_params, params) = fn_args(ctx, target_module, &call)?; 134 let (type_params, params) = fn_args(ctx, target_module, &call)?;
135 let ret_type = fn_ret_type(ctx, target_module, &call);
134 136
135 Some(Self { target, fn_name, type_params, params, file, needs_pub }) 137 Some(Self { target, fn_name, type_params, params, ret_type, file, needs_pub })
136 } 138 }
137 139
138 fn render(self) -> FunctionTemplate { 140 fn render(self) -> FunctionTemplate {
@@ -145,7 +147,7 @@ impl FunctionBuilder {
145 self.type_params, 147 self.type_params,
146 self.params, 148 self.params,
147 fn_body, 149 fn_body,
148 Some(make::ret_type(make::ty_unit())), 150 Some(self.ret_type.unwrap_or_else(|| make::ret_type(make::ty_unit()))),
149 ); 151 );
150 let leading_ws; 152 let leading_ws;
151 let trailing_ws; 153 let trailing_ws;
@@ -223,6 +225,23 @@ fn fn_args(
223 Some((None, make::param_list(None, params))) 225 Some((None, make::param_list(None, params)))
224} 226}
225 227
228fn fn_ret_type(
229 ctx: &AssistContext,
230 target_module: hir::Module,
231 call: &ast::CallExpr,
232) -> Option<ast::RetType> {
233 let ty = ctx.sema.type_of_expr(&ast::Expr::CallExpr(call.clone()))?;
234 if ty.is_unknown() {
235 return None;
236 }
237
238 if let Ok(rendered) = ty.display_source_code(ctx.db(), target_module.into()) {
239 Some(make::ret_type(make::ty(&rendered)))
240 } else {
241 None
242 }
243}
244
226/// Makes duplicate argument names unique by appending incrementing numbers. 245/// Makes duplicate argument names unique by appending incrementing numbers.
227/// 246///
228/// ``` 247/// ```
@@ -546,7 +565,7 @@ impl Baz {
546 } 565 }
547} 566}
548 567
549fn bar(baz: Baz) ${0:-> ()} { 568fn bar(baz: Baz) ${0:-> Baz} {
550 todo!() 569 todo!()
551} 570}
552", 571",
@@ -1060,6 +1079,27 @@ pub(crate) fn bar() ${0:-> ()} {
1060 } 1079 }
1061 1080
1062 #[test] 1081 #[test]
1082 fn add_function_with_return_type() {
1083 check_assist(
1084 generate_function,
1085 r"
1086fn main() {
1087 let x: u32 = foo$0();
1088}
1089",
1090 r"
1091fn main() {
1092 let x: u32 = foo();
1093}
1094
1095fn foo() ${0:-> u32} {
1096 todo!()
1097}
1098",
1099 )
1100 }
1101
1102 #[test]
1063 fn add_function_not_applicable_if_function_already_exists() { 1103 fn add_function_not_applicable_if_function_already_exists() {
1064 check_assist_not_applicable( 1104 check_assist_not_applicable(
1065 generate_function, 1105 generate_function,