perf: Only process changed obligations in ObligationForest

This rewrites most of ObligationForest forest so to avoid iterating
through the entire set of obligations on each `select` attempt. Instead
only the obligations that can actually make progress are processed. This
gives great speedups in benchmarks such as `inflate` which create a
large number of pending obligations.

To support this the unification tables where extended to to keep track
of which type inference variables that has actually changed at each
step. Which then lets `ObligationForest` get a list of only the changed
variables at each step which it can map back to its obligations.

In addition to this primary change, many of the other iterations in
`ObligationForest` were refactored to only process lists of the nodes
that they actually are interested in. The extra bookkeeping needed for
this was possible without the primary change but were a performance
regressions there as they slowed down the main loop. As the main loop is
no longer the main issue these optimizations could be re-applied.

Author: Markus Westerlind

Cargo.lock

@@ -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/lib.rs

@@ -0,0 +1,172 @@
+use std::ops::Range;
+
+use rustc_index::vec::{Idx, IndexVec};
+
+use crate::modified_set as ms;
+use crate::unify as ut;
+use crate::unify_log as ul;
+
+pub struct LoggedUnificationTable<K: ut::UnifyKey, I: Idx = K> {
+    relations: ut::UnificationTable<ut::InPlace<K>>,
+    unify_log: ul::UnifyLog<I>,
+    modified_set: ms::ModifiedSet<I>,
+    reference_counts: IndexVec<I, u32>,
+}
+
+impl<K, I> LoggedUnificationTable<K, I>
+where
+    K: ut::UnifyKey + From<I>,
+    I: Idx + From<K>,
+{
+    pub fn new() -> Self {
+        Self {
+            relations: ut::UnificationTable::new(),
+            unify_log: ul::UnifyLog::new(),
+            modified_set: ms::ModifiedSet::new(),
+            reference_counts: IndexVec::new(),
+        }
+    }
+
+    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>,
+    {
+        if !self.unify_log.get(vid).is_empty()
+            || self.reference_counts.get(vid).map_or(false, |c| *c != 0)
+        {
+            self.modified_set.set(vid);
+        }
+        let vid = vid.into();
+        debug_assert!(self.relations.find(vid) == vid);
+        self.relations.union_value(vid, ty);
+
+        vid
+    }
+
+    pub fn find(&mut self, vid: I) -> K {
+        self.relations.find(vid)
+    }
+
+    pub fn unify_var_value(
+        &mut self,
+        vid: I,
+        value: K::Value,
+    ) -> Result<(), <K::Value as ut::UnifyValue>::Error> {
+        let vid = self.find(vid).into();
+        if !self.unify_log.get(vid).is_empty()
+            || self.reference_counts.get(vid).map_or(false, |c| *c != 0)
+        {
+            self.modified_set.set(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 a_root = self.relations.find(a);
+        let b_root = self.relations.find(b);
+        if a_root == b_root {
+            return Ok(());
+        }
+
+        self.relations.unify_var_var(a_root, b_root)?;
+
+        if a_root == self.relations.find(a_root) {
+            self.unify_log.unify(a_root.into(), b_root.into());
+        } else {
+            self.unify_log.unify(b_root.into(), a_root.into());
+        }
+        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)
+    }
+
+    pub fn len(&self) -> usize {
+        self.relations.len()
+    }
+
+    pub fn snapshot(&mut self) -> Snapshot<K, I> {
+        Snapshot {
+            snapshot: self.relations.snapshot(),
+            unify_log_snapshot: self.unify_log.snapshot(),
+            modified_snapshot: self.modified_set.snapshot(),
+        }
+    }
+
+    pub fn rollback_to(&mut self, s: Snapshot<K, I>) {
+        let Snapshot { snapshot, unify_log_snapshot, modified_snapshot } = s;
+        self.relations.rollback_to(snapshot);
+        self.unify_log.rollback_to(unify_log_snapshot);
+        self.modified_set.rollback_to(modified_snapshot);
+    }
+
+    pub fn commit(&mut self, s: Snapshot<K, I>) {
+        let Snapshot { snapshot, unify_log_snapshot, modified_snapshot } = s;
+        self.relations.commit(snapshot);
+        self.unify_log.commit(unify_log_snapshot);
+        self.modified_set.commit(modified_snapshot);
+    }
+
+    pub fn vars_since_snapshot(&mut self, s: &Snapshot<K, I>) -> Range<K> {
+        self.relations.vars_since_snapshot(&s.snapshot)
+    }
+
+    pub fn register(&mut self) -> ms::Offset<I> {
+        self.modified_set.register()
+    }
+
+    pub fn deregister(&mut self, offset: ms::Offset<I>) {
+        self.modified_set.deregister(offset);
+    }
+
+    pub fn watch_variable(&mut self, index: I) {
+        self.reference_counts.ensure_contains_elem(index, || 0);
+        self.reference_counts[index] += 1;
+    }
+
+    pub fn unwatch_variable(&mut self, index: I) {
+        self.reference_counts[index] -= 1;
+    }
+
+    pub fn drain_modified_set(&mut self, offset: &ms::Offset<I>, mut f: impl FnMut(I) -> bool) {
+        let unify_log = &self.unify_log;
+        self.modified_set.drain(offset, |vid| {
+            for &unified_vid in unify_log.get(vid) {
+                f(unified_vid);
+            }
+
+            f(vid)
+        })
+    }
+}
+
+pub struct Snapshot<K: ut::UnifyKey, I: Idx = K> {
+    snapshot: ut::Snapshot<ut::InPlace<K>>,
+    unify_log_snapshot: ul::Snapshot<I>,
+    modified_snapshot: ms::Snapshot<I>,
+}

compiler/rustc_data_structures/src/logged_unification_table.rs

@@ -0,0 +1,139 @@
+use std::{collections::VecDeque, marker::PhantomData};
+
+use rustc_index::{bit_set::BitSet, vec::Idx};
+
+#[derive(Clone, Debug)]
+pub struct ModifiedSet<T: Idx> {
+    modified: VecDeque<T>,
+    snapshots: usize,
+    modified_set: BitSet<T>,
+    undo_offsets: Vec<usize>,
+    offsets: Vec<usize>,
+}
+
+impl<T: Idx> Default for ModifiedSet<T> {
+    fn default() -> Self {
+        Self {
+            modified: Default::default(),
+            snapshots: 0,
+            modified_set: BitSet::new_empty(0),
+            offsets: Vec::new(),
+            undo_offsets: Vec::new(),
+        }
+    }
+}
+
+impl<T: Idx> ModifiedSet<T> {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn set(&mut self, index: T) {
+        if index.index() >= self.modified_set.domain_size() {
+            self.modified_set.resize(index.index() + 1);
+        }
+        if self.modified_set.insert(index) {
+            self.modified.push_back(index);
+        }
+    }
+
+    pub fn drain(&mut self, offset: &Offset<T>, mut f: impl FnMut(T) -> bool) {
+        let offset = &mut self.offsets[offset.index];
+        for &index in self.modified.iter().skip(*offset) {
+            if f(index) {}
+        }
+        *offset = self.modified.len();
+    }
+
+    pub fn snapshot(&mut self) -> Snapshot<T> {
+        self.snapshots += 1;
+        let offsets_start = self.undo_offsets.len();
+        self.undo_offsets.extend_from_slice(&self.offsets);
+        Snapshot {
+            modified_len: self.modified.len(),
+            offsets_start,
+            offsets_len: self.offsets.len(),
+            _marker: PhantomData,
+        }
+    }
+
+    pub fn rollback_to(&mut self, snapshot: Snapshot<T>) {
+        self.snapshots -= 1;
+        for &index in self.modified.iter().skip(snapshot.modified_len) {
+            self.modified_set.remove(index);
+        }
+        self.modified.truncate(snapshot.modified_len);
+        let mut offsets = self.offsets.iter_mut();
+        for (offset, &saved_offset) in offsets.by_ref().zip(
+            &self.undo_offsets
+                [snapshot.offsets_start..snapshot.offsets_start + snapshot.offsets_len],
+        ) {
+            *offset = saved_offset;
+        }
+        for offset in offsets {
+            *offset = self.modified.len().min(*offset);
+        }
+        self.undo_offsets.truncate(snapshot.offsets_start);
+
+        if self.snapshots == 0 {
+            let min = self.offsets.iter().copied().min().unwrap_or(0);
+            // Any indices still in `modified` may not have been instantiated, so if we observe them again
+            // we need to notify any listeners again
+            for index in self.modified.drain(..min) {
+                self.modified_set.remove(index);
+            }
+            for offset in &mut self.offsets {
+                *offset -= min;
+            }
+        }
+    }
+
+    pub fn commit(&mut self, snapshot: Snapshot<T>) {
+        self.snapshots -= 1;
+        self.undo_offsets.truncate(snapshot.offsets_start);
+        if self.snapshots == 0 {
+            // Everything up until this point is committed, so we can forget anything before the
+            // current offsets
+            let min = self.offsets.iter().copied().min().unwrap_or(0);
+            self.modified.drain(..min);
+            for offset in &mut self.offsets {
+                *offset -= min;
+            }
+        }
+    }
+
+    pub fn register(&mut self) -> Offset<T> {
+        let index = self.offsets.len();
+        self.offsets.push(0);
+        Offset { index, _marker: PhantomData }
+    }
+
+    pub fn deregister(&mut self, offset: Offset<T>) {
+        assert_eq!(offset.index, self.offsets.len() - 1);
+        self.offsets.pop();
+        std::mem::forget(offset);
+    }
+}
+
+#[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,
+    offsets_start: usize,
+    offsets_len: usize,
+    _marker: PhantomData<T>,
+}