From bcdacfe50182089772f6e454ecb0904392dfcc21 Mon Sep 17 00:00:00 2001
From: Chayim Refael Friedman <chayimfr@gmail.com>
Date: Wed, 31 Aug 2022 01:24:36 +0000
Subject: [PATCH] Support `?` and `.await` in "Replace turbofish with explicit
 type"

Now that we use type information this is easy.
---
 .../replace_turbofish_with_explicit_type.rs   | 66 ++++++++++++++-----
 1 file changed, 50 insertions(+), 16 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 a2df56d2f69..521447c26df 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,6 +1,6 @@
 use hir::HirDisplay;
 use syntax::{
-    ast::{Expr, GenericArg},
+    ast::{Expr, GenericArg, GenericArgList},
     ast::{LetStmt, Type::InferType},
     AstNode, TextRange,
 };
@@ -35,21 +35,7 @@ pub(crate) fn replace_turbofish_with_explicit_type(
 
     let initializer = let_stmt.initializer()?;
 
-    let generic_args = match &initializer {
-        Expr::MethodCallExpr(ce) => ce.generic_arg_list()?,
-        Expr::CallExpr(ce) => {
-            if let Expr::PathExpr(pe) = ce.expr()? {
-                pe.path()?.segment()?.generic_arg_list()?
-            } else {
-                cov_mark::hit!(not_applicable_if_non_path_function_call);
-                return None;
-            }
-        }
-        _ => {
-            cov_mark::hit!(not_applicable_if_non_function_call_initializer);
-            return None;
-        }
-    };
+    let generic_args = generic_arg_list(&initializer)?;
 
     // Find range of ::<_>
     let colon2 = generic_args.coloncolon_token()?;
@@ -117,6 +103,26 @@ pub(crate) fn replace_turbofish_with_explicit_type(
     None
 }
 
+fn generic_arg_list(expr: &Expr) -> Option<GenericArgList> {
+    match expr {
+        Expr::MethodCallExpr(expr) => expr.generic_arg_list(),
+        Expr::CallExpr(expr) => {
+            if let Expr::PathExpr(pe) = expr.expr()? {
+                pe.path()?.segment()?.generic_arg_list()
+            } else {
+                cov_mark::hit!(not_applicable_if_non_path_function_call);
+                return None;
+            }
+        }
+        Expr::AwaitExpr(expr) => generic_arg_list(&expr.expr()?),
+        Expr::TryExpr(expr) => generic_arg_list(&expr.expr()?),
+        _ => {
+            cov_mark::hit!(not_applicable_if_non_function_call_initializer);
+            None
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -325,6 +331,34 @@ fn make<T>() -> HasDefault<T> {}
 fn main() {
     let a: HasDefault<bool> = make();
 }
+"#,
+        );
+    }
+
+    #[test]
+    fn replaces_turbofish_try_await() {
+        check_assist(
+            replace_turbofish_with_explicit_type,
+            r#"
+//- minicore: option, future
+struct Fut<T>(T);
+impl<T> core::future::Future for Fut<T> {
+    type Output = Option<T>;
+}
+fn make<T>() -> Fut<T> {}
+fn main() {
+    let a = make$0::<bool>().await?;
+}
+"#,
+            r#"
+struct Fut<T>(T);
+impl<T> core::future::Future for Fut<T> {
+    type Output = Option<T>;
+}
+fn make<T>() -> Fut<T> {}
+fn main() {
+    let a: bool = make().await?;
+}
 "#,
         );
     }