From bce7c4b70e75e06058ddaa14bc85b3a47b5e5306 Mon Sep 17 00:00:00 2001
From: Michael Goulet <michael@errs.io>
Date: Wed, 28 Aug 2024 00:10:45 -0400
Subject: [PATCH] Don't build by-move body when async closure is tainted

---
 .../src/coroutine/by_move_body.rs              | 11 +++++++++--
 .../async-closures/tainted-body-2.rs           | 18 ++++++++++++++++++
 .../async-closures/tainted-body-2.stderr       |  9 +++++++++
 3 files changed, 36 insertions(+), 2 deletions(-)
 create mode 100644 tests/ui/async-await/async-closures/tainted-body-2.rs
 create mode 100644 tests/ui/async-await/async-closures/tainted-body-2.stderr

diff --git a/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs b/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs
index cf39c136b01..9e35dfed0ee 100644
--- a/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs
+++ b/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs
@@ -88,6 +88,11 @@ pub fn coroutine_by_move_body_def_id<'tcx>(
 ) -> DefId {
     let body = tcx.mir_built(coroutine_def_id).borrow();
 
+    // If the typeck results are tainted, no need to make a by-ref body.
+    if body.tainted_by_errors.is_some() {
+        return coroutine_def_id.to_def_id();
+    }
+
     let Some(hir::CoroutineKind::Desugared(_, hir::CoroutineSource::Closure)) =
         tcx.coroutine_kind(coroutine_def_id)
     else {
@@ -98,7 +103,9 @@ pub fn coroutine_by_move_body_def_id<'tcx>(
     // the MIR body will be constructed well.
     let coroutine_ty = body.local_decls[ty::CAPTURE_STRUCT_LOCAL].ty;
 
-    let ty::Coroutine(_, args) = *coroutine_ty.kind() else { bug!("{body:#?}") };
+    let ty::Coroutine(_, args) = *coroutine_ty.kind() else {
+        bug!("tried to create by-move body of non-coroutine receiver");
+    };
     let args = args.as_coroutine();
 
     let coroutine_kind = args.kind_ty().to_opt_closure_kind().unwrap();
@@ -107,7 +114,7 @@ pub fn coroutine_by_move_body_def_id<'tcx>(
     let ty::CoroutineClosure(_, parent_args) =
         *tcx.type_of(parent_def_id).instantiate_identity().kind()
     else {
-        bug!();
+        bug!("coroutine's parent was not a coroutine-closure");
     };
     if parent_args.references_error() {
         return coroutine_def_id.to_def_id();
diff --git a/tests/ui/async-await/async-closures/tainted-body-2.rs b/tests/ui/async-await/async-closures/tainted-body-2.rs
new file mode 100644
index 00000000000..73c6bdc30a0
--- /dev/null
+++ b/tests/ui/async-await/async-closures/tainted-body-2.rs
@@ -0,0 +1,18 @@
+//@ edition: 2021
+
+#![feature(async_closure)]
+
+// Ensure that building a by-ref async closure body doesn't ICE when the parent
+// body is tainted.
+
+fn main() {
+    missing;
+    //~^ ERROR cannot find value `missing` in this scope
+
+    // We don't do numerical inference fallback when the body is tainted.
+    // This leads to writeback folding the type of the coroutine-closure
+    // into an error type, since its signature contains that numerical
+    // infer var.
+    let c = async |_| {};
+    c(1);
+}
diff --git a/tests/ui/async-await/async-closures/tainted-body-2.stderr b/tests/ui/async-await/async-closures/tainted-body-2.stderr
new file mode 100644
index 00000000000..798d47064d9
--- /dev/null
+++ b/tests/ui/async-await/async-closures/tainted-body-2.stderr
@@ -0,0 +1,9 @@
+error[E0425]: cannot find value `missing` in this scope
+  --> $DIR/tainted-body-2.rs:9:5
+   |
+LL |     missing;
+   |     ^^^^^^^ not found in this scope
+
+error: aborting due to 1 previous error
+
+For more information about this error, try `rustc --explain E0425`.