feat: Highlight exit points of async blocks

Async blocks act similar to async functions in that the await keywords
are related, but also act like functions where the exit points are
related.

crates/ide/src/highlight_related.rs

@@ -281,12 +281,7 @@ fn highlight_references(
     }
 }
 
-// If `file_id` is None,
+fn hl_exit_points(
-pub(crate) fn highlight_exit_points(
-    sema: &Semantics<'_, RootDatabase>,
-    token: SyntaxToken,
-) -> FxHashMap<EditionedFileId, Vec<HighlightedRange>> {
-    fn hl(
     sema: &Semantics<'_, RootDatabase>,
     def_token: Option<SyntaxToken>,
     body: ast::Expr,
@@ -310,9 +305,7 @@ pub(crate) fn highlight_exit_points(
         let file_id = sema.hir_file_for(expr.syntax());
 
         let range = match &expr {
-                ast::Expr::TryExpr(try_) => {
+            ast::Expr::TryExpr(try_) => try_.question_mark_token().map(|token| token.text_range()),
-                    try_.question_mark_token().map(|token| token.text_range())
-                }
             ast::Expr::MethodCallExpr(_) | ast::Expr::CallExpr(_) | ast::Expr::MacroExpr(_)
                 if sema.type_of_expr(&expr).map_or(false, |ty| ty.original.is_never()) =>
             {
@@ -332,9 +325,7 @@ pub(crate) fn highlight_exit_points(
             let file_id = sema.hir_file_for(expr.syntax());
 
             let range = match &expr {
-                    ast::Expr::ReturnExpr(expr) => {
+                ast::Expr::ReturnExpr(expr) => expr.return_token().map(|token| token.text_range()),
-                        expr.return_token().map(|token| token.text_range())
-                    }
                 _ => None,
             };
 
@@ -361,19 +352,24 @@ pub(crate) fn highlight_exit_points(
     Some(highlights)
 }
 
+// If `file_id` is None,
+pub(crate) fn highlight_exit_points(
+    sema: &Semantics<'_, RootDatabase>,
+    token: SyntaxToken,
+) -> FxHashMap<EditionedFileId, Vec<HighlightedRange>> {
     let mut res = FxHashMap::default();
     for def in goto_definition::find_fn_or_blocks(sema, &token) {
         let new_map = match_ast! {
             match def {
-                ast::Fn(fn_) => fn_.body().and_then(|body| hl(sema, fn_.fn_token(), body.into())),
+                ast::Fn(fn_) => fn_.body().and_then(|body| hl_exit_points(sema, fn_.fn_token(), body.into())),
                 ast::ClosureExpr(closure) => {
                     let pipe_tok = closure.param_list().and_then(|p| p.pipe_token());
-                    closure.body().and_then(|body| hl(sema, pipe_tok, body))
+                    closure.body().and_then(|body| hl_exit_points(sema, pipe_tok, body))
                 },
                 ast::BlockExpr(blk) => match blk.modifier() {
-                    Some(ast::BlockModifier::Async(t)) => hl(sema, Some(t), blk.into()),
+                    Some(ast::BlockModifier::Async(t)) => hl_exit_points(sema, Some(t), blk.into()),
                     Some(ast::BlockModifier::Try(t)) if token.kind() != T![return] => {
-                        hl(sema, Some(t), blk.into())
+                        hl_exit_points(sema, Some(t), blk.into())
                     },
                     _ => continue,
                 },
@@ -520,6 +516,12 @@ pub(crate) fn highlight_yield_points(
                     if block_expr.async_token().is_none() {
                         continue;
                     }
+
+                    // Async blocks act similar to closures. So we want to
+                    // highlight their exit points too.
+                    let exit_points = hl_exit_points(sema, block_expr.async_token(), block_expr.clone().into());
+                    merge_map(&mut res, exit_points);
+
                     hl(sema, block_expr.async_token(), Some(block_expr.into()))
                 },
                 ast::ClosureExpr(closure) => hl(sema, closure.async_token(), closure.body()),
@@ -876,6 +878,27 @@ pub async$0 fn foo() {
         );
     }
 
+    #[test]
+    fn test_hl_exit_points_of_async_blocks() {
+        check(
+            r#"
+pub fn foo() {
+    let x = async$0 {
+         // ^^^^^
+        0.await;
+       // ^^^^^
+       0?;
+     // ^
+       return 0;
+    // ^^^^^^
+       0
+    // ^
+    };
+}
+"#,
+        );
+    }
+
     #[test]
     fn test_hl_let_else_yield_points() {
         check(
@@ -925,11 +948,10 @@ async fn foo() {
 async fn foo() {
     (async {
   // ^^^^^
-        (async {
+        (async { 0.await }).await$0
-           0.await
+     // ^^^^^^^^^^^^^^^^^^^^^^^^^
-        }).await$0 }
                          // ^^^^^
-    ).await;
+    }).await;
 }
 "#,
         );