From 92cdff1483a21ecc42f3f35ab70425360df26a6c Mon Sep 17 00:00:00 2001
From: Michael Goulet <michael@errs.io>
Date: Mon, 29 Apr 2024 12:40:15 -0400
Subject: [PATCH] Don't ICE when we have leftover child captures due to
 ambiguous closure params

---
 compiler/rustc_hir_typeck/src/upvar.rs              |  1 +
 compiler/rustc_middle/src/ty/closure.rs             | 11 +++++++++--
 .../src/coroutine/by_move_body.rs                   |  1 +
 .../leftover-captures-due-to-ambiguity.rs           | 12 ++++++++++++
 .../leftover-captures-due-to-ambiguity.stderr       | 13 +++++++++++++
 5 files changed, 36 insertions(+), 2 deletions(-)
 create mode 100644 tests/ui/async-await/async-closures/leftover-captures-due-to-ambiguity.rs
 create mode 100644 tests/ui/async-await/async-closures/leftover-captures-due-to-ambiguity.stderr

diff --git a/compiler/rustc_hir_typeck/src/upvar.rs b/compiler/rustc_hir_typeck/src/upvar.rs
index 2bf4f51a8038a..9c34244f08b1e 100644
--- a/compiler/rustc_hir_typeck/src/upvar.rs
+++ b/compiler/rustc_hir_typeck/src/upvar.rs
@@ -384,6 +384,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             let tupled_upvars_ty_for_borrow = Ty::new_tup_from_iter(
                 self.tcx,
                 ty::analyze_coroutine_closure_captures(
+                    self.tcx,
                     typeck_results.closure_min_captures_flattened(closure_def_id),
                     typeck_results
                         .closure_min_captures_flattened(
diff --git a/compiler/rustc_middle/src/ty/closure.rs b/compiler/rustc_middle/src/ty/closure.rs
index b66c664e6ae1d..41d2c33b0338b 100644
--- a/compiler/rustc_middle/src/ty/closure.rs
+++ b/compiler/rustc_middle/src/ty/closure.rs
@@ -13,7 +13,7 @@ use rustc_hir::def_id::LocalDefId;
 use rustc_hir::HirId;
 use rustc_span::def_id::LocalDefIdMap;
 use rustc_span::symbol::Ident;
-use rustc_span::{Span, Symbol};
+use rustc_span::{Span, Symbol, DUMMY_SP};
 
 use super::TyCtxt;
 
@@ -418,6 +418,7 @@ impl BorrowKind {
 }
 
 pub fn analyze_coroutine_closure_captures<'a, 'tcx: 'a, T>(
+    tcx: TyCtxt<'tcx>,
     parent_captures: impl IntoIterator<Item = &'a CapturedPlace<'tcx>>,
     child_captures: impl IntoIterator<Item = &'a CapturedPlace<'tcx>>,
     mut for_each: impl FnMut((usize, &'a CapturedPlace<'tcx>), (usize, &'a CapturedPlace<'tcx>)) -> T,
@@ -466,7 +467,13 @@ pub fn analyze_coroutine_closure_captures<'a, 'tcx: 'a, T>(
                     "we captured {parent_capture:#?} but it was not used in the child coroutine?"
                 );
             }
-            assert_eq!(child_captures.next(), None, "leftover child captures?");
+
+            if let Some((_, capture)) = child_captures.next() {
+                tcx.dcx().span_delayed_bug(
+                    capture.info.path_expr_id.map_or(DUMMY_SP, |hir_id| tcx.hir().span(hir_id)),
+                    "leftover child captures: expecting an error",
+                );
+            }
         },
     )
 }
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 3d6c1a952041c..d7a9bd37f0ca4 100644
--- a/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs
+++ b/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs
@@ -125,6 +125,7 @@ impl<'tcx> MirPass<'tcx> for ByMoveBody {
             .len();
 
         let field_remapping: UnordMap<_, _> = ty::analyze_coroutine_closure_captures(
+            tcx,
             tcx.closure_captures(parent_def_id).iter().copied(),
             tcx.closure_captures(coroutine_def_id).iter().skip(num_args).copied(),
             |(parent_field_idx, parent_capture), (child_field_idx, child_capture)| {
diff --git a/tests/ui/async-await/async-closures/leftover-captures-due-to-ambiguity.rs b/tests/ui/async-await/async-closures/leftover-captures-due-to-ambiguity.rs
new file mode 100644
index 0000000000000..5a911faecbf44
--- /dev/null
+++ b/tests/ui/async-await/async-closures/leftover-captures-due-to-ambiguity.rs
@@ -0,0 +1,12 @@
+//@ edition: 2021
+
+#![feature(async_closure)]
+
+pub fn test(test: &()) {
+    async |unconstrained| {
+        //~^ ERROR type annotations needed
+        (test,)
+    };
+}
+
+fn main() {}
diff --git a/tests/ui/async-await/async-closures/leftover-captures-due-to-ambiguity.stderr b/tests/ui/async-await/async-closures/leftover-captures-due-to-ambiguity.stderr
new file mode 100644
index 0000000000000..9c57ca214c275
--- /dev/null
+++ b/tests/ui/async-await/async-closures/leftover-captures-due-to-ambiguity.stderr
@@ -0,0 +1,13 @@
+error[E0282]: type annotations needed
+  --> $DIR/leftover-captures-due-to-ambiguity.rs:6:27
+   |
+LL |       async |unconstrained| {
+   |  ___________________________^
+LL | |
+LL | |         (test,)
+LL | |     };
+   | |_____^ cannot infer type
+
+error: aborting due to 1 previous error
+
+For more information about this error, try `rustc --explain E0282`.