From 74f8fe77632b992e6b28e98b6027d4e8ffd29398 Mon Sep 17 00:00:00 2001
From: Jamie Cunliffe <Jamie.Cunliffe@outlook.com>
Date: Sat, 24 Apr 2021 18:01:36 +0100
Subject: Extract function assist will add async if required

The extract function assist will check for an AWAIT_EXPR in the body
and if found, will add async to the generated function.
---
 .../ide_assists/src/handlers/extract_function.rs   | 65 +++++++++++++++++++++-
 1 file changed, 63 insertions(+), 2 deletions(-)

diff --git a/crates/ide_assists/src/handlers/extract_function.rs b/crates/ide_assists/src/handlers/extract_function.rs
index 6311afc1f..a2dba915c 100644
--- a/crates/ide_assists/src/handlers/extract_function.rs
+++ b/crates/ide_assists/src/handlers/extract_function.rs
@@ -642,6 +642,10 @@ fn vars_used_in_body(ctx: &AssistContext, body: &FunctionBody) -> Vec<Local> {
         .collect()
 }
 
+fn body_contains_await(body: &FunctionBody) -> bool {
+    body.descendants().any(|d| matches!(d.kind(), SyntaxKind::AWAIT_EXPR))
+}
+
 /// find `self` param, that was not defined inside `body`
 ///
 /// It should skip `self` params from impls inside `body`
@@ -1123,9 +1127,10 @@ fn format_function(
     let params = make_param_list(ctx, module, fun);
     let ret_ty = make_ret_ty(ctx, module, fun);
     let body = make_body(ctx, old_indent, new_indent, fun);
+    let async_kw = if body_contains_await(&fun.body) { "async " } else { "" };
     match ctx.config.snippet_cap {
-        Some(_) => format_to!(fn_def, "\n\n{}fn $0{}{}", new_indent, fun.name, params),
-        None => format_to!(fn_def, "\n\n{}fn {}{}", new_indent, fun.name, params),
+        Some(_) => format_to!(fn_def, "\n\n{}{}fn $0{}{}", new_indent, async_kw, fun.name, params),
+        None => format_to!(fn_def, "\n\n{}{}fn {}{}", new_indent, async_kw, fun.name, params),
     }
     if let Some(ret_ty) = ret_ty {
         format_to!(fn_def, " {}", ret_ty);
@@ -3565,4 +3570,60 @@ fn $0fun_name(n: i32) -> i32 {
 }",
         );
     }
+
+    #[test]
+    fn extract_with_await() {
+        check_assist(
+            extract_function,
+            r#"fn main() {
+    $0some_function().await;$0
+}
+
+async fn some_function() {
+
+}
+"#,
+            r#"
+fn main() {
+    fun_name();
+}
+
+async fn $0fun_name() {
+    some_function().await;
+}
+
+async fn some_function() {
+
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn extract_with_await_in_args() {
+        check_assist(
+            extract_function,
+            r#"fn main() {
+    $0function_call("a", some_function().await);$0
+}
+
+async fn some_function() {
+
+}
+"#,
+            r#"
+fn main() {
+    fun_name();
+}
+
+async fn $0fun_name() {
+    function_call("a", some_function().await);
+}
+
+async fn some_function() {
+
+}
+"#,
+        );
+    }
 }
-- 
cgit v1.2.3