diff options
Diffstat (limited to 'crates/ide_assists/src/handlers')
-rw-r--r-- | crates/ide_assists/src/handlers/generate_function.rs | 46 |
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 | ||
228 | fn 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 | ||
549 | fn bar(baz: Baz) ${0:-> ()} { | 568 | fn 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" | ||
1086 | fn main() { | ||
1087 | let x: u32 = foo$0(); | ||
1088 | } | ||
1089 | ", | ||
1090 | r" | ||
1091 | fn main() { | ||
1092 | let x: u32 = foo(); | ||
1093 | } | ||
1094 | |||
1095 | fn 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, |