miden-crypto/src/utils/kv_map.rs
src/utils/kv_map.rsRaw
1use core::cell::RefCell;
2use winter_utils::{
3 collections::{btree_map::IntoIter, BTreeMap, BTreeSet},
4 Box,
5};
6
7// KEY-VALUE MAP TRAIT
8// ================================================================================================
9
10/// A trait that defines the interface for a key-value map.
11pub trait KvMap<K: Ord + Clone, V: Clone>:
12 Extend<(K, V)> + FromIterator<(K, V)> + IntoIterator<Item = (K, V)>
13{
14 fn get(&self, key: &K) -> Option<&V>;
15 fn contains_key(&self, key: &K) -> bool;
16 fn len(&self) -> usize;
17 fn is_empty(&self) -> bool {
18 self.len() == 0
19 }
20 fn insert(&mut self, key: K, value: V) -> Option<V>;
21
22 fn iter(&self) -> Box<dyn Iterator<Item = (&K, &V)> + '_>;
23}
24
25// BTREE MAP `KvMap` IMPLEMENTATION
26// ================================================================================================
27
28impl<K: Ord + Clone, V: Clone> KvMap<K, V> for BTreeMap<K, V> {
29 fn get(&self, key: &K) -> Option<&V> {
30 self.get(key)
31 }
32
33 fn contains_key(&self, key: &K) -> bool {
34 self.contains_key(key)
35 }
36
37 fn len(&self) -> usize {
38 self.len()
39 }
40
41 fn insert(&mut self, key: K, value: V) -> Option<V> {
42 self.insert(key, value)
43 }
44
45 fn iter(&self) -> Box<dyn Iterator<Item = (&K, &V)> + '_> {
46 Box::new(self.iter())
47 }
48}
49
50// RECORDING MAP
51// ================================================================================================
52
53/// A [RecordingMap] that records read requests to the underlying key-value map.
54///
55/// The data recorder is used to generate a proof for read requests.
56///
57/// The [RecordingMap] is composed of three parts:
58/// - `data`: which contains the current set of key-value pairs in the map.
59/// - `updates`: which tracks keys for which values have been since the map was instantiated.
60/// updates include both insertions and updates of values under existing keys.
61/// - `trace`: which contains the key-value pairs from the original data which have been accesses
62/// since the map was instantiated.
63#[derive(Debug, Default, Clone, Eq, PartialEq)]
64pub struct RecordingMap<K, V> {
65 data: BTreeMap<K, V>,
66 updates: BTreeSet<K>,
67 trace: RefCell<BTreeMap<K, V>>,
68}
69
70impl<K: Ord + Clone, V: Clone> RecordingMap<K, V> {
71 // CONSTRUCTOR
72 // --------------------------------------------------------------------------------------------
73 /// Returns a new [RecordingMap] instance initialized with the provided key-value pairs.
74 /// ([BTreeMap]).
75 pub fn new(init: impl IntoIterator<Item = (K, V)>) -> Self {
76 RecordingMap {
77 data: init.into_iter().collect(),
78 updates: BTreeSet::new(),
79 trace: RefCell::new(BTreeMap::new()),
80 }
81 }
82
83 // FINALIZER
84 // --------------------------------------------------------------------------------------------
85
86 /// Consumes the [RecordingMap] and returns a [BTreeMap] containing the key-value pairs from
87 /// the initial data set that were read during recording.
88 pub fn into_proof(self) -> BTreeMap<K, V> {
89 self.trace.take()
90 }
91
92 // TEST HELPERS
93 // --------------------------------------------------------------------------------------------
94
95 #[cfg(test)]
96 pub fn trace_len(&self) -> usize {
97 self.trace.borrow().len()
98 }
99
100 #[cfg(test)]
101 pub fn updates_len(&self) -> usize {
102 self.updates.len()
103 }
104}
105
106impl<K: Ord + Clone, V: Clone> KvMap<K, V> for RecordingMap<K, V> {
107 // PUBLIC ACCESSORS
108 // --------------------------------------------------------------------------------------------
109
110 /// Returns a reference to the value associated with the given key if the value exists.
111 ///
112 /// If the key is part of the initial data set, the key access is recorded.
113 fn get(&self, key: &K) -> Option<&V> {
114 self.data.get(key).map(|value| {
115 if !self.updates.contains(key) {
116 self.trace.borrow_mut().insert(key.clone(), value.clone());
117 }
118 value
119 })
120 }
121
122 /// Returns a boolean to indicate whether the given key exists in the data set.
123 ///
124 /// If the key is part of the initial data set, the key access is recorded.
125 fn contains_key(&self, key: &K) -> bool {
126 self.get(key).is_some()
127 }
128
129 /// Returns the number of key-value pairs in the data set.
130 fn len(&self) -> usize {
131 self.data.len()
132 }
133
134 // MUTATORS
135 // --------------------------------------------------------------------------------------------
136
137 /// Inserts a key-value pair into the data set.
138 ///
139 /// If the key already exists in the data set, the value is updated and the old value is
140 /// returned.
141 fn insert(&mut self, key: K, value: V) -> Option<V> {
142 let new_update = self.updates.insert(key.clone());
143 self.data.insert(key.clone(), value).map(|old_value| {
144 if new_update {
145 self.trace.borrow_mut().insert(key, old_value.clone());
146 }
147 old_value
148 })
149 }
150
151 // ITERATION
152 // --------------------------------------------------------------------------------------------
153
154 /// Returns an iterator over the key-value pairs in the data set.
155 fn iter(&self) -> Box<dyn Iterator<Item = (&K, &V)> + '_> {
156 Box::new(self.data.iter())
157 }
158}
159
160impl<K: Clone + Ord, V: Clone> Extend<(K, V)> for RecordingMap<K, V> {
161 fn extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T) {
162 iter.into_iter().for_each(move |(k, v)| {
163 self.insert(k, v);
164 });
165 }
166}
167
168impl<K: Clone + Ord, V: Clone> FromIterator<(K, V)> for RecordingMap<K, V> {
169 fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
170 Self::new(iter)
171 }
172}
173
174impl<K: Clone + Ord, V: Clone> IntoIterator for RecordingMap<K, V> {
175 type Item = (K, V);
176 type IntoIter = IntoIter<K, V>;
177
178 fn into_iter(self) -> Self::IntoIter {
179 self.data.into_iter()
180 }
181}
182
183// TESTS
184// ================================================================================================
185
186#[cfg(test)]
187mod tests {
188 use super::*;
189
190 const ITEMS: [(u64, u64); 5] = [(0, 0), (1, 1), (2, 2), (3, 3), (4, 4)];
191
192 #[test]
193 fn test_get_item() {
194 // instantiate a recording map
195 let map = RecordingMap::new(ITEMS.to_vec());
196
197 // get a few items
198 let get_items = [0, 1, 2];
199 for key in get_items.iter() {
200 map.get(key);
201 }
202
203 // convert the map into a proof
204 let proof = map.into_proof();
205
206 // check that the proof contains the expected values
207 for (key, value) in ITEMS.iter() {
208 match get_items.contains(key) {
209 true => assert_eq!(proof.get(key), Some(value)),
210 false => assert_eq!(proof.get(key), None),
211 }
212 }
213 }
214
215 #[test]
216 fn test_contains_key() {
217 // instantiate a recording map
218 let map = RecordingMap::new(ITEMS.to_vec());
219
220 // check if the map contains a few items
221 let get_items = [0, 1, 2];
222 for key in get_items.iter() {
223 map.contains_key(key);
224 }
225
226 // convert the map into a proof
227 let proof = map.into_proof();
228
229 // check that the proof contains the expected values
230 for (key, _) in ITEMS.iter() {
231 match get_items.contains(key) {
232 true => assert_eq!(proof.contains_key(key), true),
233 false => assert_eq!(proof.contains_key(key), false),
234 }
235 }
236 }
237
238 #[test]
239 fn test_len() {
240 // instantiate a recording map
241 let mut map = RecordingMap::new(ITEMS.to_vec());
242 // length of the map should be equal to the number of items
243 assert_eq!(map.len(), ITEMS.len());
244
245 // inserting entry with key that already exists should not change the length, but it does
246 // add entries to the trace and update sets
247 map.insert(4, 5);
248 assert_eq!(map.len(), ITEMS.len());
249 assert_eq!(map.trace_len(), 1);
250 assert_eq!(map.updates_len(), 1);
251
252 // inserting entry with new key should increase the length; it should also record the key
253 // as an updated key, but the trace length does not change since old values were not touched
254 map.insert(5, 5);
255 assert_eq!(map.len(), ITEMS.len() + 1);
256 assert_eq!(map.trace_len(), 1);
257 assert_eq!(map.updates_len(), 2);
258
259 // get some items so that they are saved in the trace; this should record original items
260 // in the trace, but should not affect the set of updates
261 let get_items = [0, 1, 2];
262 for key in get_items.iter() {
263 map.contains_key(key);
264 }
265 assert_eq!(map.trace_len(), 4);
266 assert_eq!(map.updates_len(), 2);
267
268 // read the same items again, this should not have any effect on either length, trace, or
269 // the set of updates
270 let get_items = [0, 1, 2];
271 for key in get_items.iter() {
272 map.contains_key(key);
273 }
274 assert_eq!(map.trace_len(), 4);
275 assert_eq!(map.updates_len(), 2);
276
277 // read a newly inserted item; this should not affect either length, trace, or the set of
278 // updates
279 let _val = map.get(&5).unwrap();
280 assert_eq!(map.trace_len(), 4);
281 assert_eq!(map.updates_len(), 2);
282
283 // update a newly inserted item; this should not affect either length, trace, or the set
284 // of updates
285 map.insert(5, 11);
286 assert_eq!(map.trace_len(), 4);
287 assert_eq!(map.updates_len(), 2);
288
289 // Note: The length reported by the proof will be different to the length originally
290 // reported by the map.
291 let proof = map.into_proof();
292
293 // length of the proof should be equal to get_items + 1. The extra item is the original
294 // value at key = 4u64
295 assert_eq!(proof.len(), get_items.len() + 1);
296 }
297
298 #[test]
299 fn test_iter() {
300 let mut map = RecordingMap::new(ITEMS.to_vec());
301 assert!(map.iter().all(|(x, y)| ITEMS.contains(&(*x, *y))));
302
303 // when inserting entry with key that already exists the iterator should return the new value
304 let new_value = 5;
305 map.insert(4, new_value);
306 assert_eq!(map.iter().count(), ITEMS.len());
307 assert!(map.iter().all(|(x, y)| if x == &4 {
308 y == &new_value
309 } else {
310 ITEMS.contains(&(*x, *y))
311 }));
312 }
313
314 #[test]
315 fn test_is_empty() {
316 // instantiate an empty recording map
317 let empty_map: RecordingMap<u64, u64> = RecordingMap::default();
318 assert!(empty_map.is_empty());
319
320 // instantiate a non-empty recording map
321 let map = RecordingMap::new(ITEMS.to_vec());
322 assert!(!map.is_empty());
323 }
324}