src/librustc_data_structures/snapshot_map/mod.rs

   1 use crate::fx::FxHashMap;
   2 use std::hash::Hash;
   3 use std::ops;
   4 use std::mem;
   5
   6 #[cfg(test)]
   7 mod tests;
   8
   9 pub struct SnapshotMap<K, V>
  10     where K: Clone + Eq
  11 {
  12     map: FxHashMap<K, V>,
  13     undo_log: Vec<UndoLog<K, V>>,
  14     num_open_snapshots: usize,
  15 }
  16
  17 // HACK(eddyb) manual impl avoids `Default` bounds on `K` and `V`.
  18 impl<K, V> Default for SnapshotMap<K, V>
  19     where K: Hash + Clone + Eq
  20 {
  21     fn default() -> Self {
  22         SnapshotMap {
  23             map: Default::default(),
  24             undo_log: Default::default(),
  25             num_open_snapshots: 0,
  26         }
  27     }
  28 }
  29
  30 pub struct Snapshot {
  31     len: usize,
  32 }
  33
  34 enum UndoLog<K, V> {
  35     Inserted(K),
  36     Overwrite(K, V),
  37     Purged,
  38 }
  39
  40 impl<K, V> SnapshotMap<K, V>
  41     where K: Hash + Clone + Eq
  42 {
  43     pub fn clear(&mut self) {
  44         self.map.clear();
  45         self.undo_log.clear();
  46         self.num_open_snapshots = 0;
  47     }
  48
  49     fn in_snapshot(&self) -> bool {
  50         self.num_open_snapshots > 0
  51     }
  52
  53     pub fn insert(&mut self, key: K, value: V) -> bool {
  54         match self.map.insert(key.clone(), value) {
  55             None => {
  56                 if self.in_snapshot() {
  57                     self.undo_log.push(UndoLog::Inserted(key));
  58                 }
  59                 true
  60             }
  61             Some(old_value) => {
  62                 if self.in_snapshot() {
  63                     self.undo_log.push(UndoLog::Overwrite(key, old_value));
  64                 }
  65                 false
  66             }
  67         }
  68     }
  69
  70     pub fn remove(&mut self, key: K) -> bool {
  71         match self.map.remove(&key) {
  72             Some(old_value) => {
  73                 if self.in_snapshot() {
  74                     self.undo_log.push(UndoLog::Overwrite(key, old_value));
  75                 }
  76                 true
  77             }
  78             None => false,
  79         }
  80     }
  81
  82     pub fn get(&self, key: &K) -> Option<&V> {
  83         self.map.get(key)
  84     }
  85
  86     pub fn snapshot(&mut self) -> Snapshot {
  87         let len = self.undo_log.len();
  88         self.num_open_snapshots += 1;
  89         Snapshot { len }
  90     }
  91
  92     fn assert_open_snapshot(&self, snapshot: &Snapshot) {
  93         assert!(self.undo_log.len() >= snapshot.len);
  94         assert!(self.num_open_snapshots > 0);
  95     }
  96
  97     pub fn commit(&mut self, snapshot: Snapshot) {
  98         self.assert_open_snapshot(&snapshot);
  99         if self.num_open_snapshots == 1 {
 100             // The root snapshot. It's safe to clear the undo log because
 101             // there's no snapshot further out that we might need to roll back
 102             // to.
 103             assert!(snapshot.len == 0);
 104             self.undo_log.clear();
 105         }
 106
 107         self.num_open_snapshots -= 1;
 108     }
 109
 110     pub fn partial_rollback<F>(&mut self,
 111                                snapshot: &Snapshot,
 112                                should_revert_key: &F)
 113         where F: Fn(&K) -> bool
 114     {
 115         self.assert_open_snapshot(snapshot);
 116         for i in (snapshot.len .. self.undo_log.len()).rev() {
 117             let reverse = match self.undo_log[i] {
 118                 UndoLog::Purged => false,
 119                 UndoLog::Inserted(ref k) => should_revert_key(k),
 120                 UndoLog::Overwrite(ref k, _) => should_revert_key(k),
 121             };
 122
 123             if reverse {
 124                 let entry = mem::replace(&mut self.undo_log[i], UndoLog::Purged);
 125                 self.reverse(entry);
 126             }
 127         }
 128     }
 129
 130     pub fn rollback_to(&mut self, snapshot: Snapshot) {
 131         self.assert_open_snapshot(&snapshot);
 132         while self.undo_log.len() > snapshot.len {
 133             let entry = self.undo_log.pop().unwrap();
 134             self.reverse(entry);
 135         }
 136
 137         self.num_open_snapshots -= 1;
 138     }
 139
 140     fn reverse(&mut self, entry: UndoLog<K, V>) {
 141         match entry {
 142             UndoLog::Inserted(key) => {
 143                 self.map.remove(&key);
 144             }
 145
 146             UndoLog::Overwrite(key, old_value) => {
 147                 self.map.insert(key, old_value);
 148             }
 149
 150             UndoLog::Purged => {}
 151         }
 152     }
 153 }
 154
 155 impl<'k, K, V> ops::Index<&'k K> for SnapshotMap<K, V>
 156     where K: Hash + Clone + Eq
 157 {
 158     type Output = V;
 159     fn index(&self, key: &'k K) -> &V {
 160         &self.map[key]
 161     }
 162 }