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