From 384aed834cd83bdd0e93a2ef5690d3e44bf15856 Mon Sep 17 00:00:00 2001
From: Michael Goulet <michael@errs.io>
Date: Sun, 1 Sep 2024 05:33:02 -0400
Subject: [PATCH] Do not call query to compute coroutine layout for synthetic
 body of async closure

---
 compiler/rustc_mir_transform/src/validate.rs  | 13 ++++++++++++-
 .../async-closures/validate-synthetic-body.rs | 19 +++++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)
 create mode 100644 tests/ui/async-await/async-closures/validate-synthetic-body.rs

diff --git a/compiler/rustc_mir_transform/src/validate.rs b/compiler/rustc_mir_transform/src/validate.rs
index 99e06f59dd04d..80410c6d690dd 100644
--- a/compiler/rustc_mir_transform/src/validate.rs
+++ b/compiler/rustc_mir_transform/src/validate.rs
@@ -1,6 +1,7 @@
 //! Validates the MIR to ensure that invariants are upheld.
 
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
+use rustc_hir as hir;
 use rustc_hir::LangItem;
 use rustc_index::bit_set::BitSet;
 use rustc_index::IndexVec;
@@ -714,7 +715,17 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                             // since we may be in the process of computing this MIR in the
                             // first place.
                             let layout = if def_id == self.caller_body.source.def_id() {
-                                // FIXME: This is not right for async closures.
+                                self.caller_body.coroutine_layout_raw()
+                            } else if let Some(hir::CoroutineKind::Desugared(
+                                _,
+                                hir::CoroutineSource::Closure,
+                            )) = self.tcx.coroutine_kind(def_id)
+                                && let ty::ClosureKind::FnOnce =
+                                    args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap()
+                                && self.caller_body.source.def_id()
+                                    == self.tcx.coroutine_by_move_body_def_id(def_id)
+                            {
+                                // Same if this is the by-move body of a coroutine-closure.
                                 self.caller_body.coroutine_layout_raw()
                             } else {
                                 self.tcx.coroutine_layout(def_id, args.as_coroutine().kind_ty())
diff --git a/tests/ui/async-await/async-closures/validate-synthetic-body.rs b/tests/ui/async-await/async-closures/validate-synthetic-body.rs
new file mode 100644
index 0000000000000..67e683ac08a22
--- /dev/null
+++ b/tests/ui/async-await/async-closures/validate-synthetic-body.rs
@@ -0,0 +1,19 @@
+//@ check-pass
+//@ edition: 2021
+
+#![feature(async_closure)]
+
+// Make sure that we don't hit a query cycle when validating
+// the by-move coroutine body for an async closure.
+
+use std::future::Future;
+
+async fn test<Fut: Future>(operation: impl Fn() -> Fut) {
+    operation().await;
+}
+
+pub async fn orchestrate_simple_crud() {
+    test(async || async {}.await).await;
+}
+
+fn main() {}