Replace tuple of infer vars for upvar_tys with single infer var

This commit allows us to decide the number of captures required after
completing capture ananysis, which is required as part of implementing
RFC-2229.

Co-authored-by: Aman Arora <me@aman-arora.com>
Co-authored-by: Jenny Wills <wills.jenniferg@gmail.com>

5 files changed, 71 insertions, 21 deletions

diff --git a/compiler/rustc_trait_selection/src/opaque_types.rs b/compiler/rustc_trait_selection/src/opaque_types.rs
index 610c6fd7e35..087f6a0deec 100644
--- a/compiler/rustc_trait_selection/src/opaque_types.rs
+++ b/compiler/rustc_trait_selection/src/opaque_types.rs
@@ -441,6 +441,7 @@ impl<'a, 'tcx> InferCtxtExt<'tcx> for InferCtxt<'a, 'tcx> {
 
             for required_region in required_region_bounds {
                 concrete_ty.visit_with(&mut ConstrainOpaqueTypeRegionVisitor {
+                    infcx: self,
                     op: |r| self.sub_regions(infer::CallReturn(span), required_region, r),
                 });
             }
@@ -509,6 +510,7 @@ impl<'a, 'tcx> InferCtxtExt<'tcx> for InferCtxt<'a, 'tcx> {
             }
         }
         concrete_ty.visit_with(&mut ConstrainOpaqueTypeRegionVisitor {
+            infcx: self,
             op: |r| self.sub_regions(infer::CallReturn(span), least_region, r),
         });
     }
