region unification update universe of region vars

5 files changed, 123 insertions, 96 deletions

diff --git a/compiler/rustc_infer/src/infer/canonical/canonicalizer.rs b/compiler/rustc_infer/src/infer/canonical/canonicalizer.rs
index 99882a42abc..9f70fee993d 100644
--- a/compiler/rustc_infer/src/infer/canonical/canonicalizer.rs
+++ b/compiler/rustc_infer/src/infer/canonical/canonicalizer.rs
@@ -175,8 +175,12 @@ impl CanonicalizeMode for CanonicalizeQueryResponse {
             ),
 
             ty::ReVar(vid) => {
-                let universe =
-                    infcx.inner.borrow_mut().unwrap_region_constraints().var_universe(vid);
+                let universe = infcx
+                    .inner
+                    .borrow_mut()
+                    .unwrap_region_constraints()
+                    .probe_value(vid)
+                    .unwrap_err();
                 canonicalizer.canonical_var_for_region(
                     CanonicalVarInfo { kind: CanonicalVarKind::Region(universe) },
                     r,
diff --git a/compiler/rustc_infer/src/infer/mod.rs b/compiler/rustc_infer/src/infer/mod.rs
index 8fc71671b27..a30799fc2b4 100644
--- a/compiler/rustc_infer/src/infer/mod.rs
+++ b/compiler/rustc_infer/src/infer/mod.rs
@@ -374,7 +374,10 @@ impl<'tcx> ty::InferCtxtLike for InferCtxt<'tcx> {
     }
 
     fn universe_of_lt(&self, lt: ty::RegionVid) -> Option<ty::UniverseIndex> {
-        Some(self.universe_of_region_vid(lt))
+        match self.inner.borrow_mut().unwrap_region_constraints().probe_value(lt) {
+            Err(universe) => Some(universe),
+            Ok(_) => None,
+        }
     }
 
     fn root_ty_var(&self, vid: TyVid) -> TyVid {
@@ -1155,11 +1158,6 @@ impl<'tcx> InferCtxt<'tcx> {
         self.inner.borrow_mut().unwrap_region_constraints().universe(r)
     }
 
-    /// Return the universe that the region variable `r` was created in.
-    pub fn universe_of_region_vid(&self, vid: ty::RegionVid) -> ty::UniverseIndex {
-        self.inner.borrow_mut().unwrap_region_constraints().var_universe(vid)
-    }
-
     /// Number of region variables created so far.
     pub fn num_region_vars(&self) -> usize {
         self.inner.borrow_mut().unwrap_region_constraints().num_region_vars()
diff --git a/compiler/rustc_infer/src/infer/region_constraints/leak_check.rs b/compiler/rustc_infer/src/infer/region_constraints/leak_check.rs
index b2bcbbf2e53..9e2f3a10b6a 100644
--- a/compiler/rustc_infer/src/infer/region_constraints/leak_check.rs
+++ b/compiler/rustc_infer/src/infer/region_constraints/leak_check.rs
@@ -90,11 +90,11 @@ impl<'tcx> RegionConstraintCollector<'_, 'tcx> {
     }
 }
 
-struct LeakCheck<'me, 'tcx> {
+struct LeakCheck<'a, 'b, 'tcx> {
     tcx: TyCtxt<'tcx>,
     outer_universe: ty::UniverseIndex,
-    mini_graph: &'me MiniGraph<'tcx>,
-    rcc: &'me RegionConstraintCollector<'me, 'tcx>,
+    mini_graph: &'a MiniGraph<'tcx>,
+    rcc: &'a mut RegionConstraintCollector<'b, 'tcx>,
 
     // Initially, for each SCC S, stores a placeholder `P` such that `S = P`
     // must hold.
@@ -117,13 +117,13 @@ struct LeakCheck<'me, 'tcx> {
     scc_universes: IndexVec<LeakCheckScc, SccUniverse<'tcx>>,
 }
 
-impl<'me, 'tcx> LeakCheck<'me, 'tcx> {
+impl<'a, 'b, 'tcx> LeakCheck<'a, 'b, 'tcx> {
     fn new(
         tcx: TyCtxt<'tcx>,
         outer_universe: ty::UniverseIndex,
         max_universe: ty::UniverseIndex,
-        mini_graph: &'me MiniGraph<'tcx>,
-        rcc: &'me RegionConstraintCollector<'me, 'tcx>,
+        mini_graph: &'a MiniGraph<'tcx>,
+        rcc: &'a mut RegionConstraintCollector<'b, 'tcx>,
     ) -> Self {
         let dummy_scc_universe = SccUniverse { universe: max_universe, region: None };
         Self {
diff --git a/compiler/rustc_infer/src/infer/region_constraints/mod.rs b/compiler/rustc_infer/src/infer/region_constraints/mod.rs
index 71adf526097..ee97dd36807 100644
--- a/compiler/rustc_infer/src/infer/region_constraints/mod.rs
+++ b/compiler/rustc_infer/src/infer/region_constraints/mod.rs
@@ -8,12 +8,11 @@ use super::{
 };
 
 use rustc_data_structures::fx::FxHashMap;
-use rustc_data_structures::intern::Interned;
 use rustc_data_structures::sync::Lrc;
 use rustc_data_structures::undo_log::UndoLogs;
 use rustc_data_structures::unify as ut;
 use rustc_index::IndexVec;
-use rustc_middle::infer::unify_key::{RegionVidKey, UnifiedRegion};
+use rustc_middle::infer::unify_key::{RegionVariableValue, RegionVidKey};
 use rustc_middle::ty::ReStatic;
 use rustc_middle::ty::{self, Ty, TyCtxt};
 use rustc_middle::ty::{ReBound, ReVar};
@@ -292,6 +291,18 @@ type CombineMap<'tcx> = FxHashMap<TwoRegions<'tcx>, RegionVid>;
 #[derive(Debug, Clone, Copy)]
 pub struct RegionVariableInfo {
     pub origin: RegionVariableOrigin,
+    // FIXME: This is only necessary for `fn take_and_reset_data` and
+    // `lexical_region_resolve`. We should rework `lexical_region_resolve`
+    // in the near/medium future anyways and could move the unverse info
+    // for `fn take_and_reset_data` into a separate table which is
+    // only populated when needed.
+    //
+    // For both of these cases it is fine that this can diverge from the
+    // actual universe of the variable, which is directly stored in the
+    // unification table for unknown region variables. At some point we could
+    // stop emitting bidirectional outlives constraints if equate succeeds.
+    // This would be currently unsound as it would cause us to drop the universe
+    // changes in `lexical_region_resolve`.
     pub universe: ty::UniverseIndex,
 }
 
@@ -395,7 +406,11 @@ impl<'tcx> RegionConstraintCollector<'_, 'tcx> {
         // `RegionConstraintData` contains the relationship here.
         if *any_unifications {
             *any_unifications = false;
-            self.unification_table_mut().reset_unifications(|_| UnifiedRegion::new(None));
+            // Manually inlined `self.unification_table_mut()` as `self` is used in the closure.
+            ut::UnificationTable::with_log(&mut self.storage.unification_table, &mut self.undo_log)
+                .reset_unifications(|key| RegionVariableValue::Unknown {
+                    universe: self.storage.var_infos[key.vid].universe,
+                });
         }
 
         data
@@ -422,18 +437,13 @@ impl<'tcx> RegionConstraintCollector<'_, 'tcx> {
     ) -> RegionVid {
         let vid = self.var_infos.push(RegionVariableInfo { origin, universe });
 
-        let u_vid = self.unification_table_mut().new_key(UnifiedRegion::new(None));
+        let u_vid = self.unification_table_mut().new_key(RegionVariableValue::Unknown { universe });
         assert_eq!(vid, u_vid.vid);
         self.undo_log.push(AddVar(vid));
         debug!("created new region variable {:?} in {:?} with origin {:?}", vid, universe, origin);
         vid
     }
 
-    /// Returns the universe for the given variable.
-    pub(super) fn var_universe(&self, vid: RegionVid) -> ty::UniverseIndex {
-        self.var_infos[vid].universe
-    }
-
     /// Returns the origin for the given variable.
     pub(super) fn var_origin(&self, vid: RegionVid) -> RegionVariableOrigin {
         self.var_infos[vid].origin
@@ -467,26 +477,41 @@ impl<'tcx> RegionConstraintCollector<'_, 'tcx> {
     pub(super) fn make_eqregion(
         &mut self,
         origin: SubregionOrigin<'tcx>,
-        sub: Region<'tcx>,
-        sup: Region<'tcx>,
+        a: Region<'tcx>,
+        b: Region<'tcx>,
     ) {
-        if sub != sup {
+        if a != b {
             // Eventually, it would be nice to add direct support for
             // equating regions.
-            self.make_subregion(origin.clone(), sub, sup);
-            self.make_subregion(origin, sup, sub);
-
-            match (sub, sup) {
-                (Region(Interned(ReVar(sub), _)), Region(Interned(ReVar(sup), _))) => {
-                    debug!("make_eqregion: unifying {:?} with {:?}", sub, sup);
-                    self.unification_table_mut().union(*sub, *sup);
-                    self.any_unifications = true;
+            self.make_subregion(origin.clone(), a, b);
+            self.make_subregion(origin, b, a);
+
+            match (a.kind(), b.kind()) {
+                (ty::ReVar(a), ty::ReVar(b)) => {
+                    debug!("make_eqregion: unifying {:?} with {:?}", a, b);
+                    if self.unification_table_mut().unify_var_var(a, b).is_ok() {
+                        self.any_unifications = true;
+                    }
+                }
+                (ty::ReVar(vid), _) => {
+                    debug!("make_eqregion: unifying {:?} with {:?}", vid, b);
+                    if self
+                        .unification_table_mut()
+                        .unify_var_value(vid, RegionVariableValue::Known { value: b })
+                        .is_ok()
+                    {
+                        self.any_unifications = true;
+                    };
                 }
-                (Region(Interned(ReVar(vid), _)), value)
-                | (value, Region(Interned(ReVar(vid), _))) => {
-                    debug!("make_eqregion: unifying {:?} with {:?}", vid, value);
-                    self.unification_table_mut().union_value(*vid, UnifiedRegion::new(Some(value)));
-                    self.any_unifications = true;
+                (_, ty::ReVar(vid)) => {
+                    debug!("make_eqregion: unifying {:?} with {:?}", a, vid);
+                    if self
+                        .unification_table_mut()
+                        .unify_var_value(vid, RegionVariableValue::Known { value: a })
+                        .is_ok()
+                    {
+                        self.any_unifications = true;
+                    };
                 }
                 (_, _) => {}
             }
@@ -603,18 +628,21 @@ impl<'tcx> RegionConstraintCollector<'_, 'tcx> {
         tcx: TyCtxt<'tcx>,
         vid: ty::RegionVid,
     ) -> ty::Region<'tcx> {
-        let mut ut = self.unification_table_mut(); // FIXME(rust-lang/ena#42): unnecessary mut
+        let mut ut = self.unification_table_mut();
         let root_vid = ut.find(vid).vid;
-        let resolved = ut
-            .probe_value(root_vid)
-            .get_value_ignoring_universes()
-            .unwrap_or_else(|| ty::Region::new_var(tcx, root_vid));
-
-        // Don't resolve a variable to a region that it cannot name.
-        if self.var_universe(vid).can_name(self.universe(resolved)) {
-            resolved
-        } else {
-            ty::Region::new_var(tcx, vid)
+        match ut.probe_value(root_vid) {
+            RegionVariableValue::Known { value } => value,
+            RegionVariableValue::Unknown { .. } => ty::Region::new_var(tcx, root_vid),
+        }
+    }
+
+    pub fn probe_value(
+        &mut self,
+        vid: ty::RegionVid,
+    ) -> Result<ty::Region<'tcx>, ty::UniverseIndex> {
+        match self.unification_table_mut().probe_value(vid) {
+            RegionVariableValue::Known { value } => Ok(value),
+            RegionVariableValue::Unknown { universe } => Err(universe),
         }
     }
 
@@ -654,7 +682,7 @@ impl<'tcx> RegionConstraintCollector<'_, 'tcx> {
         new_r
     }
 
-    pub fn universe(&self, region: Region<'tcx>) -> ty::UniverseIndex {
+    pub fn universe(&mut self, region: Region<'tcx>) -> ty::UniverseIndex {
         match *region {
             ty::ReStatic
             | ty::ReErased
@@ -662,7 +690,10 @@ impl<'tcx> RegionConstraintCollector<'_, 'tcx> {
             | ty::ReEarlyParam(..)
             | ty::ReError(_) => ty::UniverseIndex::ROOT,
             ty::RePlaceholder(placeholder) => placeholder.universe,
-            ty::ReVar(vid) => self.var_universe(vid),
+            ty::ReVar(vid) => match self.probe_value(vid) {
+                Ok(value) => self.universe(value),
+                Err(universe) => universe,
+            },
             ty::ReBound(..) => bug!("universe(): encountered bound region {:?}", region),
         }
     }
diff --git a/compiler/rustc_middle/src/infer/unify_key.rs b/compiler/rustc_middle/src/infer/unify_key.rs
index 63c0ebd5f6b..8d8ff48a79e 100644
--- a/compiler/rustc_middle/src/infer/unify_key.rs
+++ b/compiler/rustc_middle/src/infer/unify_key.rs
@@ -1,4 +1,4 @@
-use crate::ty::{self, Region, Ty, TyCtxt};
+use crate::ty::{self, Ty, TyCtxt};
 use rustc_data_structures::unify::{NoError, UnifyKey, UnifyValue};
 use rustc_span::def_id::DefId;
 use rustc_span::symbol::Symbol;
@@ -10,26 +10,16 @@ pub trait ToType {
     fn to_type<'tcx>(&self, tcx: TyCtxt<'tcx>) -> Ty<'tcx>;
 }
 
-#[derive(PartialEq, Copy, Clone, Debug)]
-pub struct UnifiedRegion<'tcx> {
-    value: Option<ty::Region<'tcx>>,
-}
-
-impl<'tcx> UnifiedRegion<'tcx> {
-    pub fn new(value: Option<Region<'tcx>>) -> Self {
-        Self { value }
-    }
-
-    /// The caller is responsible for checking universe compatibility before using this value.
-    pub fn get_value_ignoring_universes(self) -> Option<Region<'tcx>> {
-        self.value
-    }
+#[derive(Copy, Clone, Debug)]
+pub enum RegionVariableValue<'tcx> {
+    Known { value: ty::Region<'tcx> },
+    Unknown { universe: ty::UniverseIndex },
 }
 
 #[derive(PartialEq, Copy, Clone, Debug)]
 pub struct RegionVidKey<'tcx> {
     pub vid: ty::RegionVid,
-    pub phantom: PhantomData<UnifiedRegion<'tcx>>,
+    pub phantom: PhantomData<RegionVariableValue<'tcx>>,
 }
 
 impl<'tcx> From<ty::RegionVid> for RegionVidKey<'tcx> {
@@ -39,7 +29,7 @@ impl<'tcx> From<ty::RegionVid> for RegionVidKey<'tcx> {
 }
 
 impl<'tcx> UnifyKey for RegionVidKey<'tcx> {
-    type Value = UnifiedRegion<'tcx>;
+    type Value = RegionVariableValue<'tcx>;
     #[inline]
     fn index(&self) -> u32 {
         self.vid.as_u32()
@@ -53,36 +43,40 @@ impl<'tcx> UnifyKey for RegionVidKey<'tcx> {
     }
 }
 
-impl<'tcx> UnifyValue for UnifiedRegion<'tcx> {
-    type Error = NoError;
-
-    fn unify_values(value1: &Self, value2: &Self) -> Result<Self, NoError> {
-        // We pick the value of the least universe because it is compatible with more variables.
-        // This is *not* necessary for completeness.
-        #[cold]
-        fn min_universe<'tcx>(r1: Region<'tcx>, r2: Region<'tcx>) -> Region<'tcx> {
-            cmp::min_by_key(r1, r2, |r| match r.kind() {
-                ty::ReStatic
-                | ty::ReErased
-                | ty::ReLateParam(..)
-                | ty::ReEarlyParam(..)
-                | ty::ReError(_) => ty::UniverseIndex::ROOT,
-                ty::RePlaceholder(placeholder) => placeholder.universe,
-                ty::ReVar(..) | ty::ReBound(..) => bug!("not a universal region"),
-            })
-        }
+pub struct RegionUnificationError;
+impl<'tcx> UnifyValue for RegionVariableValue<'tcx> {
+    type Error = RegionUnificationError;
 
-        Ok(match (value1.value, value2.value) {
-            // Here we can just pick one value, because the full constraints graph
-            // will be handled later. Ideally, we might want a `MultipleValues`
-            // variant or something. For now though, this is fine.
-            (Some(val1), Some(val2)) => Self { value: Some(min_universe(val1, val2)) },
+    fn unify_values(value1: &Self, value2: &Self) -> Result<Self, Self::Error> {
+        match (*value1, *value2) {
+            (RegionVariableValue::Known { .. }, RegionVariableValue::Known { .. }) => {
+                Err(RegionUnificationError)
+            }
 
-            (Some(_), _) => *value1,
-            (_, Some(_)) => *value2,
+            (RegionVariableValue::Known { value }, RegionVariableValue::Unknown { universe })
+            | (RegionVariableValue::Unknown { universe }, RegionVariableValue::Known { value }) => {
+                let universe_of_value = match value.kind() {
+                    ty::ReStatic
+                    | ty::ReErased
+                    | ty::ReLateParam(..)
+                    | ty::ReEarlyParam(..)
+                    | ty::ReError(_) => ty::UniverseIndex::ROOT,
+                    ty::RePlaceholder(placeholder) => placeholder.universe,
+                    ty::ReVar(..) | ty::ReBound(..) => bug!("not a universal region"),
+                };
+
+                if universe.can_name(universe_of_value) {
+                    Ok(RegionVariableValue::Known { value })
+                } else {
+                    Err(RegionUnificationError)
+                }
+            }
 
-            (None, None) => *value1,
-        })
+            (
+                RegionVariableValue::Unknown { universe: a },
+                RegionVariableValue::Unknown { universe: b },
+            ) => Ok(RegionVariableValue::Unknown { universe: a.min(b) }),
+        }
     }
 }