Perf: fulfillment wake by stalled var

crates/hir-ty/src/next_solver/infer/mod.rs

@@ -59,6 +59,7 @@ pub mod relate;
 pub mod resolve;
 pub mod select;
 pub(crate) mod snapshot;
+pub(crate) use snapshot::StalledVarKey;
 pub mod traits;
 mod type_variable;
 mod unify_key;
@@ -878,7 +879,7 @@ impl<'db> InferCtxt<'db> {
     pub fn take_opaque_types(
         &self,
     ) -> impl IntoIterator<Item = (OpaqueTypeKey<'db>, OpaqueHiddenType<'db>)> + use<'db> {
-        self.inner.borrow_mut().opaque_type_storage.take_opaque_types()
+        self.inner.borrow_mut().opaque_types().take_opaque_types()
     }
 
     #[instrument(level = "debug", skip(self), ret)]
@@ -1004,6 +1005,14 @@ impl<'db> InferCtxt<'db> {
         self.inner.borrow_mut().const_unification_table().find(var).vid
     }
 
+    pub(crate) fn root_int_var(&self, var: IntVid) -> IntVid {
+        self.inner.borrow_mut().int_unification_table().find(var)
+    }
+
+    pub(crate) fn root_float_var(&self, var: FloatVid) -> FloatVid {
+        self.inner.borrow_mut().float_unification_table().find(var)
+    }
+
     /// Resolves an int var to a rigid int type, if it was constrained to one,
     /// or else the root int var in the unification table.
     pub fn opportunistic_resolve_int_var(&self, vid: IntVid) -> Ty<'db> {

crates/hir-ty/src/next_solver/infer/opaque_types/table.rs

@@ -27,6 +27,15 @@ pub struct OpaqueTypeStorageEntries {
     duplicate_entries: usize,
 }
 
+impl OpaqueTypeStorageEntries {
+    /// The raw entry count, for constructing a [`rustc_next_trait_solver::solve::GoalStalledOn`]
+    /// outside of the solver's own canonicalization (which is where `num_opaques` is normally
+    /// computed from).
+    pub(crate) fn opaque_type_count(self) -> usize {
+        self.opaque_types
+    }
+}
+
 impl rustc_type_ir::inherent::OpaqueTypeStorageEntries for OpaqueTypeStorageEntries {
     fn needs_reevaluation(self, canonicalized: usize) -> bool {
         self.opaque_types != canonicalized
@@ -159,4 +168,15 @@ impl<'a, 'db> OpaqueTypeTable<'a, 'db> {
         self.storage.duplicate_entries.push((key, hidden_type));
         self.undo_log.push(UndoLog::DuplicateOpaqueType);
     }
+
+    pub(crate) fn take_opaque_types(
+        &mut self,
+    ) -> impl IntoIterator<Item = (OpaqueTypeKey<'db>, OpaqueHiddenType<'db>)> + use<'db> {
+        if !self.storage.is_empty() {
+            // Draining the storage shrinks the entry count without going through the undo
+            // machinery; record it so obligations stalled on the opaque count get rechecked.
+            self.undo_log.mark_opaques_changed();
+        }
+        self.storage.take_opaque_types()
+    }
 }

crates/hir-ty/src/next_solver/infer/region_constraints/mod.rs

@@ -323,7 +323,12 @@ impl<'db> RegionConstraintCollector<'db, '_> {
     ///
     /// Not legal during a snapshot.
     pub fn take_and_reset_data(&mut self) -> RegionConstraintData<'db> {
