src/tools/rust-analyzer/crates/ide-assists/src/handlers/unnecessary_async.rs

   1 use ide_db::{
   2     assists::{AssistId, AssistKind},
   3     base_db::FileId,
   4     defs::Definition,
   5     search::FileReference,
   6     syntax_helpers::node_ext::full_path_of_name_ref,
   7 };
   8 use syntax::{
   9     ast::{self, NameLike, NameRef},
  10     AstNode, SyntaxKind, TextRange,
  11 };
  12
  13 use crate::{AssistContext, Assists};
  14
  15 // Assist: unnecessary_async
  16 //
  17 // Removes the `async` mark from functions which have no `.await` in their body.
  18 // Looks for calls to the functions and removes the `.await` on the call site.
  19 //
  20 // ```
  21 // pub async f$0n foo() {}
  22 // pub async fn bar() { foo().await }
  23 // ```
  24 // ->
  25 // ```
  26 // pub fn foo() {}
  27 // pub async fn bar() { foo() }
  28 // ```
  29 pub(crate) fn unnecessary_async(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
  30     let function: ast::Fn = ctx.find_node_at_offset()?;
  31
  32     // Do nothing if the cursor is not on the prototype. This is so that the check does not pollute
  33     // when the user asks us for assists when in the middle of the function body.
  34     // We consider the prototype to be anything that is before the body of the function.
  35     let cursor_position = ctx.offset();
  36     if cursor_position >= function.body()?.syntax().text_range().start() {
  37         return None;
  38     }
  39     // Do nothing if the function isn't async.
  40     if let None = function.async_token() {
  41         return None;
  42     }
  43     // Do nothing if the function has an `await` expression in its body.
  44     if function.body()?.syntax().descendants().find_map(ast::AwaitExpr::cast).is_some() {
  45         return None;
  46     }
  47
  48     // Remove the `async` keyword plus whitespace after it, if any.
  49     let async_range = {
  50         let async_token = function.async_token()?;
  51         let next_token = async_token.next_token()?;
  52         if matches!(next_token.kind(), SyntaxKind::WHITESPACE) {
  53             TextRange::new(async_token.text_range().start(), next_token.text_range().end())
  54         } else {
  55             async_token.text_range()
  56         }
  57     };
  58
  59     // Otherwise, we may remove the `async` keyword.
  60     acc.add(
  61         AssistId("unnecessary_async", AssistKind::QuickFix),
  62         "Remove unnecessary async",
  63         async_range,
  64         |edit| {
  65             // Remove async on the function definition.
  66             edit.replace(async_range, "");
  67
  68             // Remove all `.await`s from calls to the function we remove `async` from.
  69             if let Some(fn_def) = ctx.sema.to_def(&function) {
  70                 for await_expr in find_all_references(ctx, &Definition::Function(fn_def))
  71                     // Keep only references that correspond NameRefs.
  72                     .filter_map(|(_, reference)| match reference.name {
  73                         NameLike::NameRef(nameref) => Some(nameref),
  74                         _ => None,
  75                     })
  76                     // Keep only references that correspond to await expressions
  77                     .filter_map(|nameref| find_await_expression(ctx, &nameref))
  78                 {
  79                     if let Some(await_token) = &await_expr.await_token() {
  80                         edit.replace(await_token.text_range(), "");
  81                     }
  82                     if let Some(dot_token) = &await_expr.dot_token() {
  83                         edit.replace(dot_token.text_range(), "");
  84                     }
  85                 }
  86             }
  87         },
  88     )
  89 }
  90
  91 fn find_all_references(
  92     ctx: &AssistContext<'_>,
  93     def: &Definition,
  94 ) -> impl Iterator<Item = (FileId, FileReference)> {
  95     def.usages(&ctx.sema).all().into_iter().flat_map(|(file_id, references)| {
  96         references.into_iter().map(move |reference| (file_id, reference))
  97     })
  98 }
  99
 100 /// Finds the await expression for the given `NameRef`.
 101 /// If no await expression is found, returns None.
 102 fn find_await_expression(ctx: &AssistContext<'_>, nameref: &NameRef) -> Option<ast::AwaitExpr> {
 103     // From the nameref, walk up the tree to the await expression.
 104     let await_expr = if let Some(path) = full_path_of_name_ref(&nameref) {
 105         // Function calls.
 106         path.syntax()
 107             .parent()
 108             .and_then(ast::PathExpr::cast)?
 109             .syntax()
 110             .parent()
 111             .and_then(ast::CallExpr::cast)?
 112             .syntax()
 113             .parent()
 114             .and_then(ast::AwaitExpr::cast)
 115     } else {
 116         // Method calls.
 117         nameref
 118             .syntax()
 119             .parent()
 120             .and_then(ast::MethodCallExpr::cast)?
 121             .syntax()
 122             .parent()
 123             .and_then(ast::AwaitExpr::cast)
 124     };
 125
 126     ctx.sema.original_ast_node(await_expr?)
 127 }
 128
 129 #[cfg(test)]
 130 mod tests {
 131     use super::*;
 132
 133     use crate::tests::{check_assist, check_assist_not_applicable};
 134
 135     #[test]
 136     fn applies_on_empty_function() {
 137         check_assist(unnecessary_async, "pub async f$0n f() {}", "pub fn f() {}")
 138     }
 139
 140     #[test]
 141     fn applies_and_removes_whitespace() {
 142         check_assist(unnecessary_async, "pub async       f$0n f() {}", "pub fn f() {}")
 143     }
 144
 145     #[test]
 146     fn does_not_apply_on_non_async_function() {
 147         check_assist_not_applicable(unnecessary_async, "pub f$0n f() {}")
 148     }
 149
 150     #[test]
 151     fn applies_on_function_with_a_non_await_expr() {
 152         check_assist(unnecessary_async, "pub async f$0n f() { f2() }", "pub fn f() { f2() }")
 153     }
 154
 155     #[test]
 156     fn does_not_apply_on_function_with_an_await_expr() {
 157         check_assist_not_applicable(unnecessary_async, "pub async f$0n f() { f2().await }")
 158     }
 159
 160     #[test]
 161     fn applies_and_removes_await_on_reference() {
 162         check_assist(
 163             unnecessary_async,
 164             r#"
 165 pub async fn f4() { }
 166 pub async f$0n f2() { }
 167 pub async fn f() { f2().await }
 168 pub async fn f3() { f2().await }"#,
 169             r#"
 170 pub async fn f4() { }
 171 pub fn f2() { }
 172 pub async fn f() { f2() }
 173 pub async fn f3() { f2() }"#,
 174         )
 175     }
 176
 177     #[test]
 178     fn applies_and_removes_await_from_within_module() {
 179         check_assist(
 180             unnecessary_async,
 181             r#"
 182 pub async fn f4() { }
 183 mod a { pub async f$0n f2() { } }
 184 pub async fn f() { a::f2().await }
 185 pub async fn f3() { a::f2().await }"#,
 186             r#"
 187 pub async fn f4() { }
 188 mod a { pub fn f2() { } }
 189 pub async fn f() { a::f2() }
 190 pub async fn f3() { a::f2() }"#,
 191         )
 192     }
 193
 194     #[test]
 195     fn applies_and_removes_await_on_inner_await() {
 196         check_assist(
 197             unnecessary_async,
 198             // Ensure that it is the first await on the 3rd line that is removed
 199             r#"
 200 pub async fn f() { f2().await }
 201 pub async f$0n f2() -> i32 { 1 }
 202 pub async fn f3() { f4(f2().await).await }
 203 pub async fn f4(i: i32) { }"#,
 204             r#"
 205 pub async fn f() { f2() }
 206 pub fn f2() -> i32 { 1 }
 207 pub async fn f3() { f4(f2()).await }
 208 pub async fn f4(i: i32) { }"#,
 209         )
 210     }
 211
 212     #[test]
 213     fn applies_and_removes_await_on_outer_await() {
 214         check_assist(
 215             unnecessary_async,
 216             // Ensure that it is the second await on the 3rd line that is removed
 217             r#"
 218 pub async fn f() { f2().await }
 219 pub async f$0n f2(i: i32) { }
 220 pub async fn f3() { f2(f4().await).await }
 221 pub async fn f4() -> i32 { 1 }"#,
 222             r#"
 223 pub async fn f() { f2() }
 224 pub fn f2(i: i32) { }
 225 pub async fn f3() { f2(f4().await) }
 226 pub async fn f4() -> i32 { 1 }"#,
 227         )
 228     }
 229
 230     #[test]
 231     fn applies_on_method_call() {
 232         check_assist(
 233             unnecessary_async,
 234             r#"
 235 pub struct S { }
 236 impl S { pub async f$0n f2(&self) { } }
 237 pub async fn f(s: &S) { s.f2().await }"#,
 238             r#"
 239 pub struct S { }
 240 impl S { pub fn f2(&self) { } }
 241 pub async fn f(s: &S) { s.f2() }"#,
 242         )
 243     }
 244
 245     #[test]
 246     fn does_not_apply_on_function_with_a_nested_await_expr() {
 247         check_assist_not_applicable(
 248             unnecessary_async,
 249             "async f$0n f() { if true { loop { f2().await } } }",
 250         )
 251     }
 252
 253     #[test]
 254     fn does_not_apply_when_not_on_prototype() {
 255         check_assist_not_applicable(unnecessary_async, "pub async fn f() { $0f2() }")
 256     }
 257 }