refactor: re-use add_turbo_fish function

2 files changed, 137 insertions, 199 deletions

diff --git a/crates/ide_assists/src/handlers/add_turbo_fish.rs b/crates/ide_assists/src/handlers/add_turbo_fish.rs
index 8e9ea4fad..f18e3edf9 100644
--- a/crates/ide_assists/src/handlers/add_turbo_fish.rs
+++ b/crates/ide_assists/src/handlers/add_turbo_fish.rs
@@ -31,11 +31,13 @@ pub(crate) fn add_turbo_fish(acc: &mut Assists, ctx: &AssistContext) -> Option<(
             return None;
         }
         mark::hit!(add_turbo_fish_after_call);
+        mark::hit!(add_type_ascription_after_call);
         arg_list.l_paren_token()?.prev_token().filter(|it| it.kind() == SyntaxKind::IDENT)
     })?;
     let next_token = ident.next_token()?;
     if next_token.kind() == T![::] {
         mark::hit!(add_turbo_fish_one_fish_is_enough);
+        mark::hit!(add_type_ascription_turbofished);
         return None;
     }
     let name_ref = ast::NameRef::cast(ident.parent())?;
@@ -50,8 +52,27 @@ pub(crate) fn add_turbo_fish(acc: &mut Assists, ctx: &AssistContext) -> Option<(
     let generics = hir::GenericDef::Function(fun).params(ctx.sema.db);
     if generics.is_empty() {
         mark::hit!(add_turbo_fish_non_generic);
+        mark::hit!(add_type_ascription_non_generic);
         return None;
     }
+
+    if let Some(let_stmt) = ctx.find_node_at_offset::<ast::LetStmt>() {
+        if let_stmt.colon_token().is_none() {
+            let type_pos = let_stmt.pat()?.syntax().last_token()?.text_range().end();
+            acc.add(
+                AssistId("add_type_ascription", AssistKind::RefactorRewrite),
+                "Add `: _` before assignment operator",
+                ident.text_range(),
+                |builder| match ctx.config.snippet_cap {
+                    Some(cap) => builder.insert_snippet(cap, type_pos, ": ${0:_}"),
+                    None => builder.insert(type_pos, ": _"),
+                },
+            )?
+        } else {
+            mark::hit!(add_type_ascription_already_typed);
+        }
+    }
+
     acc.add(
         AssistId("add_turbo_fish", AssistKind::RefactorRewrite),
         "Add `::<>`",
@@ -65,7 +86,7 @@ pub(crate) fn add_turbo_fish(acc: &mut Assists, ctx: &AssistContext) -> Option<(
 
 #[cfg(test)]
 mod tests {
-    use crate::tests::{check_assist, check_assist_not_applicable};
+    use crate::tests::{check_assist, check_assist_by_label, check_assist_not_applicable};
 
     use super::*;
     use test_utils::mark;
@@ -161,4 +182,119 @@ fn main() {
 "#,
         );
     }
+
+    #[test]
+    fn add_type_ascription_function() {
+        check_assist_by_label(
+            add_turbo_fish,
+            r#"
+fn make<T>() -> T {}
+fn main() {
+    let x = make$0();
+}
+"#,
+            r#"
+fn make<T>() -> T {}
+fn main() {
+    let x: ${0:_} = make();
+}
+"#,
+            "Add `: _` before assignment operator",
+        );
+    }
+
+    #[test]
+    fn add_type_ascription_after_call() {
+        mark::check!(add_type_ascription_after_call);
+        check_assist_by_label(
+            add_turbo_fish,
+            r#"
+fn make<T>() -> T {}
+fn main() {
+    let x = make()$0;
+}
+"#,
+            r#"
+fn make<T>() -> T {}
+fn main() {
+    let x: ${0:_} = make();
+}
+"#,
+            "Add `: _` before assignment operator",
+        );
+    }
+
+    #[test]
+    fn add_type_ascription_method() {
+        check_assist_by_label(
+            add_turbo_fish,
+            r#"
+struct S;
+impl S {
+    fn make<T>(&self) -> T {}
+}
+fn main() {
+    let x = S.make$0();
+}
+"#,
+            r#"
+struct S;
+impl S {
+    fn make<T>(&self) -> T {}
+}
+fn main() {
+    let x: ${0:_} = S.make();
+}
+"#,
+            "Add `: _` before assignment operator",
+        );
+    }
+
+    #[test]
+    fn add_type_ascription_turbofished() {
+        mark::check!(add_type_ascription_turbofished);
+        check_assist_not_applicable(
+            add_turbo_fish,
+            r#"
+fn make<T>() -> T {}
+fn main() {
+    let x = make$0::<()>();
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn add_type_ascription_already_typed() {
+        mark::check!(add_type_ascription_already_typed);
+        check_assist(
+            add_turbo_fish,
+            r#"
+fn make<T>() -> T {}
+fn main() {
+    let x: () = make$0();
+}
+"#,
+            r#"
+fn make<T>() -> T {}
+fn main() {
+    let x: () = make::<${0:_}>();
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn add_type_ascription_non_generic() {
+        mark::check!(add_type_ascription_non_generic);
+        check_assist_not_applicable(
+            add_turbo_fish,
+            r#"
+fn make() -> () {}
+fn main() {
+    let x = make$0();
+}
+"#,
+        );
+    }
 }
diff --git a/crates/ide_assists/src/handlers/add_type_ascription.rs b/crates/ide_assists/src/handlers/add_type_ascription.rs
deleted file mode 100644
index 060c326bf..000000000
--- a/crates/ide_assists/src/handlers/add_type_ascription.rs
+++ /dev/null
@@ -1,198 +0,0 @@
-use ide_db::defs::{Definition, NameRefClass};
-use syntax::{ast, AstNode, SyntaxKind, T};
-use test_utils::mark;
-
-use crate::{
-    assist_context::{AssistContext, Assists},
-    AssistId, AssistKind,
-};
-
-// Assist: add_type_ascription
-//
-// Adds `: _` before the assignment operator to prompt the user for a type.
-//
-// ```
-// fn make<T>() -> T { todo!() }
-// fn main() {
-//     let x = make$0();
-// }
-// ```
-// ->
-// ```
-// fn make<T>() -> T { todo!() }
-// fn main() {
-//     let x: ${0:_} = make();
-// }
-// ```
-pub(crate) fn add_type_ascription(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
-    let let_stmt = ctx.find_node_at_offset::<ast::LetStmt>()?;
-    if let_stmt.colon_token().is_some() {
-        mark::hit!(add_type_ascription_already_typed);
-        return None;
-    }
-    let type_pos = let_stmt.pat()?.syntax().last_token()?.text_range().end();
-
-    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 {
-            return None;
-        }
-        mark::hit!(add_type_ascription_after_call);
-        arg_list.l_paren_token()?.prev_token().filter(|it| it.kind() == SyntaxKind::IDENT)
-    })?;
-    let next_token = ident.next_token()?;
-    if next_token.kind() == T![::] {
-        mark::hit!(add_type_ascription_turbofished);
-        return None;
-    }
-    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,
-    };
-    let fun = match def {
-        Definition::ModuleDef(hir::ModuleDef::Function(it)) => it,
-        _ => return None,
-    };
-    let generics = hir::GenericDef::Function(fun).params(ctx.sema.db);
-    if generics.is_empty() {
-        mark::hit!(add_type_ascription_non_generic);
-        return None;
-    }
-    acc.add(
-        AssistId("add_type_ascription", AssistKind::RefactorRewrite),
-        "Add `: _` before assignment operator",
-        ident.text_range(),
-        |builder| match ctx.config.snippet_cap {
-            Some(cap) => builder.insert_snippet(cap, type_pos, ": ${0:_}"),
-            None => builder.insert(type_pos, ": _"),
-        },
-    )
-}
-
-#[cfg(test)]
-mod tests {
-    use crate::tests::{check_assist, check_assist_not_applicable};
-
-    use super::*;
-    use test_utils::mark;
-
-    #[test]
-    fn add_type_ascription_function() {
-        check_assist(
-            add_type_ascription,
-            r#"
-fn make<T>() -> T {}
-fn main() {
-    let x = make$0();
-}
-"#,
-            r#"
-fn make<T>() -> T {}
-fn main() {
-    let x: ${0:_} = make();
-}
-"#,
-        );
-    }
-
-    #[test]
-    fn add_type_ascription_after_call() {
-        mark::check!(add_type_ascription_after_call);
-        check_assist(
-            add_type_ascription,
-            r#"
-fn make<T>() -> T {}
-fn main() {
-    let x = make()$0;
-}
-"#,
-            r#"
-fn make<T>() -> T {}
-fn main() {
-    let x: ${0:_} = make();
-}
-"#,
-        );
-    }
-
-    #[test]
-    fn add_type_ascription_method() {
-        check_assist(
-            add_type_ascription,
-            r#"
-struct S;
-impl S {
-    fn make<T>(&self) -> T {}
-}
-fn main() {
-    let x = S.make$0();
-}
-"#,
-            r#"
-struct S;
-impl S {
-    fn make<T>(&self) -> T {}
-}
-fn main() {
-    let x: ${0:_} = S.make();
-}
-"#,
-        );
-    }
-
-    #[test]
-    fn add_type_ascription_turbofished() {
-        mark::check!(add_type_ascription_turbofished);
-        check_assist_not_applicable(
-            add_type_ascription,
-            r#"
-fn make<T>() -> T {}
-fn main() {
-    let x = make$0::<()>();
-}
-"#,
-        );
-    }
-
-    #[test]
-    fn add_type_ascription_already_typed() {
-        mark::check!(add_type_ascription_already_typed);
-        check_assist_not_applicable(
-            add_type_ascription,
-            r#"
-fn make<T>() -> T {}
-fn main() {
-    let x: () = make$0();
-}
-"#,
-        );
-    }
-
-    #[test]
-    fn add_type_ascription_non_generic() {
-        mark::check!(add_type_ascription_non_generic);
-        check_assist_not_applicable(
-            add_type_ascription,
-            r#"
-fn make() -> () {}
-fn main() {
-    let x = make$0();
-}
-"#,
-        );
-    }
-
-    #[test]
-    fn add_type_ascription_no_let() {
-        check_assist_not_applicable(
-            add_type_ascription,
-            r#"
-fn make<T>() -> T {}
-fn main() {
-    make$0();
-}
-"#,
-        );
-    }
-}