Don't ICE in ByMoveBody when coroutine is tainted

4 files changed, 53 insertions, 12 deletions

diff --git a/compiler/rustc_mir_build/src/build/mod.rs b/compiler/rustc_mir_build/src/build/mod.rs
index c24f2e9f53a..2b1a9fef360 100644
--- a/compiler/rustc_mir_build/src/build/mod.rs
+++ b/compiler/rustc_mir_build/src/build/mod.rs
@@ -675,16 +675,32 @@ fn construct_error(tcx: TyCtxt<'_>, def_id: LocalDefId, guar: ErrorGuaranteed) -
                         ))),
                     )
                 }
-                ty::CoroutineClosure(did, _args) => {
-                    // FIXME(async_closures): Recover the proper error signature
-                    let inputs = tcx
-                        .closure_user_provided_sig(did.expect_local())
-                        .value
-                        .skip_binder()
-                        .inputs();
-
-                    let err = Ty::new_error(tcx, guar);
-                    (inputs.iter().map(|_| err).collect(), err, None)
+                ty::CoroutineClosure(did, args) => {
+                    let args = args.as_coroutine_closure();
+                    let sig = tcx.liberate_late_bound_regions(
+                        def_id.to_def_id(),
+                        args.coroutine_closure_sig(),
+                    );
+                    let self_ty = match args.kind() {
+                        ty::ClosureKind::Fn => {
+                            Ty::new_imm_ref(tcx, tcx.lifetimes.re_erased, closure_ty)
+                        }
+                        ty::ClosureKind::FnMut => {
+                            Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, closure_ty)
+                        }
+                        ty::ClosureKind::FnOnce => closure_ty,
+                    };
+                    (
+                        [self_ty].into_iter().chain(sig.tupled_inputs_ty.tuple_fields()).collect(),
+                        sig.to_coroutine(
+                            tcx,
+                            args.parent_args(),
+                            args.kind_ty(),
+                            tcx.coroutine_for_closure(*did),
+                            Ty::new_error(tcx, guar),
+                        ),
+                        None,
+                    )
                 }
                 ty::Error(_) => (vec![closure_ty, closure_ty], closure_ty, None),
                 kind => {
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 fcd4715b9e8..e40f4520671 100644
--- a/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs
+++ b/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs
@@ -7,7 +7,7 @@ use rustc_data_structures::fx::FxIndexSet;
 use rustc_hir as hir;
 use rustc_middle::mir::visit::MutVisitor;
 use rustc_middle::mir::{self, dump_mir, MirPass};
-use rustc_middle::ty::{self, InstanceDef, Ty, TyCtxt};
+use rustc_middle::ty::{self, InstanceDef, Ty, TyCtxt, TypeVisitableExt};
 use rustc_target::abi::FieldIdx;
 
 pub struct ByMoveBody;
@@ -23,7 +23,10 @@ impl<'tcx> MirPass<'tcx> for ByMoveBody {
             return;
         };
         let coroutine_ty = body.local_decls[ty::CAPTURE_STRUCT_LOCAL].ty;
-        let ty::Coroutine(_, args) = *coroutine_ty.kind() else { bug!() };
+        if coroutine_ty.references_error() {
+            return;
+        }
+        let ty::Coroutine(_, args) = *coroutine_ty.kind() else { bug!("{body:#?}") };
 
         let coroutine_kind = args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap();
         if coroutine_kind == ty::ClosureKind::FnOnce {
diff --git a/tests/ui/async-await/async-closures/tainted-body.rs b/tests/ui/async-await/async-closures/tainted-body.rs
new file mode 100644
index 00000000000..62c28e7e585
--- /dev/null
+++ b/tests/ui/async-await/async-closures/tainted-body.rs
@@ -0,0 +1,13 @@
+// edition:2021
+
+#![feature(async_closure)]
+
+// Don't ICE in ByMove shim builder when MIR body is tainted by writeback errors
+
+fn main() {
+    let _ = async || {
+        used_fn();
+        //~^ ERROR cannot find function `used_fn` in this scope
+        0
+    };
+}
diff --git a/tests/ui/async-await/async-closures/tainted-body.stderr b/tests/ui/async-await/async-closures/tainted-body.stderr
new file mode 100644
index 00000000000..b06a896b81f
--- /dev/null
+++ b/tests/ui/async-await/async-closures/tainted-body.stderr
@@ -0,0 +1,9 @@
+error[E0425]: cannot find function `used_fn` in this scope
+  --> $DIR/tainted-body.rs:9:9
+   |
+LL |         used_fn();
+   |         ^^^^^^^ not found in this scope
+
+error: aborting due to 1 previous error
+
+For more information about this error, try `rustc --explain E0425`.