Merge #11481

[rust.git] / crates / ide_assists / src / handlers / add_turbo_fish.rs

diff --git a/crates/ide_assists/src/handlers/add_turbo_fish.rs b/crates/ide_assists/src/handlers/add_turbo_fish.rs
index ee879c15197bad428c57294f8884f5d858a71e0d..c3d27f7ea65bec30f905f9cc41ad391db867cbb1 100644 (file)
--- a/crates/ide_assists/src/handlers/add_turbo_fish.rs
+++ b/crates/ide_assists/src/handlers/add_turbo_fish.rs
@@ -1,4 +1,5 @@
 use ide_db::defs::{Definition, NameRefClass};
+use itertools::Itertools;
 use syntax::{ast, AstNode, SyntaxKind, T};
 
 use crate::{
@@ -26,7 +27,7 @@
 pub(crate) fn add_turbo_fish(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
     let ident = ctx.find_token_syntax_at_offset(SyntaxKind::IDENT).or_else(|| {
         let arg_list = ctx.find_node_at_offset::<ast::ArgList>()?;
-        if arg_list.args().count() > 0 {
+        if arg_list.args().next().is_some() {
             return None;
         }
         cov_mark::hit!(add_turbo_fish_after_call);
@@ -38,13 +39,13 @@ pub(crate) fn add_turbo_fish(acc: &mut Assists, ctx: &AssistContext) -> Option<(
         cov_mark::hit!(add_turbo_fish_one_fish_is_enough);
         return None;
     }
-    let name_ref = ast::NameRef::cast(ident.parent())?;
+    let name_ref = ast::NameRef::cast(ident.parent()?)?;
     let def = match NameRefClass::classify(&ctx.sema, &name_ref)? {
         NameRefClass::Definition(def) => def,
-        NameRefClass::ExternCrate(_) | NameRefClass::FieldShorthand { .. } => return None,
+        NameRefClass::FieldShorthand { .. } => return None,
     };
     let fun = match def {
-        Definition::ModuleDef(hir::ModuleDef::Function(it)) => it,
+        Definition::Function(it) => it,
         _ => return None,
     };
     let generics = hir::GenericDef::Function(fun).params(ctx.sema.db);
@@ -77,17 +78,42 @@ pub(crate) fn add_turbo_fish(acc: &mut Assists, ctx: &AssistContext) -> Option<(
         }
     }
 
+    let number_of_arguments = generics
+        .iter()
+        .filter(|param| {
+            matches!(param, hir::GenericParam::TypeParam(_) | hir::GenericParam::ConstParam(_))
+        })
+        .count();
+
     acc.add(
         AssistId("add_turbo_fish", AssistKind::RefactorRewrite),
         "Add `::<>`",
         ident.text_range(),
         |builder| match ctx.config.snippet_cap {
-            Some(cap) => builder.insert_snippet(cap, ident.text_range().end(), "::<${0:_}>"),
-            None => builder.insert(ident.text_range().end(), "::<_>"),
+            Some(cap) => {
+                let snip = format!("::<{}>", get_snippet_fish_head(number_of_arguments));
+                builder.insert_snippet(cap, ident.text_range().end(), snip)
+            }
+            None => {
+                let fish_head = std::iter::repeat("_").take(number_of_arguments).format(", ");
+                let snip = format!("::<{}>", fish_head);
+                builder.insert(ident.text_range().end(), snip);
+            }
         },
     )
 }
 
+/// This will create a snippet string with tabstops marked
+fn get_snippet_fish_head(number_of_arguments: usize) -> String {
+    let mut fish_head = (1..number_of_arguments)
+        .format_with("", |i, f| f(&format_args!("${{{}:_}}, ", i)))
+        .to_string();
+
+    // tabstop 0 is a special case and always the last one
+    fish_head.push_str("${0:_}");
+    fish_head
+}
+
 #[cfg(test)]
 mod tests {
     use crate::tests::{check_assist, check_assist_by_label, check_assist_not_applicable};
@@ -113,6 +139,44 @@ fn main() {
         );
     }
 
+    #[test]
+    fn add_turbo_fish_function_multiple_generic_types() {
+        check_assist(
+            add_turbo_fish,
+            r#"
+fn make<T, A>() -> T {}
+fn main() {
+    make$0();
+}
+"#,
+            r#"
+fn make<T, A>() -> T {}
+fn main() {
+    make::<${1:_}, ${0:_}>();
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn add_turbo_fish_function_many_generic_types() {
+        check_assist(
+            add_turbo_fish,
+            r#"
+fn make<T, A, B, C, D, E, F>() -> T {}
+fn main() {
+    make$0();
+}
+"#,
+            r#"
+fn make<T, A, B, C, D, E, F>() -> T {}
+fn main() {
+    make::<${1:_}, ${2:_}, ${3:_}, ${4:_}, ${5:_}, ${6:_}, ${0:_}>();
+}
+"#,
+        );
+    }
+
     #[test]
     fn add_turbo_fish_after_call() {
         cov_mark::check!(add_turbo_fish_after_call);
@@ -292,4 +356,42 @@ fn main() {
             "Add `: _` before assignment operator",
         );
     }
+
+    #[test]
+    fn add_turbo_fish_function_lifetime_parameter() {
+        check_assist(
+            add_turbo_fish,
+            r#"
+fn make<'a, T, A>(t: T, a: A) {}
+fn main() {
+    make$0(5, 2);
+}
+"#,
+            r#"
+fn make<'a, T, A>(t: T, a: A) {}
+fn main() {
+    make::<${1:_}, ${0:_}>(5, 2);
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn add_turbo_fish_function_const_parameter() {
+        check_assist(
+            add_turbo_fish,
+            r#"
+fn make<T, const N: usize>(t: T) {}
+fn main() {
+    make$0(3);
+}
+"#,
+            r#"
+fn make<T, const N: usize>(t: T) {}
+fn main() {
+    make::<${1:_}, ${0:_}>(3);
+}
+"#,
+        );
+    }
 }