ra_ap_ide_assists/handlers/
add_turbo_fish.rs

1use either::Either;
2use ide_db::defs::{Definition, NameRefClass};
3use syntax::{
4    AstNode,
5    ast::{self, HasArgList, HasGenericArgs, make, syntax_factory::SyntaxFactory},
6    syntax_editor::Position,
7};
8
9use crate::{
10    AssistId,
11    assist_context::{AssistContext, Assists},
12};
13
14// Assist: add_turbo_fish
15//
16// Adds `::<_>` to a call of a generic method or function.
17//
18// ```
19// fn make<T>() -> T { todo!() }
20// fn main() {
21//     let x = make$0();
22// }
23// ```
24// ->
25// ```
26// fn make<T>() -> T { todo!() }
27// fn main() {
28//     let x = make::<${0:_}>();
29// }
30// ```
31pub(crate) fn add_turbo_fish(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
32    let turbofish_target =
33        ctx.find_node_at_offset::<ast::PathSegment>().map(Either::Left).or_else(|| {
34            let callable_expr = ctx.find_node_at_offset::<ast::CallableExpr>()?;
35
36            if callable_expr.arg_list()?.args().next().is_some() {
37                return None;
38            }
39
40            cov_mark::hit!(add_turbo_fish_after_call);
41            cov_mark::hit!(add_type_ascription_after_call);
42
43            match callable_expr {
44                ast::CallableExpr::Call(it) => {
45                    let ast::Expr::PathExpr(path) = it.expr()? else {
46                        return None;
47                    };
48
49                    Some(Either::Left(path.path()?.segment()?))
50                }
51                ast::CallableExpr::MethodCall(it) => Some(Either::Right(it)),
52            }
53        })?;
54
55    let already_has_turbofish = match &turbofish_target {
56        Either::Left(path_segment) => path_segment.generic_arg_list().is_some(),
57        Either::Right(method_call) => method_call.generic_arg_list().is_some(),
58    };
59
60    if already_has_turbofish {
61        cov_mark::hit!(add_turbo_fish_one_fish_is_enough);
62        return None;
63    }
64
65    let name_ref = match &turbofish_target {
66        Either::Left(path_segment) => path_segment.name_ref()?,
67        Either::Right(method_call) => method_call.name_ref()?,
68    };
69    let ident = name_ref.ident_token()?;
70
71    let def = match NameRefClass::classify(&ctx.sema, &name_ref)? {
72        NameRefClass::Definition(def, _) => def,
73        NameRefClass::FieldShorthand { .. } | NameRefClass::ExternCrateShorthand { .. } => {
74            return None;
75        }
76    };
77    let fun = match def {
78        Definition::Function(it) => it,
79        _ => return None,
80    };
81    let generics = hir::GenericDef::Function(fun).params(ctx.sema.db);
82    if generics.is_empty() {
83        cov_mark::hit!(add_turbo_fish_non_generic);
84        return None;
85    }
86
87    if let Some(let_stmt) = ctx.find_node_at_offset::<ast::LetStmt>() {
88        if let_stmt.colon_token().is_none() {
89            let_stmt.pat()?;
90
91            acc.add(
92                AssistId::refactor_rewrite("add_type_ascription"),
93                "Add `: _` before assignment operator",
94                ident.text_range(),
95                |builder| {
96                    let mut editor = builder.make_editor(let_stmt.syntax());
97
98                    if let_stmt.semicolon_token().is_none() {
99                        editor.insert(
100                            Position::last_child_of(let_stmt.syntax()),
101                            make::tokens::semicolon(),
102                        );
103                    }
104
105                    let placeholder_ty = make::ty_placeholder().clone_for_update();
106
107                    if let Some(pat) = let_stmt.pat() {
108                        let elements = vec![
109                            make::token(syntax::SyntaxKind::COLON).into(),
110                            make::token(syntax::SyntaxKind::WHITESPACE).into(),
111                            placeholder_ty.syntax().clone().into(),
112                        ];
113                        editor.insert_all(Position::after(pat.syntax()), elements);
114                        if let Some(cap) = ctx.config.snippet_cap {
115                            editor.add_annotation(
116                                placeholder_ty.syntax(),
117                                builder.make_placeholder_snippet(cap),
118                            );
119                        }
120                    }
121
122                    builder.add_file_edits(ctx.vfs_file_id(), editor);
123                },
124            )?
125        } else {
126            cov_mark::hit!(add_type_ascription_already_typed);
127        }
128    }
129
130    let number_of_arguments = generics
131        .iter()
132        .filter(|param| {
133            matches!(param, hir::GenericParam::TypeParam(_) | hir::GenericParam::ConstParam(_))
134        })
135        .count();
136
137    acc.add(
138        AssistId::refactor_rewrite("add_turbo_fish"),
139        "Add `::<>`",
140        ident.text_range(),
141        |builder| {
142            builder.trigger_parameter_hints();
143
144            let make = SyntaxFactory::with_mappings();
145            let mut editor = match &turbofish_target {
146                Either::Left(it) => builder.make_editor(it.syntax()),
147                Either::Right(it) => builder.make_editor(it.syntax()),
148            };
149
150            let fish_head = get_fish_head(&make, number_of_arguments);
151
152            match turbofish_target {
153                Either::Left(path_segment) => {
154                    if let Some(generic_arg_list) = path_segment.generic_arg_list() {
155                        editor.replace(generic_arg_list.syntax(), fish_head.syntax());
156                    } else {
157                        editor.insert(
158                            Position::last_child_of(path_segment.syntax()),
159                            fish_head.syntax(),
160                        );
161                    }
162                }
163                Either::Right(method_call) => {
164                    if let Some(generic_arg_list) = method_call.generic_arg_list() {
165                        editor.replace(generic_arg_list.syntax(), fish_head.syntax());
166                    } else {
167                        let position = if let Some(arg_list) = method_call.arg_list() {
168                            Position::before(arg_list.syntax())
169                        } else {
170                            Position::last_child_of(method_call.syntax())
171                        };
172                        editor.insert(position, fish_head.syntax());
173                    }
174                }
175            };
176
177            if let Some(cap) = ctx.config.snippet_cap {
178                for arg in fish_head.generic_args() {
179                    editor.add_annotation(arg.syntax(), builder.make_placeholder_snippet(cap));
180                }
181            }
182
183            editor.add_mappings(make.finish_with_mappings());
184            builder.add_file_edits(ctx.vfs_file_id(), editor);
185        },
186    )
187}
188
189/// This will create a turbofish generic arg list corresponding to the number of arguments
190fn get_fish_head(make: &SyntaxFactory, number_of_arguments: usize) -> ast::GenericArgList {
191    let args = (0..number_of_arguments).map(|_| make::type_arg(make::ty_placeholder()).into());
192    make.generic_arg_list(args, true)
193}
194
195#[cfg(test)]
196mod tests {
197    use crate::tests::{
198        check_assist, check_assist_by_label, check_assist_not_applicable,
199        check_assist_not_applicable_by_label,
200    };
201
202    use super::*;
203
204    #[test]
205    fn add_turbo_fish_function() {
206        check_assist(
207            add_turbo_fish,
208            r#"
209fn make<T>() -> T {}
210fn main() {
211    make$0();
212}
213"#,
214            r#"
215fn make<T>() -> T {}
216fn main() {
217    make::<${0:_}>();
218}
219"#,
220        );
221    }
222
223    #[test]
224    fn add_turbo_fish_function_multiple_generic_types() {
225        check_assist(
226            add_turbo_fish,
227            r#"
228fn make<T, A>() -> T {}
229fn main() {
230    make$0();
231}
232"#,
233            r#"
234fn make<T, A>() -> T {}
235fn main() {
236    make::<${1:_}, ${0:_}>();
237}
238"#,
239        );
240    }
241
242    #[test]
243    fn add_turbo_fish_function_many_generic_types() {
244        check_assist(
245            add_turbo_fish,
246            r#"
247fn make<T, A, B, C, D, E, F>() -> T {}
248fn main() {
249    make$0();
250}
251"#,
252            r#"
253fn make<T, A, B, C, D, E, F>() -> T {}
254fn main() {
255    make::<${1:_}, ${2:_}, ${3:_}, ${4:_}, ${5:_}, ${6:_}, ${0:_}>();
256}
257"#,
258        );
259    }
260
261    #[test]
262    fn add_turbo_fish_after_call() {
263        cov_mark::check!(add_turbo_fish_after_call);
264        check_assist(
265            add_turbo_fish,
266            r#"
267fn make<T>() -> T {}
268fn main() {
269    make()$0;
270}
271"#,
272            r#"
273fn make<T>() -> T {}
274fn main() {
275    make::<${0:_}>();
276}
277"#,
278        );
279    }
280
281    #[test]
282    fn add_turbo_fish_method() {
283        check_assist(
284            add_turbo_fish,
285            r#"
286struct S;
287impl S {
288    fn make<T>(&self) -> T {}
289}
290fn main() {
291    S.make$0();
292}
293"#,
294            r#"
295struct S;
296impl S {
297    fn make<T>(&self) -> T {}
298}
299fn main() {
300    S.make::<${0:_}>();
301}
302"#,
303        );
304    }
305
306    #[test]
307    fn add_turbo_fish_one_fish_is_enough() {
308        cov_mark::check!(add_turbo_fish_one_fish_is_enough);
309        check_assist_not_applicable(
310            add_turbo_fish,
311            r#"
312fn make<T>() -> T {}
313fn main() {
314    make$0::<()>();
315}
316"#,
317        );
318    }
319
320    #[test]
321    fn add_turbo_fish_non_generic() {
322        cov_mark::check!(add_turbo_fish_non_generic);
323        check_assist_not_applicable(
324            add_turbo_fish,
325            r#"
326fn make() -> () {}
327fn main() {
328    make$0();
329}
330"#,
331        );
332    }
333
334    #[test]
335    fn add_type_ascription_function() {
336        check_assist_by_label(
337            add_turbo_fish,
338            r#"
339fn make<T>() -> T {}
340fn main() {
341    let x = make$0();
342}
343"#,
344            r#"
345fn make<T>() -> T {}
346fn main() {
347    let x: ${0:_} = make();
348}
349"#,
350            "Add `: _` before assignment operator",
351        );
352    }
353
354    #[test]
355    fn add_type_ascription_after_call() {
356        cov_mark::check!(add_type_ascription_after_call);
357        check_assist_by_label(
358            add_turbo_fish,
359            r#"
360fn make<T>() -> T {}
361fn main() {
362    let x = make()$0;
363}
364"#,
365            r#"
366fn make<T>() -> T {}
367fn main() {
368    let x: ${0:_} = make();
369}
370"#,
371            "Add `: _` before assignment operator",
372        );
373    }
374
375    #[test]
376    fn add_type_ascription_method() {
377        check_assist_by_label(
378            add_turbo_fish,
379            r#"
380struct S;
381impl S {
382    fn make<T>(&self) -> T {}
383}
384fn main() {
385    let x = S.make$0();
386}
387"#,
388            r#"
389struct S;
390impl S {
391    fn make<T>(&self) -> T {}
392}
393fn main() {
394    let x: ${0:_} = S.make();
395}
396"#,
397            "Add `: _` before assignment operator",
398        );
399    }
400
401    #[test]
402    fn add_type_ascription_already_typed() {
403        cov_mark::check!(add_type_ascription_already_typed);
404        check_assist(
405            add_turbo_fish,
406            r#"
407fn make<T>() -> T {}
408fn main() {
409    let x: () = make$0();
410}
411"#,
412            r#"
413fn make<T>() -> T {}
414fn main() {
415    let x: () = make::<${0:_}>();
416}
417"#,
418        );
419    }
420
421    #[test]
422    fn add_type_ascription_append_semicolon() {
423        check_assist_by_label(
424            add_turbo_fish,
425            r#"
426fn make<T>() -> T {}
427fn main() {
428    let x = make$0()
429}
430"#,
431            r#"
432fn make<T>() -> T {}
433fn main() {
434    let x: ${0:_} = make();
435}
436"#,
437            "Add `: _` before assignment operator",
438        );
439    }
440
441    #[test]
442    fn add_type_ascription_missing_pattern() {
443        check_assist_not_applicable_by_label(
444            add_turbo_fish,
445            r#"
446fn make<T>() -> T {}
447fn main() {
448    let = make$0()
449}
450"#,
451            "Add `: _` before assignment operator",
452        );
453    }
454
455    #[test]
456    fn add_turbo_fish_function_lifetime_parameter() {
457        check_assist(
458            add_turbo_fish,
459            r#"
460fn make<'a, T, A>(t: T, a: A) {}
461fn main() {
462    make$0(5, 2);
463}
464"#,
465            r#"
466fn make<'a, T, A>(t: T, a: A) {}
467fn main() {
468    make::<${1:_}, ${0:_}>(5, 2);
469}
470"#,
471        );
472    }
473
474    #[test]
475    fn add_turbo_fish_function_const_parameter() {
476        check_assist(
477            add_turbo_fish,
478            r#"
479fn make<T, const N: usize>(t: T) {}
480fn main() {
481    make$0(3);
482}
483"#,
484            r#"
485fn make<T, const N: usize>(t: T) {}
486fn main() {
487    make::<${1:_}, ${0:_}>(3);
488}
489"#,
490        );
491    }
492}