perf: Only process changed obligations in ObligationForest

compiler/rustc_data_structures/src/lib.rs

@@ -73,6 +73,7 @@ pub mod flock;
 pub mod fx;
 pub mod graph;
 pub mod jobserver;
+pub mod logged_unification_table;
 pub mod macros;
 pub mod map_in_place;
 pub mod obligation_forest;
@@ -83,13 +84,15 @@ pub mod small_c_str;
 pub mod snapshot_map;
 pub mod stable_map;
 pub mod svh;
+pub mod unify_log;
 pub use ena::snapshot_vec;
 pub mod sorted_map;
 pub mod stable_set;
 #[macro_use]
 pub mod stable_hasher;
 mod atomic_ref;
 pub mod fingerprint;
+pub mod modified_set;
 pub mod profiling;
 pub mod sharded;
 pub mod stack;

compiler/rustc_data_structures/src/logged_unification_table.rs

@@ -0,0 +1,209 @@
+use rustc_index::vec::Idx;
+
+use crate::modified_set as ms;
+use crate::snapshot_vec as sv;
+use crate::unify as ut;
+use crate::unify_log as ul;
+
+use ena::undo_log::{Rollback, UndoLogs};
+
+pub enum UndoLog<K: ut::UnifyKey, I = K> {
+    Relation(sv::UndoLog<ut::Delegate<K>>),
+    UnifyLog(ul::Undo<I>),
+    ModifiedSet(ms::Undo<I>),
+}
+
+impl<K: ut::UnifyKey, I> From<sv::UndoLog<ut::Delegate<K>>> for UndoLog<K, I> {
+    fn from(l: sv::UndoLog<ut::Delegate<K>>) -> Self {
+        UndoLog::Relation(l)
+    }
+}
+
+impl<K: ut::UnifyKey, I> From<ul::Undo<I>> for UndoLog<K, I> {
+    fn from(l: ul::Undo<I>) -> Self {
+        UndoLog::UnifyLog(l)
+    }
+}
+
+impl<K: ut::UnifyKey, I> From<ms::Undo<I>> for UndoLog<K, I> {
+    fn from(l: ms::Undo<I>) -> Self {
+        UndoLog::ModifiedSet(l)
+    }
+}
+
+impl<K: ut::UnifyKey, I: Idx> Rollback<UndoLog<K, I>> for LoggedUnificationTableStorage<K, I> {
+    fn reverse(&mut self, undo: UndoLog<K, I>) {
+        match undo {
+            UndoLog::Relation(undo) => self.relations.reverse(undo),
+            UndoLog::UnifyLog(undo) => self.unify_log.reverse(undo),
+            UndoLog::ModifiedSet(undo) => self.modified_set.reverse(undo),
+        }
+    }
+}
+
+/// Storage for `LoggedUnificationTable`
+pub struct LoggedUnificationTableStorage<K: ut::UnifyKey, I: Idx = K> {
+    relations: ut::UnificationTableStorage<K>,
+    unify_log: ul::UnifyLog<I>,
+    modified_set: ms::ModifiedSet<I>,
+}
+
+/// UnificationTableStorage which logs which variables has been unified with a value, allowing watchers
+/// to only iterate over the changed variables instead of all variables
+pub struct LoggedUnificationTable<'a, K: ut::UnifyKey, I: Idx, L> {
+    storage: &'a mut LoggedUnificationTableStorage<K, I>,
+    undo_log: L,
+}
+
+impl<K, I> LoggedUnificationTableStorage<K, I>
+where
+    K: ut::UnifyKey + From<I>,
+    I: Idx + From<K>,
+{
+    pub fn new() -> Self {
+        Self {
+            relations: Default::default(),
+            unify_log: ul::UnifyLog::new(),
+            modified_set: ms::ModifiedSet::new(),
+        }
+    }
+
+    pub fn with_log<L>(&mut self, undo_log: L) -> LoggedUnificationTable<'_, K, I, L> {
+        LoggedUnificationTable { storage: self, undo_log }
+    }
+}
+
+impl<K, I, L> LoggedUnificationTable<'_, K, I, L>
+where
+    K: ut::UnifyKey,
+    I: Idx,
+{
+    pub fn len(&self) -> usize {
+        self.storage.relations.len()
+    }
+}
+
+impl<K, I, L> LoggedUnificationTable<'_, K, I, L>
+where
+    K: ut::UnifyKey + From<I>,
+    I: Idx + From<K>,
+    L: UndoLogs<ms::Undo<I>> + UndoLogs<ul::Undo<I>> + UndoLogs<sv::UndoLog<ut::Delegate<K>>>,
+{
+    fn relations(
+        &mut self,
+    ) -> ut::UnificationTable<ut::InPlace<K, &mut ut::UnificationStorage<K>, &mut L>> {
+        ut::UnificationTable::with_log(&mut self.storage.relations, &mut self.undo_log)
+    }
+
+    pub fn unify(&mut self, a: I, b: I)
+    where
+        K::Value: ut::UnifyValue<Error = ut::NoError>,
+    {
+        self.unify_var_var(a, b).unwrap();
+    }
+
+    pub fn instantiate(&mut self, vid: I, ty: K::Value) -> K
+    where
+        K::Value: ut::UnifyValue<Error = ut::NoError>,
+    {
+        let vid = vid.into();
+        let mut relations = self.relations();
+        debug_assert!(relations.find(vid) == vid);
+        relations.union_value(vid, ty);
+
+        vid
+    }
+
+    pub fn find(&mut self, vid: I) -> K {
+        self.relations().find(vid)
+    }
+
+    pub fn unioned(&mut self, l: I, r: I) -> bool {
+        let mut relations = self.relations();
+        relations.find(l) == relations.find(r)
+    }
+
+    pub fn unify_var_value(
+        &mut self,
+        vid: I,
+        value: K::Value,
+    ) -> Result<(), <K::Value as ut::UnifyValue>::Error> {
+        let vid = self.find(vid);
+        self.relations().unify_var_value(vid, value)
+    }
+
+    pub fn unify_var_var(&mut self, a: I, b: I) -> Result<(), <K::Value as ut::UnifyValue>::Error> {
+        let mut relations = self.relations();
+        let a = relations.find(a);
+        let b = relations.find(b);
+        if a == b {
+            return Ok(());
+        }
+
+        relations.unify_var_var(a, b)?;
+
+        Ok(())
+    }
+
+    pub fn union_value(&mut self, vid: I, value: K::Value)
+    where
+        K::Value: ut::UnifyValue<Error = ut::NoError>,
+    {
+        let vid = self.find(vid).into();
+        self.instantiate(vid, value);
+    }
+
+    pub fn probe_value(&mut self, vid: I) -> K::Value {
+        self.relations().probe_value(vid)
+    }
+
+    #[inline(always)]
+    pub fn inlined_probe_value(&mut self, vid: I) -> K::Value {
+        self.relations().inlined_probe_value(vid)
+    }
+
+    pub fn new_key(&mut self, value: K::Value) -> K {
+        self.relations().new_key(value)
+    }
+
+    /// Clears any modifications currently tracked. Usually this can only be done once there are no
+    /// snapshots active as the modifications may otherwise be needed after a rollback
+    pub fn clear_modified_set(&mut self) {
+        self.storage.modified_set.clear();
+    }
+
+    /// Registers a watcher on the unifications done in this table
+    pub fn register_watcher(&mut self) -> ms::Offset<I> {
+        self.storage.modified_set.register()
+    }
+
+    /// Deregisters a watcher previously registered in this table
+    pub fn deregister_watcher(&mut self, offset: ms::Offset<I>) {
+        self.storage.modified_set.deregister(offset);
+    }
+
+    /// Watches the variable at `index` allowing any watchers to be notified to unifications with
+    /// `index`
+    pub fn watch_variable(&mut self, index: I) {
+        debug_assert!(index == self.relations().find(index).into());
+        self.storage.unify_log.watch_variable(index)
+    }
+
+    /// Unwatches a previous watch at `index`
+    pub fn unwatch_variable(&mut self, index: I) {
+        self.storage.unify_log.unwatch_variable(index)
+    }
+
+    /// Iterates through all unified variables since the last call to `notify_watcher`
+    /// passing the unified variable to `f`
+    pub fn notify_watcher(&mut self, offset: &ms::Offset<I>, mut f: impl FnMut(I)) {
+        let unify_log = &self.storage.unify_log;
+        self.storage.modified_set.notify_watcher(&mut self.undo_log, offset, |vid| {
+            for &unified_vid in unify_log.get(vid) {
+                f(unified_vid);
+            }
+
+            f(vid)
+        })
+    }
+}