-        assert!(!UndoLogs::<UndoLog<'db>>::in_snapshot(&self.undo_log));
+        // The *inherent* method, not the always-`true` trait impl; see `InferCtxtUndoLogs`.
+        // The explicit deref is needed because `ena` blanket-implements `UndoLogs` for
+        // `&mut U`, so on the reference itself only the trait method exists. Removing the
+        // deref is a compile error (E0283, ambiguous `T` in `UndoLogs<T>`), not a silent
+        // switch to the trait impl.
+        assert!(!(*self.undo_log).in_snapshot());
 
         // If you add a new field to `RegionConstraintCollector`, you
         // should think carefully about whether it needs to be cleared

crates/hir-ty/src/next_solver/infer/snapshot/mod.rs

@@ -10,6 +10,7 @@ use super::region_constraints::RegionSnapshot;
 mod fudge;
 pub(crate) mod undo_log;
 
+pub(crate) use undo_log::StalledVarKey;
 use undo_log::{Snapshot, UndoLog};
 
 #[must_use = "once you start a snapshot, you should always consume it"]
@@ -40,13 +41,31 @@ impl<'db> InferCtxt<'db> {
     }
 
     pub fn in_snapshot(&self) -> bool {
-        UndoLogs::<UndoLog<'db>>::in_snapshot(&self.inner.borrow_mut().undo_log)
+        // The *inherent* method, not the always-`true` trait impl; see `InferCtxtUndoLogs`.
+        self.inner.borrow().undo_log.in_snapshot()
     }
 
     pub fn num_open_snapshots(&self) -> usize {
         UndoLogs::<UndoLog<'db>>::num_open_snapshots(&self.inner.borrow_mut().undo_log)
     }
 
+    /// The current length of the changed-inference-variables log, i.e. a cursor position from
+    /// which [`Self::changed_vars_since`] later reports what changed. See
+    /// [`undo_log::InferCtxtUndoLogs`].
+    pub(crate) fn changed_vars_len(&self) -> usize {
+        self.inner.borrow().undo_log.changed_vars_len()
+    }
+
+    /// Appends to `out` the variables that (potentially) changed since `cursor` (a position
+    /// previously obtained from [`Self::changed_vars_len`]), then advances `cursor` past them.
+    /// Reading is non-destructive, so any number of consumers can watch the same infer context
+    /// with their own cursors.
+    pub(crate) fn changed_vars_since(&self, cursor: &mut usize, out: &mut Vec<StalledVarKey>) {
+        let inner = self.inner.borrow();
+        out.extend_from_slice(inner.undo_log.changed_vars_since(*cursor));
+        *cursor = inner.undo_log.changed_vars_len();
+    }
+
     pub fn start_snapshot(&self) -> CombinedSnapshot {
         debug!("start_snapshot()");
 
@@ -123,3 +142,102 @@ impl<'db> InferCtxt<'db> {
         self.inner.borrow().undo_log.opaque_types_in_snapshot(&snapshot.undo_snapshot)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use rustc_type_ir::{TypingMode, inherent::Ty as _};
+    use test_fixture::WithFixture;
+
+    use crate::{
+        next_solver::{
+            DbInterner, Ty,
+            infer::{DbInternerInferExt, InferCtxt, StalledVarKey},
+        },
+        test_db::TestDB,
+    };
+
+    fn with_infcx(f: impl FnOnce(&InferCtxt<'_>)) {
+        let (db, file_id) = TestDB::with_single_file("fn f() {}");
+        crate::attach_db(&db, || {
+            let krate = db.module_for_file(file_id.file_id(&db)).krate(&db);
+            let interner = DbInterner::new_with(&db, krate);
+            let infcx = interner.infer_ctxt().build(TypingMode::non_body_analysis());
+            f(&infcx);
+        });
+    }
+
+    fn new_ty_var(infcx: &InferCtxt<'_>) -> rustc_type_ir::TyVid {
+        infcx.next_ty_vid(crate::Span::Dummy)
+    }
+
+    fn changed_since(infcx: &InferCtxt<'_>, cursor: &mut usize) -> Vec<StalledVarKey> {
+        let mut out = Vec::new();
+        infcx.changed_vars_since(cursor, &mut out);
+        out
+    }
+
+    /// A mutation performed while *no* snapshot is open must still be reported: `ena` gates
+    /// its undo-log `push` calls on `in_snapshot()`, which `InferCtxtUndoLogs` deliberately
+    /// short-circuits to `true` to keep this log complete. A lost event here would leave a
+    /// `FulfillmentCtxt` obligation parked on the variable with a stale `Certainty::Maybe`
+    /// forever, i.e. wrong inference results rather than a crash.
+    #[test]
+    fn zero_snapshot_mutation_is_reported() {
+        with_infcx(|infcx| {
+            assert_eq!(infcx.num_open_snapshots(), 0);
+            let vid = new_ty_var(infcx);
+            let mut cursor = infcx.changed_vars_len();
+            infcx
+                .inner
+                .borrow_mut()
+                .type_variables()
+                .instantiate(vid, Ty::new_bool(infcx.interner));
+            assert!(changed_since(infcx, &mut cursor).contains(&StalledVarKey::Ty(vid)));
+        });
+    }
+
+    /// Unioning two variables must be observable through *both* variables' identities, and
+    /// instantiating the surviving root afterwards must be observable through that root: a
+    /// watcher keyed on either original variable is woken by the union, re-registers under
+    /// the merged root, and is then woken again by the instantiation.
+    #[test]
+    fn union_then_instantiate_root_is_reported() {
+        with_infcx(|infcx| {
+            let v1 = new_ty_var(infcx);
+            let v2 = new_ty_var(infcx);
+            let mut cursor = infcx.changed_vars_len();
+            infcx.inner.borrow_mut().type_variables().equate(v1, v2);
+            let after_union = changed_since(infcx, &mut cursor);
+            assert!(after_union.contains(&StalledVarKey::Ty(v1)));
+            assert!(after_union.contains(&StalledVarKey::Ty(v2)));
+
+            let root = infcx.root_var(v1);
+            assert_eq!(root, infcx.root_var(v2));
+            infcx
+                .inner
+                .borrow_mut()
+                .type_variables()
+                .instantiate(root, Ty::new_bool(infcx.interner));
+            assert!(changed_since(infcx, &mut cursor).contains(&StalledVarKey::Ty(root)));
+        });
+    }
+
+    /// Changes made inside a probe are rolled back, but the wake events they produced must
+    /// stick around: waking an obligation spuriously is harmless, missing a wake is not, and
+    /// a consumer may only get to read the log after the rollback already happened.
+    #[test]
+    fn rolled_back_mutation_still_reported() {
+        with_infcx(|infcx| {
+            let vid = new_ty_var(infcx);
+            let mut cursor = infcx.changed_vars_len();
+            infcx.probe(|_| {
+                infcx
+                    .inner
+                    .borrow_mut()
+                    .type_variables()
+                    .instantiate(vid, Ty::new_bool(infcx.interner));
+            });
+            assert!(changed_since(infcx, &mut cursor).contains(&StalledVarKey::Ty(vid)));
+        });
+    }
+}