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