mithril_common/crypto_helper/
merkle_map.rs

1//! Merkelized map and associated proof
2
3use anyhow::{Context, anyhow};
4use serde::{Deserialize, Serialize};
5use std::{
6    collections::{BTreeMap, BTreeSet, HashMap},
7    hash::Hash,
8    sync::Arc,
9};
10
11use crate::{StdError, StdResult};
12
13use super::{MKProof, MKTree, MKTreeNode, MKTreeStorer};
14
15/// The trait implemented by the keys of a MKMap
16pub trait MKMapKey: PartialEq + Eq + PartialOrd + Ord + Clone + Hash + Into<MKTreeNode> {}
17
18/// The trait implemented by the values of a MKMap
19pub trait MKMapValue<K: MKMapKey>: Clone + TryInto<MKTreeNode> + TryFrom<MKTreeNode> {
20    /// Get the root of the merkelized map value
21    fn compute_root(&self) -> StdResult<MKTreeNode>;
22
23    /// Check if the merkelized map value contains a leaf
24    fn contains<T: Into<MKTreeNode> + Clone>(&self, leaf: &T) -> bool;
25
26    /// Can the merkelized map value compute a proof
27    fn can_compute_proof(&self) -> bool;
28
29    /// Compute the proof for a set of values of the merkelized map
30    fn compute_proof<T: Into<MKTreeNode> + Clone>(
31        &self,
32        leaves: &[T],
33    ) -> StdResult<Option<MKMapProof<K>>>;
34}
35
36/// A map, where the keys and values are merkelized and provable
37pub struct MKMap<K: MKMapKey, V: MKMapValue<K>, S: MKTreeStorer> {
38    inner_map_values: BTreeMap<K, V>,
39    inner_merkle_tree: MKTree<S>,
40    provable_keys: BTreeSet<K>,
41}
42
43impl<K: MKMapKey, V: MKMapValue<K>, S: MKTreeStorer> MKMap<K, V, S> {
44    /// MKMap factory
45    pub fn new(entries: &[(K, V)]) -> StdResult<Self> {
46        Self::new_from_iter(entries.to_vec())
47    }
48
49    /// MKMap factory
50    pub fn new_from_iter<T: IntoIterator<Item = (K, V)>>(entries: T) -> StdResult<Self> {
51        let inner_map_values = BTreeMap::default();
52        let inner_merkle_tree = MKTree::<S>::new::<MKTreeNode>(&[])?;
53        let can_compute_proof_keys = BTreeSet::default();
54        let mut mk_map = Self {
55            inner_map_values,
56            inner_merkle_tree,
57            provable_keys: can_compute_proof_keys,
58        };
59        let sorted_entries = BTreeMap::from_iter(entries);
60        for (key, value) in sorted_entries {
61            mk_map.insert_unchecked(key, value)?;
62        }
63
64        Ok(mk_map)
65    }
66
67    /// Insert a new key-value pair
68    /// Important: keys must be inserted in order to guarantee
69    /// that the same set of key/values results in the same computation for the root.
70    pub fn insert(&mut self, key: K, value: V) -> StdResult<()> {
71        if let Some(existing_value) = self.inner_map_values.get(&key) {
72            if existing_value.compute_root()? != value.compute_root()? {
73                return Err(anyhow!(
74                    "MKMap values should be replaced by entry with same root"
75                ));
76            }
77            return self.replace_unchecked(key, value);
78        } else {
79            let key_max = self.inner_map_values.keys().max();
80            if key_max > Some(&key) {
81                return Err(anyhow!("MKMap keys must be inserted in order"));
82            }
83        }
84
85        self.insert_unchecked(key, value)
86    }
87
88    /// Insert a new key-value pair without checking if the key is already present nor the order of insertion.
89    fn insert_unchecked(&mut self, key: K, value: V) -> StdResult<()> {
90        self.update_provable_keys(&key, &value)?;
91        self.inner_map_values.insert(key.clone(), value.clone());
92        let mktree_node_value = value
93            .try_into()
94            .map_err(|_| anyhow!("MKMap could not convert value to NKTreeNode"))
95            .with_context(|| "MKMap could not convert insert value")?;
96        let mktree_node_key: MKTreeNode = key.into();
97        self.inner_merkle_tree
98            .append(&[mktree_node_key + mktree_node_value])?;
99
100        Ok(())
101    }
102
103    /// Replace the value of an existing key
104    pub fn replace(&mut self, key: K, value: V) -> StdResult<()> {
105        match self.inner_map_values.get(&key) {
106            Some(existing_value) if existing_value.compute_root()? != value.compute_root()? => Err(
107                anyhow!("MKMap values should be replaced by entry with same root"),
108            ),
109            Some(_) => self.replace_unchecked(key, value),
110            None => Err(anyhow!("MKMap could not replace non-existing key")),
111        }
112    }
113
114    /// Replace the value of an existing key without checking if the key is already present
115    fn replace_unchecked(&mut self, key: K, value: V) -> StdResult<()> {
116        self.update_provable_keys(&key, &value)?;
117        self.inner_map_values.insert(key.clone(), value.clone());
118
119        Ok(())
120    }
121
122    /// Keep track of the keys that can compute a proof
123    fn update_provable_keys(&mut self, key: &K, value: &V) -> StdResult<()> {
124        if value.can_compute_proof() {
125            self.provable_keys.insert(key.clone());
126        } else if self.provable_keys.contains(key) {
127            self.provable_keys.remove(key);
128        }
129
130        Ok(())
131    }
132
133    #[cfg(test)]
134    /// Get the provable keys of the merkelized map
135    pub fn get_provable_keys(&self) -> &BTreeSet<K> {
136        &self.provable_keys
137    }
138
139    /// Check if the merkelized map contains a leaf (and returns the corresponding key and value if exists)
140    pub fn contains(&self, leaf: &MKTreeNode) -> Option<(&K, &V)> {
141        self.iter().find(|(_, v)| v.contains(leaf))
142    }
143
144    /// Get the value of the merkelized map for a given key
145    pub fn get(&self, key: &K) -> Option<&V> {
146        self.inner_map_values.get(key)
147    }
148
149    /// Get an iterator for the key and values of the merkelized map
150    pub fn iter(&self) -> impl Iterator<Item = (&K, &V)> {
151        self.inner_map_values.iter()
152    }
153
154    /// Get the length of the merkelized map
155    pub fn len(&self) -> usize {
156        self.inner_map_values.len()
157    }
158
159    /// Check if the merkelized map is empty
160    pub fn is_empty(&self) -> bool {
161        self.inner_map_values.is_empty()
162    }
163
164    /// Compress the merkelized map
165    pub fn compress(&mut self) -> StdResult<()> {
166        let keys = self.provable_keys.clone();
167        for key in keys {
168            if let Some(value) = self.get(&key) {
169                let value = value
170                    .compute_root()?
171                    .try_into()
172                    .map_err(|_| anyhow!("Merkle root could not be converted to V"))?;
173                self.replace_unchecked(key.to_owned(), value)?;
174            }
175        }
176
177        Ok(())
178    }
179
180    /// Get the root of the merkle tree of the merkelized map
181    pub fn compute_root(&self) -> StdResult<MKTreeNode> {
182        self.inner_merkle_tree.compute_root()
183    }
184
185    /// Get the proof for a set of values of the merkelized map (recursively if needed)
186    pub fn compute_proof<T: Into<MKTreeNode> + Clone>(
187        &self,
188        leaves: &[T],
189    ) -> StdResult<MKMapProof<K>> {
190        if leaves.is_empty() {
191            return Err(anyhow!("MKMap could not compute proof for empty leaves"));
192        }
193
194        let leaves_by_keys = self.group_leaves_by_keys(leaves);
195        let mut sub_proofs = BTreeMap::<K, MKMapProof<K>>::default();
196        for (key, sub_leaves) in leaves_by_keys {
197            if let Some(value) = self.get(&key) {
198                if let Some(proof) = value.compute_proof(&sub_leaves)? {
199                    sub_proofs.insert(key.to_owned(), proof);
200                }
201            }
202        }
203
204        let master_proof = self
205            .inner_merkle_tree
206            .compute_proof(
207                &sub_proofs
208                    .iter()
209                    .map(|(k, p)| k.to_owned().into() + p.compute_root().to_owned())
210                    .collect::<Vec<MKTreeNode>>(),
211            )
212            .with_context(|| "MKMap could not compute master proof")?;
213
214        Ok(MKMapProof::new(master_proof, sub_proofs))
215    }
216
217    /// Returns a map with the leaves (converted to Merkle tree nodes) grouped by keys
218    fn group_leaves_by_keys<T: Into<MKTreeNode> + Clone>(
219        &self,
220        leaves: &[T],
221    ) -> HashMap<K, Vec<MKTreeNode>> {
222        let can_compute_proof_map: HashMap<K, V> = self
223            .provable_keys
224            .iter()
225            .filter_map(|k| self.get(k).map(|v| (k.to_owned(), v.to_owned())))
226            .collect();
227        let leaves_by_keys: HashMap<K, Vec<MKTreeNode>> = can_compute_proof_map
228            .iter()
229            .map(|(key, value)| {
230                let leaves_found = leaves
231                    .iter()
232                    .filter_map(|leaf| value.contains(leaf).then_some(leaf.to_owned().into()))
233                    .collect::<Vec<_>>();
234
235                (key.to_owned(), leaves_found)
236            })
237            .fold(HashMap::default(), |mut acc, (key, leaves)| {
238                leaves.into_iter().for_each(|leaf| {
239                    acc.entry(key.to_owned()).or_default().push(leaf);
240                });
241
242                acc
243            });
244
245        leaves_by_keys
246    }
247}
248
249impl<K: MKMapKey, V: MKMapValue<K>, S: MKTreeStorer> Clone for MKMap<K, V, S> {
250    fn clone(&self) -> Self {
251        // Cloning should never fail so unwrap is safe
252        let mut clone = Self::new(&[]).unwrap();
253        for (k, v) in self.inner_map_values.iter() {
254            clone.insert(k.to_owned(), v.to_owned()).unwrap();
255        }
256
257        clone
258    }
259}
260
261impl<'a, K: MKMapKey, V: MKMapValue<K>, S: MKTreeStorer> From<&'a MKMap<K, V, S>>
262    for &'a MKTree<S>
263{
264    fn from(other: &'a MKMap<K, V, S>) -> Self {
265        &other.inner_merkle_tree
266    }
267}
268
269impl<K: MKMapKey, V: MKMapValue<K>, S: MKTreeStorer> TryFrom<MKMap<K, V, S>> for MKTreeNode {
270    type Error = StdError;
271    fn try_from(other: MKMap<K, V, S>) -> Result<Self, Self::Error> {
272        other.compute_root()
273    }
274}
275
276/// A MKMapProof that proves membership of an entry in the merkelized map
277#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
278pub struct MKMapProof<K: MKMapKey> {
279    master_proof: MKProof,
280    sub_proofs: Vec<(K, MKMapProof<K>)>,
281}
282
283impl<K: MKMapKey> MKMapProof<K> {
284    /// MKMapProof factory
285    pub fn new(master_proof: MKProof, sub_proofs: BTreeMap<K, MKMapProof<K>>) -> Self {
286        let sub_proofs = sub_proofs.into_iter().collect();
287        Self {
288            master_proof,
289            sub_proofs,
290        }
291    }
292
293    /// Get the root of the merkelized map proof
294    pub fn compute_root(&self) -> MKTreeNode {
295        self.master_proof.root().to_owned()
296    }
297
298    /// Verify the merkelized map proof
299    pub fn verify(&self) -> StdResult<()> {
300        for (_key, proof) in &self.sub_proofs {
301            proof
302                .verify()
303                .with_context(|| "MKMapProof could not verify sub proof")?;
304        }
305
306        self.master_proof
307            .verify()
308            .with_context(|| "MKMapProof could not verify master proof")?;
309        if !self.sub_proofs.is_empty() {
310            self.master_proof
311                .contains(
312                    &self
313                        .sub_proofs
314                        .iter()
315                        .map(|(k, p)| k.to_owned().into() + p.compute_root().to_owned())
316                        .collect::<Vec<_>>(),
317                )
318                .with_context(|| "MKMapProof could not match verified leaves of master proof")?;
319        }
320
321        Ok(())
322    }
323
324    /// Check if the merkelized map proof contains a leaf
325    pub fn contains(&self, leaf: &MKTreeNode) -> StdResult<()> {
326        let contains_leaf = {
327            self.master_proof.contains(&[leaf.to_owned()]).is_ok()
328                || self.sub_proofs.iter().any(|(_k, p)| p.contains(leaf).is_ok())
329        };
330
331        contains_leaf
332            .then_some(())
333            .ok_or(anyhow!("MKMapProof does not contain leaf {:?}", leaf))
334    }
335
336    /// List the leaves of the merkelized map proof
337    pub fn leaves(&self) -> Vec<MKTreeNode> {
338        if self.sub_proofs.is_empty() {
339            self.master_proof.leaves()
340        } else {
341            let mut leaves = vec![];
342            self.sub_proofs.iter().for_each(|(_k, p)| {
343                leaves.extend(p.leaves());
344            });
345
346            leaves
347        }
348    }
349}
350
351impl<K: MKMapKey + Serialize + for<'de> Deserialize<'de>> MKMapProof<K> {
352    /// Convert the proof to bytes
353    pub fn to_bytes(&self) -> StdResult<Vec<u8>> {
354        bincode::serde::encode_to_vec(self, bincode::config::standard()).map_err(|e| e.into())
355    }
356
357    /// Convert the proof from bytes
358    pub fn from_bytes(bytes: &[u8]) -> StdResult<Self> {
359        let (res, _) =
360            bincode::serde::decode_from_slice::<Self, _>(bytes, bincode::config::standard())?;
361
362        Ok(res)
363    }
364}
365
366impl<K: MKMapKey> From<MKProof> for MKMapProof<K> {
367    fn from(other: MKProof) -> Self {
368        MKMapProof::new(other, BTreeMap::default())
369    }
370}
371
372/// A merkelized map node that is used to represent multi layered merkelized map
373/// The MKMapNode can be either a MKMap (Merkle map), a MKTree (full Merkle tree) or a MKTreeNode (Merkle tree node, e.g the root of a Merkle tree)
374/// Both MKMap and MKTree can generate proofs of membership for elements that they contain, which allows for recursive proof generation for the multiple layers
375#[derive(Clone)]
376pub enum MKMapNode<K: MKMapKey, S: MKTreeStorer> {
377    /// A Merkle map
378    Map(Arc<MKMap<K, Self, S>>),
379
380    /// A full Merkle tree
381    Tree(Arc<MKTree<S>>),
382
383    /// A Merkle tree node
384    TreeNode(MKTreeNode),
385}
386
387impl<K: MKMapKey, S: MKTreeStorer> MKMapValue<K> for MKMapNode<K, S> {
388    fn compute_root(&self) -> StdResult<MKTreeNode> {
389        match self {
390            MKMapNode::Map(mk_map) => mk_map.compute_root(),
391            MKMapNode::Tree(merkle_tree) => merkle_tree.compute_root(),
392            MKMapNode::TreeNode(merkle_tree_node) => Ok(merkle_tree_node.to_owned()),
393        }
394    }
395
396    fn contains<T: Into<MKTreeNode> + Clone>(&self, leaf: &T) -> bool {
397        let leaf = leaf.to_owned().into();
398        match self {
399            MKMapNode::Map(mk_map) => mk_map.contains(&leaf).is_some(),
400            MKMapNode::Tree(merkle_tree) => merkle_tree.contains(&leaf),
401            MKMapNode::TreeNode(merkle_tree_node) => *merkle_tree_node == leaf,
402        }
403    }
404
405    fn can_compute_proof(&self) -> bool {
406        match self {
407            MKMapNode::Map(_) => true,
408            MKMapNode::Tree(_) => true,
409            MKMapNode::TreeNode(_) => false,
410        }
411    }
412
413    fn compute_proof<T: Into<MKTreeNode> + Clone>(
414        &self,
415        leaves: &[T],
416    ) -> StdResult<Option<MKMapProof<K>>> {
417        match self {
418            MKMapNode::Tree(value) => {
419                let proof = value
420                    .compute_proof(
421                        &leaves.iter().map(|leaf| leaf.to_owned().into()).collect::<Vec<_>>(),
422                    )
423                    .with_context(|| "MKMapValue could not compute sub proof for MKTree")?;
424                Ok(Some(proof.into()))
425            }
426            MKMapNode::Map(value) => {
427                let proof = value
428                    .compute_proof(
429                        &leaves.iter().map(|leaf| leaf.to_owned().into()).collect::<Vec<_>>(),
430                    )
431                    .with_context(|| "MKMapValue could not compute sub proof for MKMap")?;
432                Ok(Some(proof))
433            }
434            _ => Ok(None),
435        }
436    }
437}
438
439impl<K: MKMapKey, S: MKTreeStorer> From<MKMap<K, MKMapNode<K, S>, S>> for MKMapNode<K, S> {
440    fn from(other: MKMap<K, MKMapNode<K, S>, S>) -> Self {
441        MKMapNode::Map(Arc::new(other))
442    }
443}
444
445impl<K: MKMapKey, S: MKTreeStorer> From<MKTree<S>> for MKMapNode<K, S> {
446    fn from(other: MKTree<S>) -> Self {
447        MKMapNode::Tree(Arc::new(other))
448    }
449}
450
451impl<K: MKMapKey, S: MKTreeStorer> From<MKTreeNode> for MKMapNode<K, S> {
452    fn from(other: MKTreeNode) -> Self {
453        MKMapNode::TreeNode(other)
454    }
455}
456
457impl<K: MKMapKey, S: MKTreeStorer> TryFrom<MKMapNode<K, S>> for MKTreeNode {
458    type Error = StdError;
459    fn try_from(other: MKMapNode<K, S>) -> Result<Self, Self::Error> {
460        other.compute_root()
461    }
462}
463
464#[cfg(test)]
465mod tests {
466    use std::collections::BTreeSet;
467
468    use crate::{
469        crypto_helper::MKTreeStoreInMemory,
470        entities::{BlockNumber, BlockRange},
471        test::entities_extensions::BlockRangeTestExtension,
472    };
473
474    use super::*;
475
476    fn generate_merkle_trees(
477        total_leaves: u64,
478        block_range_length: u64,
479    ) -> Vec<(BlockRange, MKTree<MKTreeStoreInMemory>)> {
480        (0..total_leaves / block_range_length)
481            .map(|block_range_index| {
482                let block_range = BlockRange::from_block_number_and_length(
483                    BlockNumber(block_range_index),
484                    BlockNumber(block_range_length),
485                )
486                .unwrap();
487                let merkle_tree_block_range = generate_merkle_tree(&block_range);
488                (block_range, merkle_tree_block_range)
489            })
490            .collect::<Vec<_>>()
491    }
492
493    fn generate_merkle_tree(block_range: &BlockRange) -> MKTree<MKTreeStoreInMemory> {
494        let leaves = (*block_range.start..*block_range.end)
495            .map(|leaf_index| leaf_index.to_string())
496            .collect::<Vec<_>>();
497        MKTree::new(&leaves).unwrap()
498    }
499
500    fn generate_merkle_trees_for_ranges(
501        block_ranges: &[BlockRange],
502    ) -> Vec<(BlockRange, MKTree<MKTreeStoreInMemory>)> {
503        block_ranges
504            .iter()
505            .map(|block_range| (block_range.to_owned(), generate_merkle_tree(block_range)))
506            .collect()
507    }
508
509    fn into_mkmap_tree_entries(
510        entries: Vec<(BlockRange, MKTree<MKTreeStoreInMemory>)>,
511    ) -> Vec<(BlockRange, MKMapNode<BlockRange, MKTreeStoreInMemory>)> {
512        entries
513            .into_iter()
514            .map(|(range, mktree)| (range, MKMapNode::Tree(Arc::new(mktree))))
515            .collect()
516    }
517
518    fn into_mkmap_tree_node_entries(
519        entries: Vec<(BlockRange, MKTree<MKTreeStoreInMemory>)>,
520    ) -> Vec<(BlockRange, MKMapNode<BlockRange, MKTreeStoreInMemory>)> {
521        entries
522            .into_iter()
523            .map(|(range, mktree)| (range, MKMapNode::TreeNode(mktree.try_into().unwrap())))
524            .collect()
525    }
526
527    #[test]
528    fn test_mk_map_should_compute_same_root_when_replacing_entry_with_equivalent() {
529        let entries = generate_merkle_trees(10, 3);
530        let mk_map_nodes =
531            MKMap::<_, _, MKTreeStoreInMemory>::new(&into_mkmap_tree_node_entries(entries.clone()))
532                .unwrap();
533        let mk_map_full =
534            MKMap::<_, _, MKTreeStoreInMemory>::new(&into_mkmap_tree_entries(entries)).unwrap();
535
536        let mk_map_nodes_root = mk_map_nodes.compute_root().unwrap();
537        let mk_map_full_root = mk_map_full.compute_root().unwrap();
538
539        assert_eq!(mk_map_full_root, mk_map_nodes_root);
540    }
541
542    #[test]
543    fn test_mk_map_should_accept_replacement_with_same_root_value() {
544        let entries = generate_merkle_trees_for_ranges(&[
545            BlockRange::new(0, 3),
546            BlockRange::new(4, 6),
547            BlockRange::new(7, 9),
548        ]);
549        let mut mk_map =
550            MKMap::<_, _, MKTreeStoreInMemory>::new(&into_mkmap_tree_entries(entries)).unwrap();
551        let mk_map_root_expected = mk_map.compute_root().unwrap();
552        let block_range_replacement = BlockRange::new(0, 3);
553        let same_root_value = MKMapNode::TreeNode(
554            mk_map.get(&block_range_replacement).unwrap().compute_root().unwrap(),
555        );
556
557        mk_map.insert(block_range_replacement, same_root_value).unwrap();
558
559        assert_eq!(mk_map_root_expected, mk_map.compute_root().unwrap())
560    }
561
562    #[test]
563    fn test_mk_map_should_reject_replacement_with_different_root_value() {
564        let entries = generate_merkle_trees_for_ranges(&[
565            BlockRange::new(0, 3),
566            BlockRange::new(4, 6),
567            BlockRange::new(7, 9),
568        ]);
569        let mut mk_map =
570            MKMap::<_, _, MKTreeStoreInMemory>::new(&into_mkmap_tree_entries(entries)).unwrap();
571        let block_range_replacement = BlockRange::new(0, 3);
572        let value_replacement: MKTreeNode = "test-123".to_string().into();
573        let different_root_value = MKMapNode::TreeNode(value_replacement);
574
575        mk_map
576            .insert(block_range_replacement, different_root_value)
577            .expect_err("the MKMap should reject replacement with different root value");
578    }
579
580    #[test]
581    fn test_mk_map_replace_should_accept_replacement_with_same_root_value() {
582        let entries = generate_merkle_trees_for_ranges(&[
583            BlockRange::new(0, 3),
584            BlockRange::new(4, 6),
585            BlockRange::new(7, 9),
586        ]);
587        let mut mk_map =
588            MKMap::<_, _, MKTreeStoreInMemory>::new(&into_mkmap_tree_entries(entries)).unwrap();
589        let block_range_replacement = BlockRange::new(0, 3);
590        let same_root_value = MKMapNode::TreeNode(
591            mk_map.get(&block_range_replacement).unwrap().compute_root().unwrap(),
592        );
593        let mk_map_root_expected = mk_map.compute_root().unwrap();
594
595        assert!(matches!(
596            mk_map.get(&block_range_replacement).unwrap(),
597            MKMapNode::Tree(..)
598        ));
599
600        mk_map
601            .replace(block_range_replacement.clone(), same_root_value)
602            .unwrap();
603
604        assert_eq!(mk_map_root_expected, mk_map.compute_root().unwrap());
605        assert!(matches!(
606            mk_map.get(&block_range_replacement).unwrap(),
607            MKMapNode::TreeNode(..)
608        ));
609    }
610
611    #[test]
612    fn test_mk_map_replace_should_reject_replacement_if_key_doesnt_exist() {
613        let entries = generate_merkle_trees_for_ranges(&[
614            BlockRange::new(0, 3),
615            BlockRange::new(4, 6),
616            BlockRange::new(7, 9),
617        ]);
618        let mut mk_map =
619            MKMap::<_, _, MKTreeStoreInMemory>::new(&into_mkmap_tree_entries(entries)).unwrap();
620
621        let error = mk_map
622            .replace(
623                BlockRange::new(10, 12),
624                MKMapNode::TreeNode("whatever".into()),
625            )
626            .expect_err("the MKMap should reject replacement for nonexisting key");
627
628        assert!(
629            error.to_string().contains("MKMap could not replace non-existing key"),
630            "Invalid error message: `{error}`",
631        );
632    }
633
634    #[test]
635    fn test_mk_map_replace_should_reject_replacement_with_different_root_value() {
636        let entries = generate_merkle_trees_for_ranges(&[
637            BlockRange::new(0, 3),
638            BlockRange::new(4, 6),
639            BlockRange::new(7, 9),
640        ]);
641        let mut mk_map =
642            MKMap::<_, _, MKTreeStoreInMemory>::new(&into_mkmap_tree_entries(entries)).unwrap();
643
644        let error = mk_map
645            .replace(
646                BlockRange::new(0, 3),
647                MKMapNode::TreeNode("different_value".into()),
648            )
649            .expect_err("the MKMap should reject replacement with different root value");
650
651        assert!(
652            error
653                .to_string()
654                .contains("MKMap values should be replaced by entry with same root"),
655            "Invalid error message: `{error}`",
656        );
657    }
658
659    #[test]
660    fn test_mk_map_should_compress_correctly() {
661        let entries = generate_merkle_trees_for_ranges(&[
662            BlockRange::new(0, 3),
663            BlockRange::new(4, 6),
664            BlockRange::new(7, 9),
665        ]);
666        let mk_map =
667            MKMap::<_, _, MKTreeStoreInMemory>::new(&into_mkmap_tree_entries(entries)).unwrap();
668        let mk_map_root_expected = mk_map.compute_root().unwrap();
669        let mk_map_provable_keys = mk_map.get_provable_keys();
670        assert!(!mk_map_provable_keys.is_empty());
671
672        let mut mk_map_compressed = mk_map.clone();
673        mk_map_compressed.compress().unwrap();
674
675        let mk_map_compressed_root = mk_map_compressed.compute_root().unwrap();
676        let mk_map_compressed_provable_keys = mk_map_compressed.get_provable_keys();
677        assert_eq!(mk_map_root_expected, mk_map_compressed_root);
678        assert!(mk_map_compressed_provable_keys.is_empty());
679    }
680
681    #[test]
682    fn test_mk_map_should_reject_out_of_order_insertion() {
683        let entries = generate_merkle_trees_for_ranges(&[
684            BlockRange::new(0, 3),
685            BlockRange::new(4, 6),
686            BlockRange::new(7, 9),
687        ]);
688        let mut mk_map =
689            MKMap::<_, _, MKTreeStoreInMemory>::new(&into_mkmap_tree_node_entries(entries))
690                .unwrap();
691        let out_of_order_entry = (
692            BlockRange::new(0, 25),
693            MKMapNode::TreeNode("test-123".into()),
694        );
695
696        mk_map
697            .insert(out_of_order_entry.0, out_of_order_entry.1)
698            .expect_err("the MKMap should reject out of order insertion");
699    }
700
701    #[test]
702    fn test_mk_map_should_list_keys_correctly() {
703        let entries = generate_merkle_trees_for_ranges(&[
704            BlockRange::new(0, 3),
705            BlockRange::new(4, 6),
706            BlockRange::new(7, 9),
707        ]);
708        let merkle_tree_entries = &into_mkmap_tree_node_entries(entries);
709        let mk_map =
710            MKMap::<_, _, MKTreeStoreInMemory>::new(merkle_tree_entries.as_slice()).unwrap();
711
712        let keys = mk_map.iter().map(|(k, _v)| k.to_owned()).collect::<Vec<_>>();
713        let expected_keys = merkle_tree_entries
714            .iter()
715            .map(|(k, _)| k)
716            .cloned()
717            .collect::<Vec<_>>();
718
719        assert_eq!(expected_keys, keys);
720    }
721
722    #[test]
723    fn test_mk_map_should_list_values_correctly() {
724        let entries = generate_merkle_trees_for_ranges(&[
725            BlockRange::new(0, 3),
726            BlockRange::new(4, 6),
727            BlockRange::new(7, 9),
728        ]);
729        let merkle_tree_entries = &into_mkmap_tree_node_entries(entries);
730        let mk_map =
731            MKMap::<_, _, MKTreeStoreInMemory>::new(merkle_tree_entries.as_slice()).unwrap();
732
733        let values = mk_map.iter().map(|(_k, v)| v.to_owned()).collect::<Vec<_>>();
734        let expected_values = merkle_tree_entries
735            .iter()
736            .map(|(_, v)| v)
737            .cloned()
738            .collect::<Vec<_>>();
739
740        assert_eq!(
741            BTreeSet::from_iter(expected_values.iter().map(|v| v.compute_root().unwrap())),
742            BTreeSet::from_iter(values.iter().map(|v| v.compute_root().unwrap()))
743        );
744    }
745
746    #[test]
747    fn test_mk_map_should_find_value_correctly() {
748        let entries = generate_merkle_trees_for_ranges(&[
749            BlockRange::new(0, 3),
750            BlockRange::new(4, 6),
751            BlockRange::new(7, 9),
752        ]);
753        let mktree_node_to_certify = entries[2].1.leaves()[1].clone();
754        let mk_map_full =
755            MKMap::<_, _, MKTreeStoreInMemory>::new(&into_mkmap_tree_entries(entries)).unwrap();
756
757        mk_map_full.contains(&mktree_node_to_certify).unwrap();
758    }
759
760    #[test]
761    fn test_mk_map_should_clone_and_compute_same_root() {
762        let entries = generate_merkle_trees_for_ranges(&[
763            BlockRange::new(0, 3),
764            BlockRange::new(4, 6),
765            BlockRange::new(7, 9),
766        ]);
767        let mk_map =
768            MKMap::<_, _, MKTreeStoreInMemory>::new(&into_mkmap_tree_entries(entries)).unwrap();
769
770        let mk_map_clone = mk_map.clone();
771
772        assert_eq!(
773            mk_map.compute_root().unwrap(),
774            mk_map_clone.compute_root().unwrap(),
775        );
776    }
777
778    #[test]
779    fn test_mk_map_should_not_compute_proof_for_no_leaves() {
780        let entries = generate_merkle_trees(10, 3);
781        let mktree_nodes_to_certify: &[MKTreeNode] = &[];
782        let mk_map_full =
783            MKMap::<_, _, MKTreeStoreInMemory>::new(&into_mkmap_tree_entries(entries)).unwrap();
784
785        mk_map_full
786            .compute_proof(mktree_nodes_to_certify)
787            .expect_err("MKMap should not compute proof for no leaves");
788    }
789
790    #[test]
791    fn test_mk_map_should_compute_and_verify_valid_proof() {
792        let entries = generate_merkle_trees(10, 3);
793        let mktree_nodes_to_certify = [
794            entries[0].1.leaves()[0].clone(),
795            entries[1].1.leaves()[0].clone(),
796            entries[1].1.leaves()[1].clone(),
797            entries[2].1.leaves()[1].clone(),
798        ];
799        let mk_map_full =
800            MKMap::<_, _, MKTreeStoreInMemory>::new(&into_mkmap_tree_entries(entries)).unwrap();
801        let mk_map_proof = mk_map_full.compute_proof(&mktree_nodes_to_certify).unwrap();
802
803        mk_map_proof.verify().unwrap();
804
805        let map_proof_root = mk_map_proof.compute_root();
806        let map_proof_root_expected = mk_map_full.compute_root().unwrap();
807        assert_eq!(map_proof_root, map_proof_root_expected);
808
809        let mk_proof_leaves = mk_map_proof.leaves();
810        assert_eq!(mktree_nodes_to_certify.to_vec(), mk_proof_leaves);
811    }
812
813    #[test]
814    fn test_mk_map_should_serialize_deserialize_proof() {
815        let entries = generate_merkle_trees(10, 3);
816        let mktree_nodes_to_certify = [
817            entries[0].1.leaves()[0].clone(),
818            entries[1].1.leaves()[0].clone(),
819            entries[1].1.leaves()[1].clone(),
820            entries[2].1.leaves()[1].clone(),
821        ];
822        let mk_map_full =
823            MKMap::<_, _, MKTreeStoreInMemory>::new(&into_mkmap_tree_entries(entries)).unwrap();
824        let mk_map_proof = mk_map_full.compute_proof(&mktree_nodes_to_certify).unwrap();
825
826        let serialized_mk_map_proof =
827            mk_map_proof.to_bytes().expect("Serialization should not fail");
828        let deserialized_mk_map_proof =
829            MKMapProof::<BlockRange>::from_bytes(&serialized_mk_map_proof)
830                .expect("Deserialization should not fail");
831        assert_eq!(
832            mk_map_proof, deserialized_mk_map_proof,
833            "Deserialized proof should match the original"
834        );
835    }
836
837    #[test]
838    fn test_mk_map_should_compute_and_verify_valid_proof_recursively() {
839        let entries = generate_merkle_trees(100, 3);
840        let mktree_nodes_to_certify = [
841            entries[0].1.leaves()[0].clone(),
842            entries[2].1.leaves()[1].clone(),
843            entries[3].1.leaves()[2].clone(),
844            entries[20].1.leaves()[0].clone(),
845            entries[30].1.leaves()[0].clone(),
846        ];
847        let merkle_tree_node_entries = &into_mkmap_tree_entries(entries)
848            .chunks(10)
849            .map(|entries| {
850                (
851                    entries.iter().fold(BlockRange::new(0, 0), |acc, (range, _)| {
852                        acc.try_add(range).unwrap()
853                    }),
854                    MKMapNode::Map(Arc::new(MKMap::new(entries).unwrap())),
855                )
856            })
857            .collect::<Vec<_>>();
858
859        let mk_map_full =
860            MKMap::<_, _, MKTreeStoreInMemory>::new(merkle_tree_node_entries.as_slice()).unwrap();
861
862        let mk_map_proof = mk_map_full.compute_proof(&mktree_nodes_to_certify).unwrap();
863
864        mk_map_proof.verify().unwrap();
865
866        let map_proof_root = mk_map_proof.compute_root();
867        let map_proof_root_expected = mk_map_full.compute_root().unwrap();
868        assert_eq!(map_proof_root, map_proof_root_expected);
869    }
870}
mithril_common/crypto_helper/merkle_map.rs

mithril_common/crypto_helper/
merkle_map.rs
