add a snapshottable hashmap

3 files changed, 189 insertions, 0 deletions

diff --git a/src/librustc_data_structures/lib.rs b/src/librustc_data_structures/lib.rs
index 926ee85230a..00f797d1b90 100644
--- a/src/librustc_data_structures/lib.rs
+++ b/src/librustc_data_structures/lib.rs
@@ -42,6 +42,7 @@ pub mod bitvec;
 pub mod graph;
 pub mod ivar;
 pub mod obligation_forest;
+pub mod snapshot_map;
 pub mod snapshot_vec;
 pub mod transitive_relation;
 pub mod unify;
diff --git a/src/librustc_data_structures/snapshot_map/mod.rs b/src/librustc_data_structures/snapshot_map/mod.rs
new file mode 100644
index 00000000000..b3989013d21
--- /dev/null
+++ b/src/librustc_data_structures/snapshot_map/mod.rs
@@ -0,0 +1,138 @@
+// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+use fnv::FnvHashMap;
+use std::hash::Hash;
+use std::ops;
+
+#[cfg(test)]
+mod test;
+
+pub struct SnapshotMap<K, V>
+    where K: Hash + Clone + Eq
+{
+    map: FnvHashMap<K, V>,
+    undo_log: Vec<UndoLog<K, V>>,
+}
+
+pub struct Snapshot {
+    len: usize
+}
+
+enum UndoLog<K, V> {
+    OpenSnapshot,
+    CommittedSnapshot,
+    Inserted(K),
+    Overwrite(K, V),
+}
+
+impl<K, V> SnapshotMap<K, V>
+    where K: Hash + Clone + Eq
+{
+    pub fn new() -> Self {
+        SnapshotMap {
+            map: FnvHashMap(),
+            undo_log: vec![]
+        }
+    }
+
+    pub fn insert(&mut self, key: K, value: V) -> bool {
+        match self.map.insert(key.clone(), value) {
+            None => {
+                if !self.undo_log.is_empty() {
+                    self.undo_log.push(UndoLog::Inserted(key));
+                }
+                true
+            }
+            Some(old_value) => {
+                if !self.undo_log.is_empty() {
+                    self.undo_log.push(UndoLog::Overwrite(key, old_value));
+                }
+                false
+            }
+        }
+    }
+
+    pub fn remove(&mut self, key: K) -> bool {
+        match self.map.remove(&key) {
+            Some(old_value) => {
+                if !self.undo_log.is_empty() {
+                    self.undo_log.push(UndoLog::Overwrite(key, old_value));
+                }
+                true
+            }
+            None => {
+                false
+            }
+        }
+    }
+
+    pub fn get(&self, key: &K) -> Option<&V> {
+        self.map.get(key)
+    }
+
+    pub fn snapshot(&mut self) -> Snapshot {
+        self.undo_log.push(UndoLog::OpenSnapshot);
+        let len = self.undo_log.len() - 1;
+        Snapshot { len: len }
+    }
+
+    fn assert_open_snapshot(&self, snapshot: &Snapshot) {
+        assert!(snapshot.len < self.undo_log.len());
+        assert!(match self.undo_log[snapshot.len] {
+            UndoLog::OpenSnapshot => true,
+            _ => false
+        });
+    }
+
+    pub fn commit(&mut self, snapshot: Snapshot) {
+        self.assert_open_snapshot(&snapshot);
+        if snapshot.len == 0 {
+            // The root snapshot.
+            self.undo_log.truncate(0);
+        } else {
+            self.undo_log[snapshot.len] = UndoLog::CommittedSnapshot;
+        }
+    }
+
+    pub fn rollback_to(&mut self, snapshot: Snapshot) {
+        self.assert_open_snapshot(&snapshot);
+        while self.undo_log.len() > snapshot.len + 1 {
+            match self.undo_log.pop().unwrap() {
+                UndoLog::OpenSnapshot => {
+                    panic!("cannot rollback an uncommitted snapshot");
+                }
+
+                UndoLog::CommittedSnapshot => { }
+
+                UndoLog::Inserted(key) => {
+                    self.map.remove(&key);
+                }
+
+                UndoLog::Overwrite(key, old_value) => {
+                    self.map.insert(key, old_value);
+                }
+            }
+        }
+
+        let v = self.undo_log.pop().unwrap();
+        assert!(match v { UndoLog::OpenSnapshot => true, _ => false });
+        assert!(self.undo_log.len() == snapshot.len);
+    }
+}
+
+impl<'k, K, V> ops::Index<&'k K> for SnapshotMap<K, V>
+    where K: Hash + Clone + Eq
+{
+    type Output = V;
+    fn index(&self, key: &'k K) -> &V {
+        &self.map[key]
+    }
+}
diff --git a/src/librustc_data_structures/snapshot_map/test.rs b/src/librustc_data_structures/snapshot_map/test.rs
new file mode 100644
index 00000000000..4114082839b
--- /dev/null
+++ b/src/librustc_data_structures/snapshot_map/test.rs
@@ -0,0 +1,50 @@
+// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+use super::SnapshotMap;
+
+#[test]
+fn basic() {
+    let mut map = SnapshotMap::new();
+    map.insert(22, "twenty-two");
+    let snapshot = map.snapshot();
+    map.insert(22, "thirty-three");
+    assert_eq!(map[&22], "thirty-three");
+    map.insert(44, "fourty-four");
+    assert_eq!(map[&44], "fourty-four");
+    assert_eq!(map.get(&33), None);
+    map.rollback_to(snapshot);
+    assert_eq!(map[&22], "twenty-two");
+    assert_eq!(map.get(&33), None);
+    assert_eq!(map.get(&44), None);
+}
+
+#[test]
+#[should_panic]
+fn out_of_order() {
+    let mut map = SnapshotMap::new();
+    map.insert(22, "twenty-two");
+    let snapshot1 = map.snapshot();
+    let _snapshot2 = map.snapshot();
+    map.rollback_to(snapshot1);
+}
+
+#[test]
+fn nested_commit_then_rollback() {
+    let mut map = SnapshotMap::new();
+    map.insert(22, "twenty-two");
+    let snapshot1 = map.snapshot();
+    let snapshot2 = map.snapshot();
+    map.insert(22, "thirty-three");
+    map.commit(snapshot2);
+    assert_eq!(map[&22], "thirty-three");
+    map.rollback_to(snapshot1);
+    assert_eq!(map[&22], "twenty-two");
+}