src/librustc_data_structures/snapshot_map/mod.rs

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