Use type information to deduce the correct type for "Replace turbofish with explicit type", even when it is not exactly the same as the turbofish type

I implemented that by checking the expressions' type.
This could probably be implemented better by taking the function's return type and substituting the generic parameter with the provided turbofish, but this is more complicated.

1 files changed, 91 insertions, 3 deletions

diff --git a/crates/ide-assists/src/handlers/replace_turbofish_with_explicit_type.rs b/crates/ide-assists/src/handlers/replace_turbofish_with_explicit_type.rs
index 5242f3b5100..a2df56d2f69 100644
--- a/crates/ide-assists/src/handlers/replace_turbofish_with_explicit_type.rs
+++ b/crates/ide-assists/src/handlers/replace_turbofish_with_explicit_type.rs
@@ -1,3 +1,4 @@
+use hir::HirDisplay;
 use syntax::{
     ast::{Expr, GenericArg},
     ast::{LetStmt, Type::InferType},
@@ -65,7 +66,16 @@ pub(crate) fn replace_turbofish_with_explicit_type(
 
     // An improvement would be to check that this is correctly part of the return value of the
     // function call, or sub in the actual return type.
-    let turbofish_type = &turbofish_args[0];
+    let returned_type = match ctx.sema.type_of_expr(&initializer) {
+        Some(returned_type) if !returned_type.original.contains_unknown() => {
+            let module = ctx.sema.scope(let_stmt.syntax())?.module();
+            returned_type.original.display_source_code(ctx.db(), module.into()).ok()?
+        }
+        _ => {
+            cov_mark::hit!(fallback_to_turbofish_type_if_type_info_not_available);
+            turbofish_args[0].to_string()
+        }
+    };
 
     let initializer_start = initializer.syntax().text_range().start();
     if ctx.offset() > turbofish_range.end() || ctx.offset() < initializer_start {
@@ -83,7 +93,7 @@ pub(crate) fn replace_turbofish_with_explicit_type(
             "Replace turbofish with explicit type",
             TextRange::new(initializer_start, turbofish_range.end()),
             |builder| {
-                builder.insert(ident_range.end(), format!(": {}", turbofish_type));
+                builder.insert(ident_range.end(), format!(": {}", returned_type));
                 builder.delete(turbofish_range);
             },
         );
@@ -98,7 +108,7 @@ pub(crate) fn replace_turbofish_with_explicit_type(
             "Replace `_` with turbofish type",
             turbofish_range,
             |builder| {
-                builder.replace(underscore_range, turbofish_type.to_string());
+                builder.replace(underscore_range, returned_type);
                 builder.delete(turbofish_range);
             },
         );
@@ -115,6 +125,7 @@ mod tests {
 
     #[test]
     fn replaces_turbofish_for_vec_string() {
+        cov_mark::check!(fallback_to_turbofish_type_if_type_info_not_available);
         check_assist(
             replace_turbofish_with_explicit_type,
             r#"
@@ -135,6 +146,7 @@ fn main() {
     #[test]
     fn replaces_method_calls() {
         // foo.make() is a method call which uses a different expr in the let initializer
+        cov_mark::check!(fallback_to_turbofish_type_if_type_info_not_available);
         check_assist(
             replace_turbofish_with_explicit_type,
             r#"
@@ -240,4 +252,80 @@ fn main() {
 "#,
         );
     }
+
+    #[test]
+    fn replaces_turbofish_for_known_type() {
+        check_assist(
+            replace_turbofish_with_explicit_type,
+            r#"
+fn make<T>() -> T {}
+fn main() {
+    let a = make$0::<i32>();
+}
+"#,
+            r#"
+fn make<T>() -> T {}
+fn main() {
+    let a: i32 = make();
+}
+"#,
+        );
+        check_assist(
+            replace_turbofish_with_explicit_type,
+            r#"
+//- minicore: option
+fn make<T>() -> T {}
+fn main() {
+    let a = make$0::<Option<bool>>();
+}
+"#,
+            r#"
+fn make<T>() -> T {}
+fn main() {
+    let a: Option<bool> = make();
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn replaces_turbofish_not_same_type() {
+        check_assist(
+            replace_turbofish_with_explicit_type,
+            r#"
+//- minicore: option
+fn make<T>() -> Option<T> {}
+fn main() {
+    let a = make$0::<u128>();
+}
+"#,
+            r#"
+fn make<T>() -> Option<T> {}
+fn main() {
+    let a: Option<u128> = make();
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn replaces_turbofish_for_type_with_defaulted_generic_param() {
+        check_assist(
+            replace_turbofish_with_explicit_type,
+            r#"
+struct HasDefault<T, U = i32>(T, U);
+fn make<T>() -> HasDefault<T> {}
+fn main() {
+    let a = make$0::<bool>();
+}
+"#,
+            r#"
+struct HasDefault<T, U = i32>(T, U);
+fn make<T>() -> HasDefault<T> {}
+fn main() {
+    let a: HasDefault<bool> = make();
+}
+"#,
+        );
+    }
 }