merge opaque types of nested bodies

6 files changed, 87 insertions, 48 deletions

diff --git a/compiler/rustc_borrowck/src/lib.rs b/compiler/rustc_borrowck/src/lib.rs
index ed95545dea4..84b7b8c6a2d 100644
--- a/compiler/rustc_borrowck/src/lib.rs
+++ b/compiler/rustc_borrowck/src/lib.rs
@@ -73,6 +73,7 @@ mod def_use;
 mod diagnostics;
 mod member_constraints;
 mod nll;
+mod opaque_types;
 mod path_utils;
 mod place_ext;
 mod places_conflict;
@@ -192,7 +193,7 @@ fn do_mir_borrowck<'tcx>(
     // Compute non-lexical lifetimes.
     let nll::NllOutput {
         regioncx,
-        opaque_type_values,
+        concrete_opaque_types,
         polonius_input,
         polonius_output,
         opt_closure_req,
@@ -222,7 +223,7 @@ fn do_mir_borrowck<'tcx>(
         body,
         &regioncx,
         &opt_closure_req,
-        &opaque_type_values,
+        &concrete_opaque_types,
         diags_buffer,
     );
 
@@ -357,7 +358,7 @@ fn do_mir_borrowck<'tcx>(
     let tainted_by_errors = mbcx.emit_errors();
 
     let result = BorrowCheckResult {
-        concrete_opaque_types: opaque_type_values,
+        concrete_opaque_types: concrete_opaque_types.into_inner(),
         closure_requirements: opt_closure_req,
         used_mut_upvars: mbcx.used_mut_upvars,
         tainted_by_errors,
diff --git a/compiler/rustc_borrowck/src/nll.rs b/compiler/rustc_borrowck/src/nll.rs
index 2031579dfd5..d0bd364425a 100644
--- a/compiler/rustc_borrowck/src/nll.rs
+++ b/compiler/rustc_borrowck/src/nll.rs
@@ -6,8 +6,6 @@ use std::str::FromStr;
 use std::{env, io};
 
 use polonius_engine::{Algorithm, Output};
-use rustc_data_structures::fx::FxIndexMap;
-use rustc_hir::def_id::LocalDefId;
 use rustc_index::IndexSlice;
 use rustc_middle::mir::pretty::{PrettyPrintMirOptions, dump_mir_with_options};
 use rustc_middle::mir::{
@@ -15,7 +13,7 @@ use rustc_middle::mir::{
     dump_enabled, dump_mir,
 };
 use rustc_middle::ty::print::with_no_trimmed_paths;
-use rustc_middle::ty::{self, OpaqueHiddenType, TyCtxt};
+use rustc_middle::ty::{self, TyCtxt};
 use rustc_mir_dataflow::ResultsCursor;
 use rustc_mir_dataflow::impls::MaybeInitializedPlaces;
 use rustc_mir_dataflow::move_paths::MoveData;
@@ -27,6 +25,7 @@ use tracing::{debug, instrument};
 use crate::borrow_set::BorrowSet;
 use crate::consumers::ConsumerOptions;
 use crate::diagnostics::{BorrowckDiagnosticsBuffer, RegionErrors};
+use crate::opaque_types::ConcreteOpaqueTypes;
 use crate::polonius::PoloniusDiagnosticsContext;
 use crate::polonius::legacy::{
     PoloniusFacts, PoloniusFactsExt, PoloniusLocationTable, PoloniusOutput,
@@ -40,7 +39,7 @@ use crate::{BorrowckInferCtxt, polonius, renumber};
 /// closure requirements to propagate, and any generated errors.
 pub(crate) struct NllOutput<'tcx> {
     pub regioncx: RegionInferenceContext<'tcx>,
-    pub opaque_type_values: FxIndexMap<LocalDefId, OpaqueHiddenType<'tcx>>,
+    pub concrete_opaque_types: ConcreteOpaqueTypes<'tcx>,
     pub polonius_input: Option<Box<PoloniusFacts>>,
     pub polonius_output: Option<Box<PoloniusOutput>>,
     pub opt_closure_req: Option<ClosureRegionRequirements<'tcx>>,
@@ -99,6 +98,8 @@ pub(crate) fn compute_regions<'a, 'tcx>(
 
     let location_map = Rc::new(DenseLocationMap::new(body));
 
+    let mut concrete_opaque_types = ConcreteOpaqueTypes::default();
+
     // Run the MIR type-checker.
     let MirTypeckResults {
         constraints,
@@ -116,6 +117,7 @@ pub(crate) fn compute_regions<'a, 'tcx>(
         flow_inits,
         move_data,
         Rc::clone(&location_map),
+        &mut concrete_opaque_types,
     );
 
     // Create the region inference context, taking ownership of the
@@ -180,11 +182,11 @@ pub(crate) fn compute_regions<'a, 'tcx>(
         infcx.set_tainted_by_errors(guar);
     }
 
-    let remapped_opaque_tys = regioncx.infer_opaque_types(infcx, opaque_type_values);
+    regioncx.infer_opaque_types(infcx, opaque_type_values, &mut concrete_opaque_types);
 
     NllOutput {
         regioncx,
-        opaque_type_values: remapped_opaque_tys,
+        concrete_opaque_types,
         polonius_input: polonius_facts.map(Box::new),
         polonius_output,
         opt_closure_req: closure_region_requirements,
@@ -300,7 +302,7 @@ pub(super) fn dump_annotation<'tcx, 'infcx>(
     body: &Body<'tcx>,
     regioncx: &RegionInferenceContext<'tcx>,
     closure_region_requirements: &Option<ClosureRegionRequirements<'tcx>>,
-    opaque_type_values: &FxIndexMap<LocalDefId, OpaqueHiddenType<'tcx>>,
+    concrete_opaque_types: &ConcreteOpaqueTypes<'tcx>,
     diagnostics_buffer: &mut BorrowckDiagnosticsBuffer<'infcx, 'tcx>,
 ) {
     let tcx = infcx.tcx;
@@ -343,8 +345,8 @@ pub(super) fn dump_annotation<'tcx, 'infcx>(
         err
     };
 
-    if !opaque_type_values.is_empty() {
-        err.note(format!("Inferred opaque type values:\n{opaque_type_values:#?}"));
+    if !concrete_opaque_types.is_empty() {
+        err.note(format!("Inferred opaque type values:\n{concrete_opaque_types:#?}"));
     }
 
     diagnostics_buffer.buffer_non_error(err);
diff --git a/compiler/rustc_borrowck/src/opaque_types.rs b/compiler/rustc_borrowck/src/opaque_types.rs
new file mode 100644
index 00000000000..5c78814abdd
--- /dev/null
+++ b/compiler/rustc_borrowck/src/opaque_types.rs
@@ -0,0 +1,55 @@
+use rustc_data_structures::fx::FxIndexMap;
+use rustc_hir::def_id::LocalDefId;
+use rustc_middle::ty::{OpaqueHiddenType, Ty, TyCtxt};
+
+#[derive(Debug, Default)]
+pub(super) struct ConcreteOpaqueTypes<'tcx> {
+    concrete_opaque_types: FxIndexMap<LocalDefId, OpaqueHiddenType<'tcx>>,
+}
+
+impl<'tcx> ConcreteOpaqueTypes<'tcx> {
+    pub(super) fn is_empty(&self) -> bool {
+        self.concrete_opaque_types.is_empty()
+    }
+
+    pub(super) fn into_inner(self) -> FxIndexMap<LocalDefId, OpaqueHiddenType<'tcx>> {
+        self.concrete_opaque_types
+    }
+
+    /// Insert an opaque type into the list of opaque types defined by this function
+    /// after mapping the hidden type to the generic parameters of the opaque type
+    /// definition.
+    pub(super) fn insert(
+        &mut self,
+        tcx: TyCtxt<'tcx>,
+        def_id: LocalDefId,
+        hidden_ty: OpaqueHiddenType<'tcx>,
+    ) {
+        // Sometimes two opaque types are the same only after we remap the generic parameters
+        // back to the opaque type definition. E.g. we may have `OpaqueType<X, Y>` mapped to
+        // `(X, Y)` and `OpaqueType<Y, X>` mapped to `(Y, X)`, and those are the same, but we
+        // only know that once we convert the generic parameters to those of the opaque type.
+        if let Some(prev) = self.concrete_opaque_types.get_mut(&def_id) {
+            if prev.ty != hidden_ty.ty {
+                let (Ok(guar) | Err(guar)) =
+                    prev.build_mismatch_error(&hidden_ty, tcx).map(|d| d.emit());
+                prev.ty = Ty::new_error(tcx, guar);
+            }
+            // Pick a better span if there is one.
+            // FIXME(oli-obk): collect multiple spans for better diagnostics down the road.
+            prev.span = prev.span.substitute_dummy(hidden_ty.span);
+        } else {
+            self.concrete_opaque_types.insert(def_id, hidden_ty);
+        }
+    }
+
+    pub(super) fn extend_from_nested_body(
+        &mut self,
+        tcx: TyCtxt<'tcx>,
+        nested_body: &FxIndexMap<LocalDefId, OpaqueHiddenType<'tcx>>,
+    ) {
+        for (&def_id, &hidden_ty) in nested_body {
+            self.insert(tcx, def_id, hidden_ty);
+        }
+    }
+}
diff --git a/compiler/rustc_borrowck/src/region_infer/opaque_types.rs b/compiler/rustc_borrowck/src/region_infer/opaque_types.rs
index 04839448596..ca8b9fb4e9d 100644
--- a/compiler/rustc_borrowck/src/region_infer/opaque_types.rs
+++ b/compiler/rustc_borrowck/src/region_infer/opaque_types.rs
@@ -15,6 +15,7 @@ use rustc_trait_selection::traits::ObligationCtxt;
 use tracing::{debug, instrument};
 
 use super::RegionInferenceContext;
+use crate::opaque_types::ConcreteOpaqueTypes;
 use crate::session_diagnostics::{LifetimeMismatchOpaqueParam, NonGenericOpaqueTypeParam};
 use crate::universal_regions::RegionClassification;
 
@@ -67,8 +68,8 @@ impl<'tcx> RegionInferenceContext<'tcx> {
         &self,
         infcx: &InferCtxt<'tcx>,
         opaque_ty_decls: FxIndexMap<OpaqueTypeKey<'tcx>, OpaqueHiddenType<'tcx>>,
-    ) -> FxIndexMap<LocalDefId, OpaqueHiddenType<'tcx>> {
-        let mut result: FxIndexMap<LocalDefId, OpaqueHiddenType<'tcx>> = FxIndexMap::default();
+        concrete_opaque_types: &mut ConcreteOpaqueTypes<'tcx>,
+    ) {
         let mut decls_modulo_regions: FxIndexMap<OpaqueTypeKey<'tcx>, (OpaqueTypeKey<'tcx>, Span)> =
             FxIndexMap::default();
 
@@ -143,33 +144,12 @@ impl<'tcx> RegionInferenceContext<'tcx> {
                     continue;
                 }
             }
-            // Sometimes two opaque types are the same only after we remap the generic parameters
-            // back to the opaque type definition. E.g. we may have `OpaqueType<X, Y>` mapped to
-            // `(X, Y)` and `OpaqueType<Y, X>` mapped to `(Y, X)`, and those are the same, but we
-            // only know that once we convert the generic parameters to those of the opaque type.
-            if let Some(prev) = result.get_mut(&opaque_type_key.def_id) {
-                if prev.ty != ty {
-                    let guar = ty.error_reported().err().unwrap_or_else(|| {
-                        let (Ok(e) | Err(e)) = prev
-                            .build_mismatch_error(
-                                &OpaqueHiddenType { ty, span: concrete_type.span },
-                                infcx.tcx,
-                            )
-                            .map(|d| d.emit());
-                        e
-                    });
-                    prev.ty = Ty::new_error(infcx.tcx, guar);
-                }
-                // Pick a better span if there is one.
-                // FIXME(oli-obk): collect multiple spans for better diagnostics down the road.
-                prev.span = prev.span.substitute_dummy(concrete_type.span);
-            } else {
-                result.insert(
-                    opaque_type_key.def_id,
-                    OpaqueHiddenType { ty, span: concrete_type.span },
-                );
-            }
 
+            concrete_opaque_types.insert(
+                infcx.tcx,
+                opaque_type_key.def_id,
+                OpaqueHiddenType { ty, span: concrete_type.span },
+            );
             // Check that all opaque types have the same region parameters if they have the same
             // non-region parameters. This is necessary because within the new solver we perform
             // various query operations modulo regions, and thus could unsoundly select some impls
@@ -193,7 +173,6 @@ impl<'tcx> RegionInferenceContext<'tcx> {
                 });
             }
         }
-        result
     }
 
     /// Map the regions in the type to named regions. This is similar to what