compiler/rustc_data_structures/src/modified_set.rs

@@ -0,0 +1,133 @@
+use std::marker::PhantomData;
+
+use rustc_index::vec::Idx;
+
+use ena::undo_log::{Rollback, UndoLogs};
+
+#[derive(Copy, Clone, Debug)]
+enum UndoInner {
+    Add,
+    Drain { index: usize, offset: usize },
+}
+
+#[derive(Copy, Clone, Debug)]
+pub struct Undo<I>(UndoInner, PhantomData<I>);
+
+/// Tracks which indices have been modified and allows watchers to registered and notified of these
+/// changes.
+#[derive(Clone, Debug)]
+pub struct ModifiedSet<T: Idx> {
+    modified: Vec<T>,
+    offsets: Vec<usize>,
+}
+
+impl<T: Idx> Default for ModifiedSet<T> {
+    fn default() -> Self {
+        Self { modified: Default::default(), offsets: Vec::new() }
+    }
+}
+
+impl<T: Idx> ModifiedSet<T> {
+    /// Creates a new `ModifiedSet`
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Marks `index` as "modified". A subsequent call to `drain` will notify the callback with
+    /// `index`
+    pub fn set(&mut self, undo_log: &mut impl UndoLogs<Undo<T>>, index: T) {
+        self.modified.push(index);
+        undo_log.push(Undo(UndoInner::Add, PhantomData));
+    }
+
+    /// Calls `f` with all the indices that have been modified since the last call to
+    /// `notify_watcher`
+    pub fn notify_watcher(
+        &mut self,
+        undo_log: &mut impl UndoLogs<Undo<T>>,
+        watcher_offset: &Offset<T>,
+        mut f: impl FnMut(T),
+    ) {
+        let offset = &mut self.offsets[watcher_offset.index];
+        if *offset < self.modified.len() {
+            for &index in &self.modified[*offset..] {
+                f(index);
+            }
+            undo_log.push(Undo(
+                UndoInner::Drain { index: watcher_offset.index, offset: *offset },
+                PhantomData,
+            ));
+            *offset = self.modified.len();
+        }
+    }
+
+    /// Clears the set of all modifications that have been drained by all watchers
+    pub fn clear(&mut self) {
+        let min = self.offsets.iter().copied().min().unwrap_or_else(|| self.modified.len());
+        self.modified.drain(..min);
+        for offset in &mut self.offsets {
+            *offset -= min;
+        }
+    }
+
+    /// Registers a new watcher on this set.
+    ///
+    /// NOTE: Watchers must be removed in the reverse order that they were registered
+    pub fn register(&mut self) -> Offset<T> {
+        let index = self.offsets.len();
+        self.offsets.push(self.modified.len());
+        Offset { index, _marker: PhantomData }
+    }
+
+    /// De-registers a watcher on this set.
+    ///
+    /// NOTE: Watchers must be removed in the reverse order that they were registered
+    pub fn deregister(&mut self, offset: Offset<T>) {
+        assert_eq!(
+            offset.index,
+            self.offsets.len() - 1,
+            "Watchers must be removed in the reverse order that they were registered"
+        );
+        self.offsets.pop();
+        std::mem::forget(offset);
+    }
+}
+
+impl<I: Idx> Rollback<Undo<I>> for ModifiedSet<I> {
+    fn reverse(&mut self, undo: Undo<I>) {
+        match undo.0 {
+            UndoInner::Add => {
+                self.modified.pop();
+            }
+            UndoInner::Drain { index, offset } => {
+                if let Some(o) = self.offsets.get_mut(index) {
+                    *o = offset;
+                }
+            }
+        }
+    }
+}
+
+/// A registered offset into a `ModifiedSet`. Tracks how much a watcher has seen so far to avoid
+/// being notified of the same event twice.
+#[must_use]
+pub struct Offset<T> {
+    index: usize,
+    _marker: PhantomData<T>,
+}
+
+impl<T> Drop for Offset<T> {
+    fn drop(&mut self) {
+        if !std::thread::panicking() {
+            panic!("Offsets should be deregistered")
+        }
+    }
+}
+
+#[must_use]
+#[derive(Debug)]
+pub struct Snapshot<T> {
+    modified_len: usize,
+    undo_log_len: usize,
+    _marker: PhantomData<T>,
+}