refactor: Make UnifyLog have lower overhead snapshots

Author: Markus Westerlind

compiler/rustc_data_structures/src/logged_unification_table.rs

@@ -1,6 +1,6 @@
 use std::ops::Range;
 
-use rustc_index::vec::{Idx, IndexVec};
+use rustc_index::vec::Idx;
 
 use crate::modified_set as ms;
 use crate::unify as ut;
@@ -10,7 +10,6 @@ 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>
@@ -23,7 +22,6 @@ where
             relations: ut::UnificationTable::new(),
             unify_log: ul::UnifyLog::new(),
             modified_set: ms::ModifiedSet::new(),
-            reference_counts: IndexVec::new(),
         }
     }
 
@@ -38,7 +36,7 @@ where
     where
         K::Value: ut::UnifyValue<Error = ut::NoError>,
     {
-        if self.needs_log(vid) {
+        if self.unify_log.needs_log(vid) {
             warn!("ModifiedSet {:?} => {:?}", vid, ty);
             self.modified_set.set(vid);
         }
@@ -59,7 +57,7 @@ where
         value: K::Value,
     ) -> Result<(), <K::Value as ut::UnifyValue>::Error> {
         let vid = self.find(vid).into();
-        if self.needs_log(vid) {
+        if self.unify_log.needs_log(vid) {
             self.modified_set.set(vid);
         }
         self.relations.unify_var_value(vid, value)
@@ -75,24 +73,13 @@ where
         self.relations.unify_var_var(a, b)?;
 
         if a == self.relations.find(a) {
-            if self.needs_log(b.into()) {
-                warn!("Log: {:?} {:?} => {:?}", a, b, I::from(self.relations.find(a)));
-                self.unify_log.unify(a.into(), b.into());
-            }
+            self.unify_log.unify(a.into(), b.into());
         } else {
-            if self.needs_log(a.into()) {
-                warn!("Log: {:?} {:?} => {:?}", a, b, I::from(self.relations.find(a)));
-                self.unify_log.unify(b.into(), a.into());
-            }
+            self.unify_log.unify(b.into(), a.into());
         }
         Ok(())
     }
 
-    fn needs_log(&self, vid: I) -> bool {
-        !self.unify_log.get(vid).is_empty()
-            || self.reference_counts.get(vid).map_or(false, |c| *c != 0)
-    }
-
     pub fn union_value(&mut self, vid: I, value: K::Value)
     where
         K::Value: ut::UnifyValue<Error = ut::NoError>,
@@ -154,12 +141,11 @@ where
 
     pub fn watch_variable(&mut self, index: I) {
         debug_assert!(index == self.relations.find(index).into());
-        self.reference_counts.ensure_contains_elem(index, || 0);
-        self.reference_counts[index] += 1;
+        self.unify_log.watch_variable(index)
     }
 
     pub fn unwatch_variable(&mut self, index: I) {
-        self.reference_counts[index] -= 1;
+        self.unify_log.unwatch_variable(index)
     }
 
     pub fn drain_modified_set(&mut self, offset: &ms::Offset<I>, mut f: impl FnMut(I) -> bool) {

compiler/rustc_data_structures/src/unify_log.rs

@@ -2,32 +2,105 @@ use std::marker::PhantomData;
 
 use rustc_index::vec::{Idx, IndexVec};
 
+enum Undo<T> {
+    Move { index: T, old: usize },
+    Extend { group_index: usize, len: usize },
+    NewGroup { index: T },
+}
+
 pub struct UnifyLog<T: Idx> {
-    unified_vars: IndexVec<T, Vec<T>>,
-    undo_log: Vec<(T, u32)>,
+    unified_vars: IndexVec<T, usize>,
+    groups: Vec<Vec<T>>,
+    undo_log: Vec<Undo<T>>,
+    reference_counts: IndexVec<T, u32>,
     snapshots: usize,
 }
 
+fn pick2_mut<T, I: Idx>(self_: &mut [T], a: I, b: I) -> (&mut T, &mut T) {
+    let (ai, bi) = (a.index(), b.index());
+    assert!(ai != bi);
+
+    if ai < bi {
+        let (c1, c2) = self_.split_at_mut(bi);
+        (&mut c1[ai], &mut c2[0])
+    } else {
+        let (c2, c1) = pick2_mut(self_, b, a);
+        (c1, c2)
+    }
+}
+
 impl<T: Idx> UnifyLog<T> {
     pub fn new() -> Self {
-        UnifyLog { unified_vars: IndexVec::new(), undo_log: Vec::new(), snapshots: 0 }
+        UnifyLog {
+            unified_vars: IndexVec::new(),
+            groups: Vec::new(),
+            undo_log: Vec::new(),
+            reference_counts: IndexVec::new(),
+            snapshots: 0,
+        }
     }
 
     pub fn unify(&mut self, root: T, other: T) {
-        self.unified_vars.ensure_contains_elem(root, Vec::new);
-        self.unified_vars.ensure_contains_elem(other, Vec::new);
-        let (root_ids, other_ids) = self.unified_vars.pick2_mut(root, other);
-        self.undo_log.push((root, root_ids.len() as u32));
-        for &other in &*other_ids {
-            if !root_ids.contains(&other) {
-                root_ids.push(other);
+        if !self.needs_log(other) {
+            return;
+        }
+        self.unified_vars.ensure_contains_elem(root, usize::max_value);
+        self.unified_vars.ensure_contains_elem(other, usize::max_value);
+        let mut root_group = self.unified_vars[root];
+        let other_group = self.unified_vars[other];
+
+        if other_group == usize::max_value() {
+            let root_vec = if root_group == usize::max_value() {
+                root_group = self.groups.len();
+                self.unified_vars[root] = root_group;
+                self.groups.push(Vec::new());
+                self.undo_log.push(Undo::NewGroup { index: root });
+                self.groups.last_mut().unwrap()
+            } else {
+                let root_vec = &mut self.groups[root_group];
+                self.undo_log.push(Undo::Extend { group_index: root_group, len: root_vec.len() });
+                root_vec
+            };
+            root_vec.push(other);
+        } else {
+            if root_group == usize::max_value() {
+                let group = &mut self.unified_vars[root];
+                self.undo_log.push(Undo::Move { index: root, old: *group });
+                *group = other_group;
+                self.groups[other_group].push(other);
+            } else {
+                let (root_vec, other_vec) = pick2_mut(&mut self.groups, root_group, other_group);
+                self.undo_log.push(Undo::Extend { group_index: root_group, len: root_vec.len() });
+                root_vec.extend_from_slice(other_vec);
+
+                if self.reference_counts.get(other).map_or(false, |c| *c != 0) {
+                    root_vec.push(other);
+                }
             }
         }
-        root_ids.push(other);
     }
 
     pub fn get(&self, root: T) -> &[T] {
-        self.unified_vars.get(root).map(|v| &v[..]).unwrap_or(&[][..])
+        match self.unified_vars.get(root) {
+            Some(group) => match self.groups.get(*group) {
+                Some(v) => v,
+                None => &[],
+            },
+            None => &[],
+        }
+    }
+
+    pub fn needs_log(&self, vid: T) -> bool {
+        !self.get(vid).is_empty() || self.reference_counts.get(vid).map_or(false, |c| *c != 0)
+    }
+
+    pub fn watch_variable(&mut self, index: T) {
+        self.reference_counts.ensure_contains_elem(index, || 0);
+        self.reference_counts[index] += 1;
+    }
+
+    pub fn unwatch_variable(&mut self, index: T) {
+        self.reference_counts[index] -= 1;
     }
 
     pub fn snapshot(&mut self) -> Snapshot<T> {
@@ -44,8 +117,17 @@ impl<T: Idx> UnifyLog<T> {
 
     pub fn rollback_to(&mut self, snapshot: Snapshot<T>) {
         self.snapshots -= 1;
-        for (index, len) in self.undo_log.drain(snapshot.undo_log_len as usize..) {
-            self.unified_vars[index].truncate(len as usize);
+        for undo in self.undo_log.drain(snapshot.undo_log_len as usize..).rev() {
+            match undo {
+                Undo::Extend { group_index, len } => {
+                    self.groups[group_index].truncate(len as usize)
+                }
+                Undo::Move { index, old } => self.unified_vars[index] = old,
+                Undo::NewGroup { index } => {
+                    self.groups.pop();
+                    self.unified_vars[index] = usize::max_value();
+                }
+            }
         }
     }
 }