diff --git a/compiler/rustc_borrowck/src/type_check/mod.rs b/compiler/rustc_borrowck/src/type_check/mod.rs
index 0fda3e31690..f6144a25938 100644
--- a/compiler/rustc_borrowck/src/type_check/mod.rs
+++ b/compiler/rustc_borrowck/src/type_check/mod.rs
@@ -45,6 +45,7 @@ use crate::borrow_set::BorrowSet;
 use crate::constraints::{OutlivesConstraint, OutlivesConstraintSet};
 use crate::diagnostics::UniverseInfo;
 use crate::member_constraints::MemberConstraintSet;
+use crate::opaque_types::ConcreteOpaqueTypes;
 use crate::polonius::legacy::{PoloniusFacts, PoloniusLocationTable};
 use crate::polonius::{PoloniusContext, PoloniusLivenessContext};
 use crate::region_infer::TypeTest;
@@ -111,6 +112,7 @@ pub(crate) fn type_check<'a, 'tcx>(
     flow_inits: ResultsCursor<'a, 'tcx, MaybeInitializedPlaces<'a, 'tcx>>,
     move_data: &MoveData<'tcx>,
     location_map: Rc<DenseLocationMap>,
+    concrete_opaque_types: &mut ConcreteOpaqueTypes<'tcx>,
 ) -> MirTypeckResults<'tcx> {
     let implicit_region_bound = ty::Region::new_var(infcx.tcx, universal_regions.fr_fn_body);
     let mut constraints = MirTypeckRegionConstraints {
@@ -165,6 +167,7 @@ pub(crate) fn type_check<'a, 'tcx>(
         polonius_facts,
         borrow_set,
         constraints: &mut constraints,
+        concrete_opaque_types,
         polonius_liveness,
     };
 
@@ -230,6 +233,7 @@ struct TypeChecker<'a, 'tcx> {
     polonius_facts: &'a mut Option<PoloniusFacts>,
     borrow_set: &'a BorrowSet<'tcx>,
     constraints: &'a mut MirTypeckRegionConstraints<'tcx>,
+    concrete_opaque_types: &'a mut ConcreteOpaqueTypes<'tcx>,
     /// When using `-Zpolonius=next`, the liveness helper data used to create polonius constraints.
     polonius_liveness: Option<PoloniusLivenessContext>,
 }
@@ -2499,7 +2503,11 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
         args: GenericArgsRef<'tcx>,
         locations: Locations,
     ) -> ty::InstantiatedPredicates<'tcx> {