@@ -543,6 +545,7 @@ impl<'a, 'tcx> InferCtxtExt<'tcx> for InferCtxt<'a, 'tcx> {
         );
 
         concrete_ty.visit_with(&mut ConstrainOpaqueTypeRegionVisitor {
+            infcx: self,
             op: |r| {
                 self.member_constraint(
                     opaque_type_def_id,
@@ -683,11 +686,12 @@ impl<'a, 'tcx> InferCtxtExt<'tcx> for InferCtxt<'a, 'tcx> {
 //
 // We ignore any type parameters because impl trait values are assumed to
 // capture all the in-scope type parameters.
-struct ConstrainOpaqueTypeRegionVisitor<OP> {
+struct ConstrainOpaqueTypeRegionVisitor<'cx, 'tcx, OP> {
+    infcx: &'cx InferCtxt<'cx, 'tcx>,
     op: OP,
 }
 
-impl<'tcx, OP> TypeVisitor<'tcx> for ConstrainOpaqueTypeRegionVisitor<OP>
+impl<'cx, 'tcx, OP> TypeVisitor<'tcx> for ConstrainOpaqueTypeRegionVisitor<'cx, 'tcx, OP>
 where
     OP: FnMut(ty::Region<'tcx>),
 {
@@ -717,24 +721,36 @@ where
             ty::Closure(_, ref substs) => {
                 // Skip lifetime parameters of the enclosing item(s)
 
-                for upvar_ty in substs.as_closure().upvar_tys() {
-                    upvar_ty.visit_with(self);
-                }
+                let ty = self.infcx.shallow_resolve(substs.as_closure().tupled_upvars_ty());
+                if let ty::Infer(ty::TyVar(_)) = ty.kind() {
+                    // Not yet resolved.
+                    ty.super_visit_with(self);
+                } else {
+                    for upvar_ty in substs.as_closure().upvar_tys() {
+                        upvar_ty.visit_with(self);
+                    }
 
-                substs.as_closure().sig_as_fn_ptr_ty().visit_with(self);
+                    substs.as_closure().sig_as_fn_ptr_ty().visit_with(self);
+                }
             }
 
             ty::Generator(_, ref substs, _) => {
                 // Skip lifetime parameters of the enclosing item(s)
                 // Also skip the witness type, because that has no free regions.
 
-                for upvar_ty in substs.as_generator().upvar_tys() {
-                    upvar_ty.visit_with(self);
-                }
+                let ty = self.infcx.shallow_resolve(substs.as_generator().tupled_upvars_ty());
+                if let ty::Infer(ty::TyVar(_)) = ty.kind() {
+                    // Not yet resolved.
+                    ty.super_visit_with(self);
+                } else {
+                    for upvar_ty in substs.as_generator().upvar_tys() {
+                        upvar_ty.visit_with(self);
+                    }
 
-                substs.as_generator().return_ty().visit_with(self);
-                substs.as_generator().yield_ty().visit_with(self);
-                substs.as_generator().resume_ty().visit_with(self);
+                    substs.as_generator().return_ty().visit_with(self);
+                    substs.as_generator().yield_ty().visit_with(self);
+                    substs.as_generator().resume_ty().visit_with(self);
+                }
             }
             _ => {
                 ty.super_visit_with(self);
diff --git a/compiler/rustc_trait_selection/src/traits/error_reporting/suggestions.rs b/compiler/rustc_trait_selection/src/traits/error_reporting/suggestions.rs
index 967374ffdc2..5908e73d483 100644
--- a/compiler/rustc_trait_selection/src/traits/error_reporting/suggestions.rs
+++ b/compiler/rustc_trait_selection/src/traits/error_reporting/suggestions.rs
@@ -1307,6 +1307,9 @@ impl<'a, 'tcx> InferCtxtExt<'tcx> for InferCtxt<'a, 'tcx> {
         let mut generator = None;
         let mut outer_generator = None;
         let mut next_code = Some(&obligation.cause.code);
+
+        let mut seen_upvar_tys_infer_tuple = false;
+
         while let Some(code) = next_code {
             debug!("maybe_note_obligation_cause_for_async_await: code={:?}", code);
             match code {
@@ -1327,6 +1330,13 @@ impl<'a, 'tcx> InferCtxtExt<'tcx> for InferCtxt<'a, 'tcx> {
                             outer_generator = Some(did);
                         }
                         ty::GeneratorWitness(..) => {}
+                        ty::Tuple(_) if !seen_upvar_tys_infer_tuple => {
+                            // By introducing a tuple of upvar types into the chain of obligations
+                            // of a generator, the first non-generator item is now the tuple itself,
+                            // we shall ignore this.
+
+                            seen_upvar_tys_infer_tuple = true;
+                        }
                         _ if generator.is_none() => {
                             trait_ref = Some(derived_obligation.parent_trait_ref.skip_binder());
                             target_ty = Some(ty);
diff --git a/compiler/rustc_trait_selection/src/traits/query/dropck_outlives.rs b/compiler/rustc_trait_selection/src/traits/query/dropck_outlives.rs
index 424b3bd67ff..8212823a6db 100644
--- a/compiler/rustc_trait_selection/src/traits/query/dropck_outlives.rs
+++ b/compiler/rustc_trait_selection/src/traits/query/dropck_outlives.rs
@@ -110,7 +110,7 @@ pub fn trivial_dropck_outlives<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> bool {
         // check if *any* of those are trivial.
         ty::Tuple(ref tys) => tys.iter().all(|t| trivial_dropck_outlives(tcx, t.expect_ty())),
         ty::Closure(_, ref substs) => {
-            substs.as_closure().upvar_tys().all(|t| trivial_dropck_outlives(tcx, t))
+            trivial_dropck_outlives(tcx, substs.as_closure().tupled_upvars_ty())
         }
 
         ty::Adt(def, _) => {
diff --git a/compiler/rustc_trait_selection/src/traits/select/mod.rs b/compiler/rustc_trait_selection/src/traits/select/mod.rs
index a142ba58a69..efe01191e58 100644
--- a/compiler/rustc_trait_selection/src/traits/select/mod.rs
+++ b/compiler/rustc_trait_selection/src/traits/select/mod.rs
@@ -1631,7 +1631,13 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
 
             ty::Closure(_, substs) => {
                 // (*) binder moved here
-                Where(ty::Binder::bind(substs.as_closure().upvar_tys().collect()))
+                let ty = self.infcx.shallow_resolve(substs.as_closure().tupled_upvars_ty());
+                if let ty::Infer(ty::TyVar(_)) = ty.kind() {
+                    // Not yet resolved.
+                    Ambiguous
+                } else {
+                    Where(ty::Binder::bind(substs.as_closure().upvar_tys().collect()))
+                }
             }
 
             ty::Adt(..) | ty::Projection(..) | ty::Param(..) | ty::Opaque(..) => {
@@ -1700,11 +1706,31 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
                 tys.iter().map(|k| k.expect_ty()).collect()
             }
 
-            ty::Closure(_, ref substs) => substs.as_closure().upvar_tys().collect(),
+            ty::Closure(_, ref substs) => {
+                let ty = self.infcx.shallow_resolve(substs.as_closure().tupled_upvars_ty());
+                if let ty::Infer(ty::TyVar(_)) = ty.kind() {
+                    // The inference variable will be replaced by a tuple once capture analysis
+                    // completes. If the tuple meets a bound, so do all the elements within it.
+                    vec![ty]
+                } else {
+                    substs.as_closure().upvar_tys().collect()
+                }
+            }
 
             ty::Generator(_, ref substs, _) => {
-                let witness = substs.as_generator().witness();
-                substs.as_generator().upvar_tys().chain(iter::once(witness)).collect()
+                let upvar_tys_resolved =
+                    self.infcx.shallow_resolve(substs.as_generator().tupled_upvars_ty());
+
+                if let ty::Infer(ty::TyVar(_)) = upvar_tys_resolved.kind() {
+                    // The inference variable will be replaced by a tuple once capture analysis
+                    // completes, if the tuple meets a bound, so do all the elements within it.
+                    let witness_resolved =
+                        self.infcx.shallow_resolve(substs.as_generator().witness());
+                    vec![upvar_tys_resolved, witness_resolved]
+                } else {
+                    let witness = substs.as_generator().witness();
+                    substs.as_generator().upvar_tys().chain(iter::once(witness)).collect()
+                }
             }
 
             ty::GeneratorWitness(types) => {
diff --git a/compiler/rustc_trait_selection/src/traits/wf.rs b/compiler/rustc_trait_selection/src/traits/wf.rs
index d66bfd48206..496dff6c5b2 100644
--- a/compiler/rustc_trait_selection/src/traits/wf.rs
+++ b/compiler/rustc_trait_selection/src/traits/wf.rs
@@ -592,10 +592,8 @@ impl<'a, 'tcx> WfPredicates<'a, 'tcx> {
                     // anyway, except via auto trait matching (which
                     // only inspects the upvar types).
                     walker.skip_current_subtree(); // subtree handled below
-                    for upvar_ty in substs.as_closure().upvar_tys() {
-                        // FIXME(eddyb) add the type to `walker` instead of recursing.
-                        self.compute(upvar_ty.into());
-                    }
+                    // FIXME(eddyb) add the type to `walker` instead of recursing.
+                    self.compute(substs.as_closure().tupled_upvars_ty().into());
                 }
 
                 ty::FnPtr(_) => {