Merge pull request #10 from rust-lang-nursery/ndm-reset-unifications

add `reset_unifications` API

Author: nikomatsakis

Cargo.toml

@@ -4,7 +4,7 @@ description = "Union-find, congruence closure, and other unification code. Based
 license = "MIT/Apache-2.0"
 homepage = "https://github.com/nikomatsakis/ena"
 repository = "https://github.com/nikomatsakis/ena"
-version = "0.9.2"
+version = "0.9.3"
 authors = ["Niko Matsakis <[email protected]>"]
 readme = "README.md"
 keywords = ["unification", "union-find"]

src/snapshot_vec.rs

@@ -119,6 +119,12 @@ impl<D: SnapshotVecDelegate> SnapshotVec<D> {
         &self.values[index]
     }
 
+    /// Reserve space for new values, just like an ordinary vec.
+    pub fn reserve(&mut self, additional: usize) {
+        // This is not affected by snapshots or anything.
+        self.values.reserve(additional);
+    }
+
     /// Returns a mutable pointer into the vec; whatever changes you make here cannot be undone
     /// automatically, so you should be sure call `record()` with some sort of suitable undo
     /// action.
@@ -135,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,10 +25,17 @@ 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>);
 
+    fn reserve(&mut self, num_new_values: usize);
+
     fn update<F>(&mut self, index: usize, op: F)
         where F: FnOnce(&mut VarValue<Self::Key>);
 
@@ -69,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()
@@ -79,6 +94,11 @@ impl<K: UnifyKey> UnificationStore for InPlace<K> {
         self.values.push(value);
     }
 
+    #[inline]
+    fn reserve(&mut self, num_new_values: usize) {
+        self.values.reserve(num_new_values);
+    }
+
     #[inline]
     fn update<F>(&mut self, index: usize, op: F)
         where F: FnOnce(&mut VarValue<Self::Key>)
@@ -137,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()
@@ -147,6 +180,11 @@ impl<K: UnifyKey> UnificationStore for Persistent<K> {
         self.values.push(value);
     }
 
+    #[inline]
+    fn reserve(&mut self, _num_new_values: usize) {
+        // not obviously relevant to DVec.
+    }
+
     #[inline]
     fn update<F>(&mut self, index: usize, op: F)
         where F: FnOnce(&mut VarValue<Self::Key>)

src/unify/mod.rs

@@ -274,6 +274,26 @@ impl<S: UnificationStore> UnificationTable<S> {
         key
     }
 
+    /// Reserve memory for `num_new_keys` to be created. Does not
+    /// actually create the new keys; you must then invoke `new_key`.
+    pub fn reserve(&mut self, num_new_keys: usize) {
+        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()