eagerly compute `sub_relations` again

5 files changed, 131 insertions, 1 deletions

diff --git a/compiler/rustc_infer/src/infer/context.rs b/compiler/rustc_infer/src/infer/context.rs
index 8265fccabc9..efd623fac78 100644
--- a/compiler/rustc_infer/src/infer/context.rs
+++ b/compiler/rustc_infer/src/infer/context.rs
@@ -179,6 +179,10 @@ impl<'tcx> rustc_type_ir::InferCtxtLike for InferCtxt<'tcx> {
         self.inner.borrow_mut().type_variables().equate(a, b);
     }
 
+    fn sub_ty_vids_raw(&self, a: ty::TyVid, b: ty::TyVid) {
+        self.sub_ty_vids_raw(a, b);
+    }
+
     fn equate_int_vids_raw(&self, a: ty::IntVid, b: ty::IntVid) {
         self.inner.borrow_mut().int_unification_table().union(a, b);
     }
diff --git a/compiler/rustc_infer/src/infer/mod.rs b/compiler/rustc_infer/src/infer/mod.rs
index d105d24bed7..48dd699cd50 100644
--- a/compiler/rustc_infer/src/infer/mod.rs
+++ b/compiler/rustc_infer/src/infer/mod.rs
@@ -764,6 +764,7 @@ impl<'tcx> InferCtxt<'tcx> {
         let r_b = self.shallow_resolve(predicate.skip_binder().b);
         match (r_a.kind(), r_b.kind()) {
             (&ty::Infer(ty::TyVar(a_vid)), &ty::Infer(ty::TyVar(b_vid))) => {
+                self.sub_ty_vids_raw(a_vid, b_vid);
                 return Err((a_vid, b_vid));
             }
             _ => {}
@@ -1128,6 +1129,14 @@ impl<'tcx> InferCtxt<'tcx> {
         self.inner.borrow_mut().type_variables().root_var(var)
     }
 
+    pub fn sub_ty_vids_raw(&self, a: ty::TyVid, b: ty::TyVid) {
+        self.inner.borrow_mut().type_variables().sub(a, b);
+    }
+
+    pub fn sub_root_var(&self, var: ty::TyVid) -> ty::TyVid {
+        self.inner.borrow_mut().type_variables().sub_root_var(var)
+    }
+
     pub fn root_const_var(&self, var: ty::ConstVid) -> ty::ConstVid {
         self.inner.borrow_mut().const_unification_table().find(var).vid
     }
diff --git a/compiler/rustc_infer/src/infer/relate/generalize.rs b/compiler/rustc_infer/src/infer/relate/generalize.rs
index a000bb1123c..2323add0b32 100644
--- a/compiler/rustc_infer/src/infer/relate/generalize.rs
+++ b/compiler/rustc_infer/src/infer/relate/generalize.rs
@@ -519,6 +519,10 @@ impl<'tcx> TypeRelation<TyCtxt<'tcx>> for Generalizer<'_, 'tcx> {
                             let origin = inner.type_variables().var_origin(vid);
                             let new_var_id =
                                 inner.type_variables().new_var(self.for_universe, origin);
+                            // Record that `vid` and `new_var_id` have to be subtypes
+                            // of each other. This is currently only used for diagnostics.
+                            // To see why, see the docs in the `type_variables` module.
+                            inner.type_variables().sub(vid, new_var_id);
                             // If we're in the new solver and create a new inference
                             // variable inside of an alias we eagerly constrain that
                             // inference variable to prevent unexpected ambiguity errors.
diff --git a/compiler/rustc_infer/src/infer/snapshot/undo_log.rs b/compiler/rustc_infer/src/infer/snapshot/undo_log.rs
index fcc0ab3af41..22c815fb87c 100644
--- a/compiler/rustc_infer/src/infer/snapshot/undo_log.rs
+++ b/compiler/rustc_infer/src/infer/snapshot/undo_log.rs
@@ -20,7 +20,7 @@ pub struct Snapshot<'tcx> {
 pub(crate) enum UndoLog<'tcx> {
     DuplicateOpaqueType,
     OpaqueTypes(OpaqueTypeKey<'tcx>, Option<OpaqueHiddenType<'tcx>>),
-    TypeVariables(sv::UndoLog<ut::Delegate<type_variable::TyVidEqKey<'tcx>>>),
+    TypeVariables(type_variable::UndoLog<'tcx>),
     ConstUnificationTable(sv::UndoLog<ut::Delegate<ConstVidKey<'tcx>>>),
     IntUnificationTable(sv::UndoLog<ut::Delegate<ty::IntVid>>),
     FloatUnificationTable(sv::UndoLog<ut::Delegate<ty::FloatVid>>),
@@ -49,6 +49,8 @@ impl_from! {
     RegionConstraintCollector(region_constraints::UndoLog<'tcx>),
 
     TypeVariables(sv::UndoLog<ut::Delegate<type_variable::TyVidEqKey<'tcx>>>),
+    TypeVariables(sv::UndoLog<ut::Delegate<type_variable::TyVidSubKey>>),
+    TypeVariables(type_variable::UndoLog<'tcx>),
     IntUnificationTable(sv::UndoLog<ut::Delegate<ty::IntVid>>),
     FloatUnificationTable(sv::UndoLog<ut::Delegate<ty::FloatVid>>),
 
diff --git a/compiler/rustc_infer/src/infer/type_variable.rs b/compiler/rustc_infer/src/infer/type_variable.rs
index 6f6791804d3..613e5a5883a 100644
--- a/compiler/rustc_infer/src/infer/type_variable.rs
+++ b/compiler/rustc_infer/src/infer/type_variable.rs
@@ -13,12 +13,48 @@ use tracing::debug;
 
 use crate::infer::InferCtxtUndoLogs;
 
+/// Represents a single undo-able action that affects a type inference variable.
+#[derive(Clone)]
+pub(crate) enum UndoLog<'tcx> {
+    EqRelation(sv::UndoLog<ut::Delegate<TyVidEqKey<'tcx>>>),
+    SubRelation(sv::UndoLog<ut::Delegate<TyVidSubKey>>),
+}
+
+/// Convert from a specific kind of undo to the more general UndoLog
+impl<'tcx> From<sv::UndoLog<ut::Delegate<TyVidEqKey<'tcx>>>> for UndoLog<'tcx> {
+    fn from(l: sv::UndoLog<ut::Delegate<TyVidEqKey<'tcx>>>) -> Self {
+        UndoLog::EqRelation(l)
+    }
+}
+
+/// Convert from a specific kind of undo to the more general UndoLog
+impl<'tcx> From<sv::UndoLog<ut::Delegate<TyVidSubKey>>> for UndoLog<'tcx> {
+    fn from(l: sv::UndoLog<ut::Delegate<TyVidSubKey>>) -> Self {
+        UndoLog::SubRelation(l)
+    }
+}
+
 impl<'tcx> Rollback<sv::UndoLog<ut::Delegate<TyVidEqKey<'tcx>>>> for TypeVariableStorage<'tcx> {
     fn reverse(&mut self, undo: sv::UndoLog<ut::Delegate<TyVidEqKey<'tcx>>>) {
         self.eq_relations.reverse(undo)
     }
 }
 
+impl<'tcx> Rollback<sv::UndoLog<ut::Delegate<TyVidSubKey>>> for TypeVariableStorage<'tcx> {
+    fn reverse(&mut self, undo: sv::UndoLog<ut::Delegate<TyVidSubKey>>) {
+        self.sub_relations.reverse(undo)
+    }
+}
+
+impl<'tcx> Rollback<UndoLog<'tcx>> for TypeVariableStorage<'tcx> {
+    fn reverse(&mut self, undo: UndoLog<'tcx>) {
+        match undo {
+            UndoLog::EqRelation(undo) => self.eq_relations.reverse(undo),
+            UndoLog::SubRelation(undo) => self.sub_relations.reverse(undo),
+        }
+    }
+}
+
 #[derive(Clone, Default)]
 pub(crate) struct TypeVariableStorage<'tcx> {
     /// The origins of each type variable.
@@ -27,6 +63,23 @@ pub(crate) struct TypeVariableStorage<'tcx> {
     /// constraint `?X == ?Y`. This table also stores, for each key,
     /// the known value.
     eq_relations: ut::UnificationTableStorage<TyVidEqKey<'tcx>>,
+    /// Only used by `-Znext-solver` and for diagnostics.
+    ///
+    /// When reporting ambiguity errors, we sometimes want to
+    /// treat all inference vars which are subtypes of each
+    /// others as if they are equal. For this case we compute
+    /// the transitive closure of our subtype obligations here.
+    ///
+    /// E.g. when encountering ambiguity errors, we want to suggest
+    /// specifying some method argument or to add a type annotation
+    /// to a local variable. Because subtyping cannot change the
+    /// shape of a type, it's fine if the cause of the ambiguity error
+    /// is only related to the suggested variable via subtyping.
+    ///
+    /// Even for something like `let x = returns_arg(); x.method();` the
+    /// type of `x` is only a supertype of the argument of `returns_arg`. We
+    /// still want to suggest specifying the type of the argument.
+    sub_relations: ut::UnificationTableStorage<TyVidSubKey>,
 }
 
 pub(crate) struct TypeVariableTable<'a, 'tcx> {
@@ -109,6 +162,16 @@ impl<'tcx> TypeVariableTable<'_, 'tcx> {
         debug_assert!(self.probe(a).is_unknown());
         debug_assert!(self.probe(b).is_unknown());
         self.eq_relations().union(a, b);
+        self.sub_relations().union(a, b);
+    }
+
+    /// Records that `a <: b`, depending on `dir`.
+    ///
+    /// Precondition: neither `a` nor `b` are known.
+    pub(crate) fn sub(&mut self, a: ty::TyVid, b: ty::TyVid) {
+        debug_assert!(self.probe(a).is_unknown());
+        debug_assert!(self.probe(b).is_unknown());
+        self.sub_relations().union(a, b);
     }
 
     /// Instantiates `vid` with the type `ty`.
@@ -142,6 +205,10 @@ impl<'tcx> TypeVariableTable<'_, 'tcx> {
         origin: TypeVariableOrigin,
     ) -> ty::TyVid {
         let eq_key = self.eq_relations().new_key(TypeVariableValue::Unknown { universe });
+
+        let sub_key = self.sub_relations().new_key(());
+        debug_assert_eq!(eq_key.vid, sub_key.vid);
+
         let index = self.storage.values.push(TypeVariableData { origin });
         debug_assert_eq!(eq_key.vid, index);
 
@@ -164,6 +231,18 @@ impl<'tcx> TypeVariableTable<'_, 'tcx> {
         self.eq_relations().find(vid).vid
     }
 
+    /// Returns the "root" variable of `vid` in the `sub_relations`
+    /// equivalence table. All type variables that have been are
+    /// related via equality or subtyping will yield the same root
+    /// variable (per the union-find algorithm), so `sub_root_var(a)
+    /// == sub_root_var(b)` implies that:
+    /// ```text
+    /// exists X. (a <: X || X <: a) && (b <: X || X <: b)
+    /// ```
+    pub(crate) fn sub_root_var(&mut self, vid: ty::TyVid) -> ty::TyVid {
+        self.sub_relations().find(vid).vid
+    }
+
     /// Retrieves the type to which `vid` has been instantiated, if
     /// any.
     pub(crate) fn probe(&mut self, vid: ty::TyVid) -> TypeVariableValue<'tcx> {
@@ -181,6 +260,11 @@ impl<'tcx> TypeVariableTable<'_, 'tcx> {
         self.storage.eq_relations.with_log(self.undo_log)
     }
 
+    #[inline]
+    fn sub_relations(&mut self) -> super::UnificationTable<'_, 'tcx, TyVidSubKey> {
+        self.storage.sub_relations.with_log(self.undo_log)
+    }
+
     /// Returns a range of the type variables created during the snapshot.
     pub(crate) fn vars_since_snapshot(
         &mut self,
@@ -243,6 +327,33 @@ impl<'tcx> ut::UnifyKey for TyVidEqKey<'tcx> {
     }
 }
 
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub(crate) struct TyVidSubKey {
+    vid: ty::TyVid,
+}
+
+impl From<ty::TyVid> for TyVidSubKey {
+    #[inline] // make this function eligible for inlining - it is quite hot.
+    fn from(vid: ty::TyVid) -> Self {
+        TyVidSubKey { vid }
+    }
+}
+
+impl ut::UnifyKey for TyVidSubKey {
+    type Value = ();
+    #[inline]
+    fn index(&self) -> u32 {
+        self.vid.as_u32()
+    }
+    #[inline]
+    fn from_index(i: u32) -> TyVidSubKey {
+        TyVidSubKey { vid: ty::TyVid::from_u32(i) }
+    }
+    fn tag() -> &'static str {
+        "TyVidSubKey"
+    }
+}
+
 impl<'tcx> ut::UnifyValue for TypeVariableValue<'tcx> {
     type Error = ut::NoError;