]>
Commit | Line | Data |
---|---|---|
9fa01778 | 1 | use crate::fx::FxHashMap; |
3157f602 XL |
2 | use std::hash::Hash; |
3 | use std::ops; | |
c30ab7b3 | 4 | use std::mem; |
3157f602 XL |
5 | |
6 | #[cfg(test)] | |
416331ca | 7 | mod tests; |
3157f602 XL |
8 | |
9 | pub struct SnapshotMap<K, V> | |
e74abb32 | 10 | where K: Clone + Eq |
3157f602 | 11 | { |
476ff2be | 12 | map: FxHashMap<K, V>, |
3157f602 | 13 | undo_log: Vec<UndoLog<K, V>>, |
a1dfa0c6 | 14 | num_open_snapshots: usize, |
3157f602 XL |
15 | } |
16 | ||
a1dfa0c6 | 17 | // HACK(eddyb) manual impl avoids `Default` bounds on `K` and `V`. |
0bf4aa26 | 18 | impl<K, V> Default for SnapshotMap<K, V> |
3157f602 XL |
19 | where K: Hash + Clone + Eq |
20 | { | |
0bf4aa26 | 21 | fn default() -> Self { |
3157f602 | 22 | SnapshotMap { |
a1dfa0c6 XL |
23 | map: Default::default(), |
24 | undo_log: Default::default(), | |
25 | num_open_snapshots: 0, | |
3157f602 XL |
26 | } |
27 | } | |
0bf4aa26 | 28 | } |
3157f602 | 29 | |
a1dfa0c6 XL |
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 | ||
0bf4aa26 XL |
40 | impl<K, V> SnapshotMap<K, V> |
41 | where K: Hash + Clone + Eq | |
42 | { | |
0531ce1d XL |
43 | pub fn clear(&mut self) { |
44 | self.map.clear(); | |
45 | self.undo_log.clear(); | |
a1dfa0c6 XL |
46 | self.num_open_snapshots = 0; |
47 | } | |
48 | ||
49 | fn in_snapshot(&self) -> bool { | |
50 | self.num_open_snapshots > 0 | |
0531ce1d XL |
51 | } |
52 | ||
3157f602 XL |
53 | pub fn insert(&mut self, key: K, value: V) -> bool { |
54 | match self.map.insert(key.clone(), value) { | |
55 | None => { | |
a1dfa0c6 | 56 | if self.in_snapshot() { |
3157f602 XL |
57 | self.undo_log.push(UndoLog::Inserted(key)); |
58 | } | |
59 | true | |
60 | } | |
61 | Some(old_value) => { | |
a1dfa0c6 | 62 | if self.in_snapshot() { |
3157f602 XL |
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) => { | |
a1dfa0c6 | 73 | if self.in_snapshot() { |
3157f602 XL |
74 | self.undo_log.push(UndoLog::Overwrite(key, old_value)); |
75 | } | |
76 | true | |
77 | } | |
c30ab7b3 | 78 | None => false, |
3157f602 XL |
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 { | |
a1dfa0c6 XL |
87 | let len = self.undo_log.len(); |
88 | self.num_open_snapshots += 1; | |
b7449926 | 89 | Snapshot { len } |
3157f602 XL |
90 | } |
91 | ||
92 | fn assert_open_snapshot(&self, snapshot: &Snapshot) { | |
a1dfa0c6 XL |
93 | assert!(self.undo_log.len() >= snapshot.len); |
94 | assert!(self.num_open_snapshots > 0); | |
3157f602 XL |
95 | } |
96 | ||
a1dfa0c6 XL |
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(); | |
3157f602 | 105 | } |
a1dfa0c6 XL |
106 | |
107 | self.num_open_snapshots -= 1; | |
3157f602 XL |
108 | } |
109 | ||
c30ab7b3 SL |
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); | |
a1dfa0c6 | 116 | for i in (snapshot.len .. self.undo_log.len()).rev() { |
c30ab7b3 | 117 | let reverse = match self.undo_log[i] { |
a1dfa0c6 | 118 | UndoLog::Purged => false, |
c30ab7b3 SL |
119 | UndoLog::Inserted(ref k) => should_revert_key(k), |
120 | UndoLog::Overwrite(ref k, _) => should_revert_key(k), | |
121 | }; | |
122 | ||
123 | if reverse { | |
a1dfa0c6 | 124 | let entry = mem::replace(&mut self.undo_log[i], UndoLog::Purged); |
c30ab7b3 SL |
125 | self.reverse(entry); |
126 | } | |
127 | } | |
128 | } | |
129 | ||
a1dfa0c6 XL |
130 | pub fn rollback_to(&mut self, snapshot: Snapshot) { |
131 | self.assert_open_snapshot(&snapshot); | |
132 | while self.undo_log.len() > snapshot.len { | |
c30ab7b3 SL |
133 | let entry = self.undo_log.pop().unwrap(); |
134 | self.reverse(entry); | |
135 | } | |
136 | ||
a1dfa0c6 | 137 | self.num_open_snapshots -= 1; |
c30ab7b3 | 138 | } |
3157f602 | 139 | |
c30ab7b3 SL |
140 | fn reverse(&mut self, entry: UndoLog<K, V>) { |
141 | match entry { | |
c30ab7b3 SL |
142 | UndoLog::Inserted(key) => { |
143 | self.map.remove(&key); | |
3157f602 | 144 | } |
3157f602 | 145 | |
c30ab7b3 SL |
146 | UndoLog::Overwrite(key, old_value) => { |
147 | self.map.insert(key, old_value); | |
148 | } | |
149 | ||
a1dfa0c6 | 150 | UndoLog::Purged => {} |
c30ab7b3 | 151 | } |
3157f602 XL |
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 | } |