From 61643513b620ea6758d45665db26344d9a9b01b7 Mon Sep 17 00:00:00 2001
From: vi_mi <fenil.jain2018@vitstudent.ac.in>
Date: Sun, 3 Oct 2021 08:26:19 +0000
Subject: [PATCH] fix: await insertion with try_expr during extract_function

---
 .../src/handlers/extract_function.rs          | 73 ++++++++++++++++++-
 crates/syntax/src/ast/make.rs                 |  3 +
 2 files changed, 72 insertions(+), 4 deletions(-)

diff --git a/crates/ide_assists/src/handlers/extract_function.rs b/crates/ide_assists/src/handlers/extract_function.rs
index 585ef31daf..2e91f8d4ab 100644
--- a/crates/ide_assists/src/handlers/extract_function.rs
+++ b/crates/ide_assists/src/handlers/extract_function.rs
@@ -1090,7 +1090,7 @@ fn make_call(ctx: &AssistContext, fun: &Function, indent: IndentLevel) -> String
 
     let args = make::arg_list(fun.params.iter().map(|param| param.to_arg(ctx)));
     let name = fun.name.clone();
-    let call_expr = if fun.self_param.is_some() {
+    let mut call_expr = if fun.self_param.is_some() {
         let self_arg = make::expr_path(make::ext::ident_path("self"));
         make::expr_method_call(self_arg, name, args)
     } else {
@@ -1100,6 +1100,9 @@ fn make_call(ctx: &AssistContext, fun: &Function, indent: IndentLevel) -> String
 
     let handler = FlowHandler::from_ret_ty(fun, &ret_ty);
 
+    if fun.control_flow.is_async {
+        call_expr = make::expr_await(call_expr);
+    }
     let expr = handler.make_call_expr(call_expr).indent(indent);
 
     let mut_modifier = |var: &OutlivedLocal| if var.mut_usage_outside_body { "mut " } else { "" };
@@ -1119,10 +1122,8 @@ fn make_call(ctx: &AssistContext, fun: &Function, indent: IndentLevel) -> String
             buf.push_str(") = ");
         }
     }
+
     format_to!(buf, "{}", expr);
-    if fun.control_flow.is_async {
-        buf.push_str(".await");
-    }
     let insert_comma = fun
         .body
         .parent()
@@ -3870,6 +3871,70 @@ async fn $0fun_name() {
 
 async fn some_function() {
 
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn extract_with_await_and_result_not_producing_match_expr() {
+        check_assist(
+            extract_function,
+            r#"
+async fn foo() -> Result<(), ()> {
+    $0async {}.await;
+    Err(())?$0
+}
+"#,
+            r#"
+async fn foo() -> Result<(), ()> {
+    fun_name().await?
+}
+
+async fn $0fun_name() -> _ {
+    async {}.await;
+    Err(())?
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn extract_with_await_and_result_producing_match_expr() {
+        check_assist(
+            extract_function,
+            r#"
+async fn foo() -> i32 {
+    loop {
+        let n = 1;$0
+        let k = async { 1 }.await;
+        if k == 42 {
+            break 3;
+        }
+        let m = k + 1;$0
+        let h = 1 + m;
+    }
+}
+"#,
+            r#"
+async fn foo() -> i32 {
+    loop {
+        let n = 1;
+        let m = match fun_name().await {
+            Ok(value) => value,
+            Err(value) => break value,
+        };
+        let h = 1 + m;
+    }
+}
+
+async fn $0fun_name() -> Result<i32, i32> {
+    let k = async { 1 }.await;
+    if k == 42 {
+        return Err(3);
+    }
+    let m = k + 1;
+    Ok(m)
 }
 "#,
         );
diff --git a/crates/syntax/src/ast/make.rs b/crates/syntax/src/ast/make.rs
index dc13e916c8..40c0cd7a37 100644
--- a/crates/syntax/src/ast/make.rs
+++ b/crates/syntax/src/ast/make.rs
@@ -299,6 +299,9 @@ pub fn expr_return(expr: Option<ast::Expr>) -> ast::Expr {
 pub fn expr_try(expr: ast::Expr) -> ast::Expr {
     expr_from_text(&format!("{}?", expr))
 }
+pub fn expr_await(expr: ast::Expr) -> ast::Expr {
+    expr_from_text(&format!("{}.await", expr))
+}
 pub fn expr_match(expr: ast::Expr, match_arm_list: ast::MatchArmList) -> ast::Expr {
     expr_from_text(&format!("match {} {}", expr, match_arm_list))
 }