From f4e35c60b04b381395150669f1e8656dccc5f94f Mon Sep 17 00:00:00 2001
From: Michael Goulet <michael@errs.io>
Date: Mon, 15 Jan 2024 19:24:03 +0000
Subject: [PATCH] Fix async closure call suggestion

---
 compiler/rustc_hir_typeck/src/callee.rs       | 82 +++++++++++--------
 .../suggest-on-bare-closure-call.stderr       |  5 ++
 2 files changed, 51 insertions(+), 36 deletions(-)

diff --git a/compiler/rustc_hir_typeck/src/callee.rs b/compiler/rustc_hir_typeck/src/callee.rs
index 1a4e03d50cae0..d486f069989ee 100644
--- a/compiler/rustc_hir_typeck/src/callee.rs
+++ b/compiler/rustc_hir_typeck/src/callee.rs
@@ -293,49 +293,59 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         callee_node: &hir::ExprKind<'_>,
         callee_span: Span,
     ) {
+        let hir::ExprKind::Block(..) = callee_node else {
+            // Only calls on blocks suggested here.
+            return;
+        };
+
         let hir = self.tcx.hir();
-        let parent_hir_id = hir.parent_id(hir_id);
-        let parent_node = self.tcx.hir_node(parent_hir_id);
-        if let (
-            hir::Node::Expr(hir::Expr {
-                kind: hir::ExprKind::Closure(&hir::Closure { fn_decl_span, kind, .. }),
+        let fn_decl_span = if let hir::Node::Expr(hir::Expr {
+            kind: hir::ExprKind::Closure(&hir::Closure { fn_decl_span, .. }),
+            ..
+        }) = hir.get_parent(hir_id)
+        {
+            fn_decl_span
+        } else if let Some((
+            _,
+            hir::Node::Expr(&hir::Expr {
+                hir_id: parent_hir_id,
+                kind:
+                    hir::ExprKind::Closure(&hir::Closure {
+                        kind:
+                            hir::ClosureKind::Coroutine(hir::CoroutineKind::Desugared(
+                                hir::CoroutineDesugaring::Async,
+                                hir::CoroutineSource::Closure,
+                            )),
+                        ..
+                    }),
                 ..
             }),
-            hir::ExprKind::Block(..),
-        ) = (parent_node, callee_node)
+        )) = hir.parent_iter(hir_id).nth(3)
         {
-            let fn_decl_span = if matches!(
-                kind,
-                hir::ClosureKind::Coroutine(hir::CoroutineKind::Desugared(
-                    hir::CoroutineDesugaring::Async,
-                    hir::CoroutineSource::Closure
-                ),)
-            ) {
-                // Actually need to unwrap one more layer of HIR to get to
-                // the _real_ closure...
-                let async_closure = hir.parent_id(parent_hir_id);
-                if let hir::Node::Expr(hir::Expr {
-                    kind: hir::ExprKind::Closure(&hir::Closure { fn_decl_span, .. }),
-                    ..
-                }) = self.tcx.hir_node(async_closure)
-                {
-                    fn_decl_span
-                } else {
-                    return;
-                }
-            } else {
+            // Actually need to unwrap one more layer of HIR to get to
+            // the _real_ closure...
+            let async_closure = hir.parent_id(parent_hir_id);
+            if let hir::Node::Expr(hir::Expr {
+                kind: hir::ExprKind::Closure(&hir::Closure { fn_decl_span, .. }),
+                ..
+            }) = self.tcx.hir_node(async_closure)
+            {
                 fn_decl_span
-            };
+            } else {
+                return;
+            }
+        } else {
+            return;
+        };
 
-            let start = fn_decl_span.shrink_to_lo();
-            let end = callee_span.shrink_to_hi();
-            err.multipart_suggestion(
-                "if you meant to create this closure and immediately call it, surround the \
+        let start = fn_decl_span.shrink_to_lo();
+        let end = callee_span.shrink_to_hi();
+        err.multipart_suggestion(
+            "if you meant to create this closure and immediately call it, surround the \
                 closure with parentheses",
-                vec![(start, "(".to_string()), (end, ")".to_string())],
-                Applicability::MaybeIncorrect,
-            );
-        }
+            vec![(start, "(".to_string()), (end, ")".to_string())],
+            Applicability::MaybeIncorrect,
+        );
     }
 
     /// Give appropriate suggestion when encountering `[("a", 0) ("b", 1)]`, where the
diff --git a/tests/ui/suggestions/suggest-on-bare-closure-call.stderr b/tests/ui/suggestions/suggest-on-bare-closure-call.stderr
index 77631b36e0f30..e65a6eb4939d9 100644
--- a/tests/ui/suggestions/suggest-on-bare-closure-call.stderr
+++ b/tests/ui/suggestions/suggest-on-bare-closure-call.stderr
@@ -18,6 +18,11 @@ LL |     let _ = async ||{}();
    |                     ^^--
    |                     |
    |                     call expression requires function
+   |
+help: if you meant to create this closure and immediately call it, surround the closure with parentheses
+   |
+LL |     let _ = (async ||{})();
+   |             +          +
 
 error: aborting due to 2 previous errors