add a reset_unifications method

Author: nikomatsakis

src/snapshot_vec.rs

@@ -141,6 +141,20 @@ impl<D: SnapshotVecDelegate> SnapshotVec<D> {
         }
     }
 
+    /// Updates all elements. Potentially more efficient -- but
+    /// otherwise equivalent to -- invoking `set` for each element.
+    pub fn set_all(&mut self, mut new_elems: impl FnMut(usize) -> D::Value) {
+        if !self.in_snapshot() {
+            for (slot, index) in self.values.iter_mut().zip(0..) {
+                *slot = new_elems(index);
+            }
+        } else {
+            for i in 0..self.values.len() {
+                self.set(i, new_elems(i));
+            }
+        }
+    }
+
     pub fn update<OP>(&mut self, index: usize, op: OP)
     where
         OP: FnOnce(&mut D::Value),

src/unify/backing_vec.rs

@@ -25,6 +25,11 @@ pub trait UnificationStore: ops::Index<usize, Output = VarValue<Key<Self>>> + Cl
 
     fn commit(&mut self, snapshot: Self::Snapshot);
 
+    fn reset_unifications(
+        &mut self,
+        value: impl FnMut(u32) -> VarValue<Self::Key>,
+    );
+
     fn len(&self) -> usize;
 
     fn push(&mut self, value: VarValue<Self::Key>);
@@ -71,6 +76,14 @@ impl<K: UnifyKey> UnificationStore for InPlace<K> {
         self.values.commit(snapshot);
     }
 
+    #[inline]
+    fn reset_unifications(
+        &mut self,
+        mut value: impl FnMut(u32) -> VarValue<Self::Key>,
+    ) {
+        self.values.set_all(|i| value(i as u32));
+    }
+
     #[inline]
     fn len(&self) -> usize {
         self.values.len()
@@ -144,6 +157,19 @@ impl<K: UnifyKey> UnificationStore for Persistent<K> {
     fn commit(&mut self, _snapshot: Self::Snapshot) {
     }
 
+    #[inline]
+    fn reset_unifications(
+        &mut self,
+        mut value: impl FnMut(u32) -> VarValue<Self::Key>,
+    ) {
+        // Without extending dogged, there isn't obviously a more
+        // efficient way to do this. But it's pretty dumb. Maybe
+        // dogged needs a `map`.
+        for i in 0 .. self.values.len() {
+            self.values[i] = value(i as u32);
+        }
+    }
+
     #[inline]
     fn len(&self) -> usize {
         self.values.len()

src/unify/mod.rs

@@ -280,6 +280,20 @@ impl<S: UnificationStore> UnificationTable<S> {
         self.values.reserve(num_new_keys);
     }
 
+    /// Clears all unifications that have been performed, resetting to
+    /// the initial state. The values of each variable are given by
+    /// the closure.
+    pub fn reset_unifications(
+        &mut self,
+        mut value: impl FnMut(S::Key) -> S::Value,
+    ) {
+        self.values.reset_unifications(|i| {
+            let key = UnifyKey::from_index(i as u32);
+            let value = value(key);
+            VarValue::new_var(key, value)
+        });
+    }
+
     /// Returns the number of keys created so far.
     pub fn len(&self) -> usize {
         self.values.len()