Store Option<Region> as value for RegionVid

4 files changed, 55 insertions, 46 deletions

diff --git a/compiler/rustc_infer/src/infer/canonical/canonicalizer.rs b/compiler/rustc_infer/src/infer/canonical/canonicalizer.rs
index c68705da413..27ac817463f 100644
--- a/compiler/rustc_infer/src/infer/canonical/canonicalizer.rs
+++ b/compiler/rustc_infer/src/infer/canonical/canonicalizer.rs
@@ -304,6 +304,15 @@ impl<'cx, 'tcx> TypeFolder<'tcx> for Canonicalizer<'cx, 'tcx> {
     }
 
     fn fold_region(&mut self, r: ty::Region<'tcx>) -> ty::Region<'tcx> {
+        let tcx = self.tcx();
+        let r = self
+            .infcx
+            .unwrap()
+            .inner
+            .borrow_mut()
+            .unwrap_region_constraints()
+            .opportunistic_resolve_region(tcx, r);
+
         match *r {
             ty::ReLateBound(index, ..) => {
                 if index >= self.binder_index {
@@ -313,22 +322,7 @@ impl<'cx, 'tcx> TypeFolder<'tcx> for Canonicalizer<'cx, 'tcx> {
                 }
             }
 
-            ty::ReVar(vid) => {
-                let resolved_vid = self
-                    .infcx
-                    .unwrap()
-                    .inner
-                    .borrow_mut()
-                    .unwrap_region_constraints()
-                    .opportunistic_resolve_var(vid);
-                debug!(
-                    "canonical: region var found with vid {:?}, \
-                     opportunistically resolved to {:?}",
-                    vid, r
-                );
-                let r = self.tcx.reuse_or_mk_region(r, ty::ReVar(resolved_vid));
-                self.canonicalize_region_mode.canonicalize_free_region(self, r)
-            }
+            ty::ReVar(_) => self.canonicalize_region_mode.canonicalize_free_region(self, r),
 
             ty::ReStatic
             | ty::ReEarlyBound(..)
diff --git a/compiler/rustc_infer/src/infer/region_constraints/mod.rs b/compiler/rustc_infer/src/infer/region_constraints/mod.rs
index abbcff30cda..dfbb82bd657 100644
--- a/compiler/rustc_infer/src/infer/region_constraints/mod.rs
+++ b/compiler/rustc_infer/src/infer/region_constraints/mod.rs
@@ -11,9 +11,9 @@ use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_data_structures::sync::Lrc;
 use rustc_data_structures::undo_log::UndoLogs;
 use rustc_data_structures::unify as ut;
-use rustc_data_structures::unify::UnifyKey;
 use rustc_hir::def_id::DefId;
 use rustc_index::vec::IndexVec;
+use rustc_middle::infer::unify_key::{RegionVidKey, UnifiedRegion};
 use rustc_middle::ty::ReStatic;
 use rustc_middle::ty::{self, Ty, TyCtxt};
 use rustc_middle::ty::{ReLateBound, ReVar};
@@ -47,13 +47,13 @@ pub struct RegionConstraintStorage<'tcx> {
 
     /// When we add a R1 == R2 constriant, we currently add (a) edges
     /// R1 <= R2 and R2 <= R1 and (b) we unify the two regions in this
-    /// table. You can then call `opportunistic_resolve_var` early
+    /// table. You can then call `opportunistic_resolve_region` early
     /// which will map R1 and R2 to some common region (i.e., either
     /// R1 or R2). This is important when fulfillment, dropck and other such
     /// code is iterating to a fixed point, because otherwise we sometimes
     /// would wind up with a fresh stream of region variables that have been
     /// equated but appear distinct.
-    pub(super) unification_table: ut::UnificationTableStorage<ty::RegionVid>,
+    pub(super) unification_table: ut::UnificationTableStorage<RegionVidKey<'tcx>>,
 
     /// a flag set to true when we perform any unifications; this is used
     /// to micro-optimize `take_and_reset_data`
@@ -406,8 +406,7 @@ impl<'tcx> RegionConstraintCollector<'_, 'tcx> {
         // `RegionConstraintData` contains the relationship here.
         if *any_unifications {
             *any_unifications = false;
-            self.unification_table()
-                .reset_unifications(|_| ());
+            self.unification_table().reset_unifications(|_| UnifiedRegion(None));
         }
 
         data
@@ -434,8 +433,8 @@ impl<'tcx> RegionConstraintCollector<'_, 'tcx> {
     ) -> RegionVid {
         let vid = self.var_infos.push(RegionVariableInfo { origin, universe });
 
-        let u_vid = self.unification_table().new_key(());
-        assert_eq!(vid, u_vid);
+        let u_vid = self.unification_table().new_key(UnifiedRegion(None));
+        assert_eq!(vid, u_vid.vid);
         self.undo_log.push(AddVar(vid));
         debug!("created new region variable {:?} in {:?} with origin {:?}", vid, universe, origin);
         vid
@@ -497,10 +496,18 @@ impl<'tcx> RegionConstraintCollector<'_, 'tcx> {
             self.make_subregion(origin.clone(), sub, sup);
             self.make_subregion(origin, sup, sub);
 
-            if let (ty::ReVar(sub), ty::ReVar(sup)) = (*sub, *sup) {
-                debug!("make_eqregion: uniying {:?} with {:?}", sub, sup);
-                self.unification_table().union(sub, sup);
-                self.any_unifications = true;
+            match (sub, sup) {
+                (&ty::ReVar(sub), &ty::ReVar(sup)) => {
+                    debug!("make_eqregion: unifying {:?} with {:?}", sub, sup);
+                    self.unification_table().union(sub, sup);
+                    self.any_unifications = true;
+                }
+                (&ty::ReVar(vid), value) | (value, &ty::ReVar(vid)) => {
+                    debug!("make_eqregion: unifying {:?} with {:?}", vid, value);
+                    self.unification_table().union_value(vid, UnifiedRegion(Some(value)));
+                    self.any_unifications = true;
+                }
+                (_, _) => {}
             }
         }
     }
@@ -616,8 +623,21 @@ impl<'tcx> RegionConstraintCollector<'_, 'tcx> {
         }
     }
 
