diff options
Diffstat (limited to 'crates/ide_assists/src/handlers/infer_function_return_type.rs')
-rw-r--r-- | crates/ide_assists/src/handlers/infer_function_return_type.rs | 345 |
1 files changed, 345 insertions, 0 deletions
diff --git a/crates/ide_assists/src/handlers/infer_function_return_type.rs b/crates/ide_assists/src/handlers/infer_function_return_type.rs new file mode 100644 index 000000000..5279af1f3 --- /dev/null +++ b/crates/ide_assists/src/handlers/infer_function_return_type.rs | |||
@@ -0,0 +1,345 @@ | |||
1 | use hir::HirDisplay; | ||
2 | use syntax::{ast, AstNode, TextRange, TextSize}; | ||
3 | use test_utils::mark; | ||
4 | |||
5 | use crate::{AssistContext, AssistId, AssistKind, Assists}; | ||
6 | |||
7 | // Assist: infer_function_return_type | ||
8 | // | ||
9 | // Adds the return type to a function or closure inferred from its tail expression if it doesn't have a return | ||
10 | // type specified. This assists is useable in a functions or closures tail expression or return type position. | ||
11 | // | ||
12 | // ``` | ||
13 | // fn foo() { 4$02i32 } | ||
14 | // ``` | ||
15 | // -> | ||
16 | // ``` | ||
17 | // fn foo() -> i32 { 42i32 } | ||
18 | // ``` | ||
19 | pub(crate) fn infer_function_return_type(acc: &mut Assists, ctx: &AssistContext) -> Option<()> { | ||
20 | let (fn_type, tail_expr, builder_edit_pos) = extract_tail(ctx)?; | ||
21 | let module = ctx.sema.scope(tail_expr.syntax()).module()?; | ||
22 | let ty = ctx.sema.type_of_expr(&tail_expr)?; | ||
23 | if ty.is_unit() { | ||
24 | return None; | ||
25 | } | ||
26 | let ty = ty.display_source_code(ctx.db(), module.into()).ok()?; | ||
27 | |||
28 | acc.add( | ||
29 | AssistId("infer_function_return_type", AssistKind::RefactorRewrite), | ||
30 | match fn_type { | ||
31 | FnType::Function => "Add this function's return type", | ||
32 | FnType::Closure { .. } => "Add this closure's return type", | ||
33 | }, | ||
34 | tail_expr.syntax().text_range(), | ||
35 | |builder| { | ||
36 | match builder_edit_pos { | ||
37 | InsertOrReplace::Insert(insert_pos) => { | ||
38 | builder.insert(insert_pos, &format!("-> {} ", ty)) | ||
39 | } | ||
40 | InsertOrReplace::Replace(text_range) => { | ||
41 | builder.replace(text_range, &format!("-> {}", ty)) | ||
42 | } | ||
43 | } | ||
44 | if let FnType::Closure { wrap_expr: true } = fn_type { | ||
45 | mark::hit!(wrap_closure_non_block_expr); | ||
46 | // `|x| x` becomes `|x| -> T x` which is invalid, so wrap it in a block | ||
47 | builder.replace(tail_expr.syntax().text_range(), &format!("{{{}}}", tail_expr)); | ||
48 | } | ||
49 | }, | ||
50 | ) | ||
51 | } | ||
52 | |||
53 | enum InsertOrReplace { | ||
54 | Insert(TextSize), | ||
55 | Replace(TextRange), | ||
56 | } | ||
57 | |||
58 | /// Check the potentially already specified return type and reject it or turn it into a builder command | ||
59 | /// if allowed. | ||
60 | fn ret_ty_to_action(ret_ty: Option<ast::RetType>, insert_pos: TextSize) -> Option<InsertOrReplace> { | ||
61 | match ret_ty { | ||
62 | Some(ret_ty) => match ret_ty.ty() { | ||
63 | Some(ast::Type::InferType(_)) | None => { | ||
64 | mark::hit!(existing_infer_ret_type); | ||
65 | mark::hit!(existing_infer_ret_type_closure); | ||
66 | Some(InsertOrReplace::Replace(ret_ty.syntax().text_range())) | ||
67 | } | ||
68 | _ => { | ||
69 | mark::hit!(existing_ret_type); | ||
70 | mark::hit!(existing_ret_type_closure); | ||
71 | None | ||
72 | } | ||
73 | }, | ||
74 | None => Some(InsertOrReplace::Insert(insert_pos + TextSize::from(1))), | ||
75 | } | ||
76 | } | ||
77 | |||
78 | enum FnType { | ||
79 | Function, | ||
80 | Closure { wrap_expr: bool }, | ||
81 | } | ||
82 | |||
83 | fn extract_tail(ctx: &AssistContext) -> Option<(FnType, ast::Expr, InsertOrReplace)> { | ||
84 | let (fn_type, tail_expr, return_type_range, action) = | ||
85 | if let Some(closure) = ctx.find_node_at_offset::<ast::ClosureExpr>() { | ||
86 | let rpipe_pos = closure.param_list()?.syntax().last_token()?.text_range().end(); | ||
87 | let action = ret_ty_to_action(closure.ret_type(), rpipe_pos)?; | ||
88 | |||
89 | let body = closure.body()?; | ||
90 | let body_start = body.syntax().first_token()?.text_range().start(); | ||
91 | let (tail_expr, wrap_expr) = match body { | ||
92 | ast::Expr::BlockExpr(block) => (block.tail_expr()?, false), | ||
93 | body => (body, true), | ||
94 | }; | ||
95 | |||
96 | let ret_range = TextRange::new(rpipe_pos, body_start); | ||
97 | (FnType::Closure { wrap_expr }, tail_expr, ret_range, action) | ||
98 | } else { | ||
99 | let func = ctx.find_node_at_offset::<ast::Fn>()?; | ||
100 | let rparen_pos = func.param_list()?.r_paren_token()?.text_range().end(); | ||
101 | let action = ret_ty_to_action(func.ret_type(), rparen_pos)?; | ||
102 | |||
103 | let body = func.body()?; | ||
104 | let tail_expr = body.tail_expr()?; | ||
105 | |||
106 | let ret_range_end = body.l_curly_token()?.text_range().start(); | ||
107 | let ret_range = TextRange::new(rparen_pos, ret_range_end); | ||
108 | (FnType::Function, tail_expr, ret_range, action) | ||
109 | }; | ||
110 | let frange = ctx.frange.range; | ||
111 | if return_type_range.contains_range(frange) { | ||
112 | mark::hit!(cursor_in_ret_position); | ||
113 | mark::hit!(cursor_in_ret_position_closure); | ||
114 | } else if tail_expr.syntax().text_range().contains_range(frange) { | ||
115 | mark::hit!(cursor_on_tail); | ||
116 | mark::hit!(cursor_on_tail_closure); | ||
117 | } else { | ||
118 | return None; | ||
119 | } | ||
120 | Some((fn_type, tail_expr, action)) | ||
121 | } | ||
122 | |||
123 | #[cfg(test)] | ||
124 | mod tests { | ||
125 | use crate::tests::{check_assist, check_assist_not_applicable}; | ||
126 | |||
127 | use super::*; | ||
128 | |||
129 | #[test] | ||
130 | fn infer_return_type_specified_inferred() { | ||
131 | mark::check!(existing_infer_ret_type); | ||
132 | check_assist( | ||
133 | infer_function_return_type, | ||
134 | r#"fn foo() -> $0_ { | ||
135 | 45 | ||
136 | }"#, | ||
137 | r#"fn foo() -> i32 { | ||
138 | 45 | ||
139 | }"#, | ||
140 | ); | ||
141 | } | ||
142 | |||
143 | #[test] | ||
144 | fn infer_return_type_specified_inferred_closure() { | ||
145 | mark::check!(existing_infer_ret_type_closure); | ||
146 | check_assist( | ||
147 | infer_function_return_type, | ||
148 | r#"fn foo() { | ||
149 | || -> _ {$045}; | ||
150 | }"#, | ||
151 | r#"fn foo() { | ||
152 | || -> i32 {45}; | ||
153 | }"#, | ||
154 | ); | ||
155 | } | ||
156 | |||
157 | #[test] | ||
158 | fn infer_return_type_cursor_at_return_type_pos() { | ||
159 | mark::check!(cursor_in_ret_position); | ||
160 | check_assist( | ||
161 | infer_function_return_type, | ||
162 | r#"fn foo() $0{ | ||
163 | 45 | ||
164 | }"#, | ||
165 | r#"fn foo() -> i32 { | ||
166 | 45 | ||
167 | }"#, | ||
168 | ); | ||
169 | } | ||
170 | |||
171 | #[test] | ||
172 | fn infer_return_type_cursor_at_return_type_pos_closure() { | ||
173 | mark::check!(cursor_in_ret_position_closure); | ||
174 | check_assist( | ||
175 | infer_function_return_type, | ||
176 | r#"fn foo() { | ||
177 | || $045 | ||
178 | }"#, | ||
179 | r#"fn foo() { | ||
180 | || -> i32 {45} | ||
181 | }"#, | ||
182 | ); | ||
183 | } | ||
184 | |||
185 | #[test] | ||
186 | fn infer_return_type() { | ||
187 | mark::check!(cursor_on_tail); | ||
188 | check_assist( | ||
189 | infer_function_return_type, | ||
190 | r#"fn foo() { | ||
191 | 45$0 | ||
192 | }"#, | ||
193 | r#"fn foo() -> i32 { | ||
194 | 45 | ||
195 | }"#, | ||
196 | ); | ||
197 | } | ||
198 | |||
199 | #[test] | ||
200 | fn infer_return_type_nested() { | ||
201 | check_assist( | ||
202 | infer_function_return_type, | ||
203 | r#"fn foo() { | ||
204 | if true { | ||
205 | 3$0 | ||
206 | } else { | ||
207 | 5 | ||
208 | } | ||
209 | }"#, | ||
210 | r#"fn foo() -> i32 { | ||
211 | if true { | ||
212 | 3 | ||
213 | } else { | ||
214 | 5 | ||
215 | } | ||
216 | }"#, | ||
217 | ); | ||
218 | } | ||
219 | |||
220 | #[test] | ||
221 | fn not_applicable_ret_type_specified() { | ||
222 | mark::check!(existing_ret_type); | ||
223 | check_assist_not_applicable( | ||
224 | infer_function_return_type, | ||
225 | r#"fn foo() -> i32 { | ||
226 | ( 45$0 + 32 ) * 123 | ||
227 | }"#, | ||
228 | ); | ||
229 | } | ||
230 | |||
231 | #[test] | ||
232 | fn not_applicable_non_tail_expr() { | ||
233 | check_assist_not_applicable( | ||
234 | infer_function_return_type, | ||
235 | r#"fn foo() { | ||
236 | let x = $03; | ||
237 | ( 45 + 32 ) * 123 | ||
238 | }"#, | ||
239 | ); | ||
240 | } | ||
241 | |||
242 | #[test] | ||
243 | fn not_applicable_unit_return_type() { | ||
244 | check_assist_not_applicable( | ||
245 | infer_function_return_type, | ||
246 | r#"fn foo() { | ||
247 | ($0) | ||
248 | }"#, | ||
249 | ); | ||
250 | } | ||
251 | |||
252 | #[test] | ||
253 | fn infer_return_type_closure_block() { | ||
254 | mark::check!(cursor_on_tail_closure); | ||
255 | check_assist( | ||
256 | infer_function_return_type, | ||
257 | r#"fn foo() { | ||
258 | |x: i32| { | ||
259 | x$0 | ||
260 | }; | ||
261 | }"#, | ||
262 | r#"fn foo() { | ||
263 | |x: i32| -> i32 { | ||
264 | x | ||
265 | }; | ||
266 | }"#, | ||
267 | ); | ||
268 | } | ||
269 | |||
270 | #[test] | ||
271 | fn infer_return_type_closure() { | ||
272 | check_assist( | ||
273 | infer_function_return_type, | ||
274 | r#"fn foo() { | ||
275 | |x: i32| { x$0 }; | ||
276 | }"#, | ||
277 | r#"fn foo() { | ||
278 | |x: i32| -> i32 { x }; | ||
279 | }"#, | ||
280 | ); | ||
281 | } | ||
282 | |||
283 | #[test] | ||
284 | fn infer_return_type_closure_wrap() { | ||
285 | mark::check!(wrap_closure_non_block_expr); | ||
286 | check_assist( | ||
287 | infer_function_return_type, | ||
288 | r#"fn foo() { | ||
289 | |x: i32| x$0; | ||
290 | }"#, | ||
291 | r#"fn foo() { | ||
292 | |x: i32| -> i32 {x}; | ||
293 | }"#, | ||
294 | ); | ||
295 | } | ||
296 | |||
297 | #[test] | ||
298 | fn infer_return_type_nested_closure() { | ||
299 | check_assist( | ||
300 | infer_function_return_type, | ||
301 | r#"fn foo() { | ||
302 | || { | ||
303 | if true { | ||
304 | 3$0 | ||
305 | } else { | ||
306 | 5 | ||
307 | } | ||
308 | } | ||
309 | }"#, | ||
310 | r#"fn foo() { | ||
311 | || -> i32 { | ||
312 | if true { | ||
313 | 3 | ||
314 | } else { | ||
315 | 5 | ||
316 | } | ||
317 | } | ||
318 | }"#, | ||
319 | ); | ||
320 | } | ||
321 | |||
322 | #[test] | ||
323 | fn not_applicable_ret_type_specified_closure() { | ||
324 | mark::check!(existing_ret_type_closure); | ||
325 | check_assist_not_applicable( | ||
326 | infer_function_return_type, | ||
327 | r#"fn foo() { | ||
328 | || -> i32 { 3$0 } | ||
329 | }"#, | ||
330 | ); | ||
331 | } | ||
332 | |||
333 | #[test] | ||
334 | fn not_applicable_non_tail_expr_closure() { | ||
335 | check_assist_not_applicable( | ||
336 | infer_function_return_type, | ||
337 | r#"fn foo() { | ||
338 | || -> i32 { | ||
339 | let x = 3$0; | ||
340 | 6 | ||
341 | } | ||
342 | }"#, | ||
343 | ); | ||
344 | } | ||
345 | } | ||