-        if let Some(closure_requirements) = &tcx.mir_borrowck(def_id).closure_requirements {
+        let closure_borrowck_results = tcx.mir_borrowck(def_id);
+        self.concrete_opaque_types
+            .extend_from_nested_body(tcx, &closure_borrowck_results.concrete_opaque_types);
+
+        if let Some(closure_requirements) = &closure_borrowck_results.closure_requirements {
             constraint_conversion::ConstraintConversion::new(
                 self.infcx,
                 self.universal_regions,
diff --git a/compiler/rustc_hir_analysis/src/collect/type_of/opaque.rs b/compiler/rustc_hir_analysis/src/collect/type_of/opaque.rs
index 075dc590b99..3dec1e286b4 100644
--- a/compiler/rustc_hir_analysis/src/collect/type_of/opaque.rs
+++ b/compiler/rustc_hir_analysis/src/collect/type_of/opaque.rs
@@ -258,9 +258,6 @@ impl<'tcx> intravisit::Visitor<'tcx> for TaitConstraintLocator<'tcx> {
         self.tcx
     }
     fn visit_expr(&mut self, ex: &'tcx Expr<'tcx>) {
-        if let hir::ExprKind::Closure(closure) = ex.kind {
-            self.check(closure.def_id);
-        }
         intravisit::walk_expr(self, ex);
     }
     fn visit_item(&mut self, it: &'tcx Item<'tcx>) {
@@ -389,9 +386,6 @@ impl<'tcx> intravisit::Visitor<'tcx> for RpitConstraintChecker<'tcx> {
         self.tcx
     }
     fn visit_expr(&mut self, ex: &'tcx Expr<'tcx>) {
-        if let hir::ExprKind::Closure(closure) = ex.kind {
-            self.check(closure.def_id);
-        }
         intravisit::walk_expr(self, ex);
     }
     fn visit_item(&mut self, it: &'tcx Item<'tcx>) {