-    pub fn opportunistic_resolve_var(&mut self, rid: RegionVid) -> ty::RegionVid {
-        self.unification_table().find(rid)
+    pub fn opportunistic_resolve_region(
+        &mut self,
+        tcx: TyCtxt<'tcx>,
+        region: ty::Region<'tcx>,
+    ) -> ty::Region<'tcx> {
+        match region {
+            ty::ReVar(rid) => {
+                let unified_region = self.unification_table().probe_value(*rid);
+                unified_region.0.unwrap_or_else(|| {
+                    let root = self.unification_table().find(*rid).vid;
+                    tcx.reuse_or_mk_region(region, ty::ReVar(root))
+                })
+            }
+            _ => region,
+        }
     }
 
     fn combine_map(&mut self, t: CombineMapType) -> &mut CombineMap<'tcx> {
@@ -672,8 +692,8 @@ impl<'tcx> RegionConstraintCollector<'_, 'tcx> {
         &self,
         value_count: usize,
     ) -> (Range<RegionVid>, Vec<RegionVariableOrigin>) {
-        let range = RegionVid::from_index(value_count as u32)
-            ..RegionVid::from_index(self.unification_table.len() as u32);
+        let range = RegionVid::from(value_count as u32)
+            ..RegionVid::from(self.unification_table.len() as u32);
         (
             range.clone(),
             (range.start.index()..range.end.index())
@@ -695,7 +715,7 @@ impl<'tcx> RegionConstraintCollector<'_, 'tcx> {
     }
 
     #[inline]
-    fn unification_table(&mut self) -> super::UnificationTable<'_, 'tcx, ty::RegionVid> {
+    fn unification_table(&mut self) -> super::UnificationTable<'_, 'tcx, RegionVidKey<'tcx>> {
         ut::UnificationTable::with_log(&mut self.storage.unification_table, self.undo_log)
     }
 }
diff --git a/compiler/rustc_infer/src/infer/resolve.rs b/compiler/rustc_infer/src/infer/resolve.rs
index 48b8ee17594..4eec65b3d96 100644
--- a/compiler/rustc_infer/src/infer/resolve.rs
+++ b/compiler/rustc_infer/src/infer/resolve.rs
@@ -84,18 +84,12 @@ impl<'a, 'tcx> TypeFolder<'tcx> for OpportunisticRegionResolver<'a, 'tcx> {
     }
 
     fn fold_region(&mut self, r: ty::Region<'tcx>) -> ty::Region<'tcx> {
-        match *r {
-            ty::ReVar(rid) => {
-                let resolved = self
-                    .infcx
-                    .inner
-                    .borrow_mut()
-                    .unwrap_region_constraints()
-                    .opportunistic_resolve_var(rid);
-                self.tcx().reuse_or_mk_region(r, ty::ReVar(resolved))
-            }
-            _ => r,
-        }
+        let tcx = self.tcx();
+        self.infcx
+            .inner
+            .borrow_mut()
+            .unwrap_region_constraints()
+            .opportunistic_resolve_region(tcx, r)
     }
 
     fn fold_const(&mut self, ct: &'tcx ty::Const<'tcx>) -> &'tcx ty::Const<'tcx> {
diff --git a/compiler/rustc_infer/src/infer/undo_log.rs b/compiler/rustc_infer/src/infer/undo_log.rs
index f41e872e004..5ad2519a93c 100644
--- a/compiler/rustc_infer/src/infer/undo_log.rs
+++ b/compiler/rustc_infer/src/infer/undo_log.rs
@@ -3,6 +3,7 @@ use std::marker::PhantomData;
 use rustc_data_structures::snapshot_vec as sv;
 use rustc_data_structures::undo_log::{Rollback, UndoLogs};
 use rustc_data_structures::unify as ut;
+use rustc_middle::infer::unify_key::RegionVidKey;
 use rustc_middle::ty;
 
 use crate::{
@@ -22,7 +23,7 @@ pub(crate) enum UndoLog<'tcx> {
     IntUnificationTable(sv::UndoLog<ut::Delegate<ty::IntVid>>),
     FloatUnificationTable(sv::UndoLog<ut::Delegate<ty::FloatVid>>),
     RegionConstraintCollector(region_constraints::UndoLog<'tcx>),
-    RegionUnificationTable(sv::UndoLog<ut::Delegate<ty::RegionVid>>),
+    RegionUnificationTable(sv::UndoLog<ut::Delegate<RegionVidKey<'tcx>>>),
     ProjectionCache(traits::UndoLog<'tcx>),
     PushRegionObligation,
 }
@@ -55,7 +56,7 @@ impl_from! {
 
     ConstUnificationTable(sv::UndoLog<ut::Delegate<ty::ConstVid<'tcx>>>),
 
-    RegionUnificationTable(sv::UndoLog<ut::Delegate<ty::RegionVid>>),
+    RegionUnificationTable(sv::UndoLog<ut::Delegate<RegionVidKey<'tcx>>>),
     ProjectionCache(traits::UndoLog<'tcx>),
 }