src/librustc_data_structures/snapshot_map/mod.rs

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