fuel_merkle/binary/
merkle_tree.rs

1use crate::{
2    binary::{
3        empty_sum,
4        in_memory::NodesTable,
5        Node,
6        Primitive,
7    },
8    common::{
9        Bytes32,
10        Position,
11        ProofSet,
12        StorageMap,
13    },
14    storage::{
15        Mappable,
16        StorageInspect,
17        StorageInspectInfallible,
18        StorageMutate,
19        StorageMutateInfallible,
20    },
21};
22
23use alloc::vec::Vec;
24use core::{
25    convert::Infallible,
26    marker::PhantomData,
27};
28
29use super::root_calculator::{
30    MerkleRootCalculator,
31    NodeStackPushError,
32};
33
34#[derive(Debug, Clone, derive_more::Display, PartialEq, Eq)]
35pub enum MerkleTreeError<StorageError> {
36    #[display(fmt = "proof index {_0} is not valid")]
37    InvalidProofIndex(u64),
38
39    #[display(fmt = "cannot load node with key {_0}; the key is not found in storage")]
40    LoadError(u64),
41
42    #[display(fmt = "{}", _0)]
43    StorageError(StorageError),
44
45    #[display(fmt = "the tree is too large")]
46    TooLarge,
47}
48
49impl<StorageError> From<StorageError> for MerkleTreeError<StorageError> {
50    fn from(err: StorageError) -> MerkleTreeError<StorageError> {
51        MerkleTreeError::StorageError(err)
52    }
53}
54
55#[derive(Debug, Clone)]
56pub struct MerkleTree<TableType, StorageType> {
57    storage: StorageType,
58    nodes: MerkleRootCalculator,
59    leaves_count: u64,
60    phantom_table: PhantomData<TableType>,
61}
62
63impl<TableType, StorageType> MerkleTree<TableType, StorageType> {
64    pub const fn empty_root() -> &'static Bytes32 {
65        empty_sum()
66    }
67
68    pub fn root(&self) -> Bytes32 {
69        let mut scratch_storage = StorageMap::<NodesTable>::new();
70        let root_node = self
71            .root_node::<Infallible>(&mut scratch_storage)
72            .expect("The type doesn't allow constructing invalid trees.");
73        match root_node {
74            None => *Self::empty_root(),
75            Some(ref node) => *node.hash(),
76        }
77    }
78
79    pub fn leaves_count(&self) -> u64 {
80        self.leaves_count
81    }
82
83    /// The root node is generated by joining all MMR peaks, where a peak is
84    /// defined as the head of a balanced subtree. A tree can be composed of a
85    /// single balanced subtree, in which case the tree is itself balanced, or
86    /// several balanced subtrees, in which case the tree is imbalanced. Only
87    /// nodes at the head of a balanced tree are persisted in storage; any node,
88    /// including the root node, whose child is an imbalanced child subtree will
89    /// not be saved in persistent storage. This is because node data for such
90    /// nodes is liable to change as more leaves are pushed to the tree.
91    /// Instead, intermediate nodes must be held in a temporary storage space.
92    ///
93    /// When calling `root_node`, callees must pass a mutable reference to a
94    /// temporary storage space that will be used to hold any intermediate nodes
95    /// that are created during root node calculation. At the end of the method
96    /// call, this temporary storage space will contain all intermediate nodes
97    /// not held in persistent storage, and these nodes will be available to the
98    /// callee.
99    ///
100    /// Returns `None` if the tree is empty, and the root node otherwise.
101    fn root_node<E>(
102        &self,
103        scratch_storage: &mut StorageMap<NodesTable>,
104    ) -> Result<Option<Node>, MerkleTreeError<E>> {
105        let mut nodes = self.nodes.stack().iter().rev();
106        let Some(mut head) = nodes.next().cloned() else {
107            return Ok(None); // Empty tree
108        };
109
110        for node in nodes {
111            let parent = node
112                .position()
113                .parent()
114                .map_err(|_| MerkleTreeError::TooLarge)?;
115            head = Node::create_node(parent, node, &head);
116            StorageMutateInfallible::insert(
117                scratch_storage,
118                &head.key(),
119                &(&head).into(),
120            );
121        }
122
123        Ok(Some(head))
124    }
125}
126
127impl<TableType, StorageType, StorageError> MerkleTree<TableType, StorageType>
128where
129    TableType: Mappable<Key = u64, Value = Primitive, OwnedValue = Primitive>,
130    StorageType: StorageInspect<TableType, Error = StorageError>,
131{
132    pub fn new(storage: StorageType) -> Self {
133        Self {
134            storage,
135            nodes: MerkleRootCalculator::new(),
136            leaves_count: 0,
137            phantom_table: Default::default(),
138        }
139    }
140
141    /// A binary Merkle tree can be built from a collection of Merkle Mountain
142    /// Range (MMR) peaks. The MMR structure can be accurately defined by the
143    /// number of leaves in the leaf row.
144    ///
145    /// Consider a binary Merkle tree with seven leaves, producing the following
146    /// MMR structure:
147    ///
148    /// ```text
149    ///       03
150    ///      /  \
151    ///     /    \
152    ///   01      05      09
153    ///  /  \    /  \    /  \
154    /// 00  02  04  06  08  10  12
155    /// ```
156    ///
157    /// We observe that the tree has three peaks at positions `03`, `09`, and
158    /// `12`. These peak positions are recorded in the order that they appear,
159    /// reading left to right in the tree structure, and only descend in height.
160    /// These peak positions communicate everything needed to determine the
161    /// remaining internal nodes building upwards to the root position:
162    ///
163    /// ```text
164    ///            07
165    ///           /  \
166    ///          /    \
167    ///         /      \
168    ///        /        \
169    ///       /          \
170    ///      /            \
171    ///    03              11
172    ///   /  \            /  \
173    /// ...  ...         /    \
174    ///                09      \
175    ///               /  \      \
176    ///             ...  ...    12
177    /// ```
178    ///
179    /// No additional intermediate nodes or leaves are required to calculate
180    /// the root position.
181    ///
182    /// The positions of the MMR peaks can be deterministically calculated as a
183    /// function of `n + 1` where `n` is the number of leaves in the tree. By
184    /// appending an additional leaf node to the tree, we generate a new tree
185    /// structure with additional internal nodes (N.B.: this may also change the
186    /// root position if the tree is already balanced).
187    ///
188    /// In our example, we add an additional leaf at leaf index `7` (in-order
189    /// index `14`):
190    ///
191    /// ```text
192    ///            07
193    ///           /  \
194    ///          /    \
195    ///         /      \
196    ///        /        \
197    ///       /          \
198    ///      /            \
199    ///    03              11
200    ///   /  \            /  \
201    /// ...  ...         /    \
202    ///                09      13
203    ///               /  \    /  \
204    ///             ...  ... 12  14
205    /// ```
206    ///
207    /// We observe that the path from the root position to our new leaf position
208    /// yields a set of side positions that includes our original peak
209    /// positions (see [Path Iterator](crate::common::path_iterator::PathIter)):
210    ///
211    /// | Path position | Side position |
212    /// |---------------|---------------|
213    /// |            07 |            07 |
214    /// |            11 |            03 |
215    /// |            13 |            09 |
216    /// |            14 |            12 |
217    ///
218    /// By excluding the root position `07`, we have established the set of
219    /// side positions `03`, `09`, and `12`, matching our set of MMR peaks.
220    pub fn load(
221        storage: StorageType,
222        leaves_count: u64,
223    ) -> Result<Self, MerkleTreeError<StorageError>> {
224        let mut nodes = Vec::new();
225        let peaks = peak_positions(leaves_count).ok_or(MerkleTreeError::TooLarge)?;
226        for peak in peaks.iter() {
227            let key = peak.in_order_index();
228            let node = storage
229                .get(&key)?
230                .ok_or(MerkleTreeError::LoadError(key))?
231                .into_owned()
232                .into();
233            nodes.push(node);
234        }
235
236        Ok(Self {
237            storage,
238            nodes: MerkleRootCalculator::new_with_stack(nodes),
239            leaves_count,
240            phantom_table: Default::default(),
241        })
242    }
243
244    pub fn prove(
245        &self,
246        proof_index: u64,
247    ) -> Result<(Bytes32, ProofSet), MerkleTreeError<StorageError>> {
248        if proof_index >= self.leaves_count {
249            return Err(MerkleTreeError::InvalidProofIndex(proof_index))
250        }
251
252        let root_position = root_position(self.leaves_count)
253            .expect("This tree is too large, but push should have prevented this");
254        let leaf_position = Position::from_leaf_index(proof_index)
255            .expect("leaves_count is valid, and this is less than leaves_count");
256        let (_, mut side_positions): (Vec<_>, Vec<_>) = root_position
257            .path(&leaf_position, self.leaves_count)
258            .iter()
259            .unzip();
260        side_positions.reverse(); // Reorder side positions from leaf to root.
261        side_positions.pop(); // The last side position is the root; remove it.
262
263        // Allocate scratch storage to store temporary nodes when building the
264        // root.
265        let mut scratch_storage = StorageMap::<NodesTable>::new();
266        let root_node = self
267            .root_node(&mut scratch_storage)?
268            .expect("Root node must be present, as leaves_count is nonzero");
269
270        // Get side nodes. First, we check the scratch storage. If the side node
271        // is not found in scratch storage, we then check main storage. Finally,
272        // if the side node is not found in main storage, we exit with a load
273        // error.
274        let mut proof_set = ProofSet::new();
275        for side_position in side_positions {
276            let key = side_position.in_order_index();
277            let primitive = StorageInspectInfallible::get(&scratch_storage, &key)
278                .or(StorageInspect::get(&self.storage, &key)?)
279                .ok_or(MerkleTreeError::LoadError(key))?
280                .into_owned();
281            let node = Node::from(primitive);
282            proof_set.push(*node.hash());
283        }
284
285        let root = *root_node.hash();
286        Ok((root, proof_set))
287    }
288
289    pub fn reset(&mut self) {
290        self.nodes.clear();
291    }
292}
293
294impl<TableType, StorageType, StorageError> MerkleTree<TableType, StorageType>
295where
296    TableType: Mappable<Key = u64, Value = Primitive, OwnedValue = Primitive>,
297    StorageType: StorageMutate<TableType, Error = StorageError>,
298{
299    /// Adds a new leaf node to the tree.
300    /// # WARNING
301    /// This code might modify the storage, and then return an error.
302    /// TODO: fix this issue
303    pub fn push(&mut self, data: &[u8]) -> Result<(), MerkleTreeError<StorageError>> {
304        let new_node = Node::create_leaf(self.leaves_count, data)
305            .ok_or(MerkleTreeError::TooLarge)?;
306
307        // u64 cannot overflow, as memory is finite
308        #[allow(clippy::arithmetic_side_effects)]
309        {
310            self.leaves_count += 1;
311        }
312
313        self.nodes
314            .push_with_callback(new_node, |node| {
315                self.storage
316                    .insert(&node.key(), &node.into())
317                    .map_err(MerkleTreeError::StorageError)
318                    .map(|_| ())
319            })
320            .map_err(|err| match err {
321                NodeStackPushError::Callback(err) => err,
322                NodeStackPushError::TooLarge => MerkleTreeError::TooLarge,
323            })
324    }
325}
326
327/// Calculcate root position from leaf count.
328/// Returns `None` if the tree is too large.
329fn root_position(leaves_count: u64) -> Option<Position> {
330    // The root position of a tree will always have an in-order index equal
331    // to N' - 1, where N is the leaves count and N' is N rounded (or equal)
332    // to the next power of 2.
333    #[allow(clippy::arithmetic_side_effects)] // next_power_of_two() > 0
334    Some(Position::from_in_order_index(
335        leaves_count.checked_add(1)?.next_power_of_two() - 1,
336    ))
337}
338
339/// Calculcate peak positons for given leaf count.
340/// Returns `None` if the tree is too large.
341fn peak_positions(leaves_count: u64) -> Option<Vec<Position>> {
342    let leaf_position = Position::from_leaf_index(leaves_count)?;
343    let root_position = root_position(leaves_count)?;
344
345    // Checked by root_position
346    #[allow(clippy::arithmetic_side_effects)]
347    let next_leaves_count = leaves_count + 1;
348
349    let mut peaks_itr = root_position.path(&leaf_position, next_leaves_count).iter();
350    peaks_itr.next(); // Omit the root
351
352    let (_, peaks): (Vec<_>, Vec<_>) = peaks_itr.unzip();
353
354    Some(peaks)
355}
356
357#[cfg(test)]
358mod test {
359    use super::{
360        MerkleTree,
361        MerkleTreeError,
362    };
363    use crate::{
364        binary::{
365            empty_sum,
366            leaf_sum,
367            node_sum,
368            Node,
369            Primitive,
370        },
371        common::StorageMap,
372    };
373    use fuel_merkle_test_helpers::TEST_DATA;
374    use fuel_storage::{
375        Mappable,
376        StorageInspect,
377        StorageMutate,
378    };
379
380    use alloc::vec::Vec;
381
382    #[derive(Debug)]
383    struct TestTable;
384
385    impl Mappable for TestTable {
386        type Key = Self::OwnedKey;
387        type OwnedKey = u64;
388        type OwnedValue = Primitive;
389        type Value = Self::OwnedValue;
390    }
391
392    #[test]
393    fn test_push_builds_internal_tree_structure() {
394        let mut storage_map = StorageMap::<TestTable>::new();
395        let mut tree = MerkleTree::new(&mut storage_map);
396
397        let data = &TEST_DATA[0..7]; // 7 leaves
398        for datum in data.iter() {
399            let _ = tree.push(datum);
400        }
401
402        //               07
403        //              /  \
404        //             /    \
405        //            /      \
406        //           /        \
407        //          /          \
408        //         /            \
409        //       03              11
410        //      /  \            /  \
411        //     /    \          /    \
412        //   01      05      09      \
413        //  /  \    /  \    /  \      \
414        // 00  02  04  06  08  10     12
415        // 00  01  02  03  04  05     06
416
417        let leaf_0 = leaf_sum(data[0]);
418        let leaf_1 = leaf_sum(data[1]);
419        let leaf_2 = leaf_sum(data[2]);
420        let leaf_3 = leaf_sum(data[3]);
421        let leaf_4 = leaf_sum(data[4]);
422        let leaf_5 = leaf_sum(data[5]);
423        let leaf_6 = leaf_sum(data[6]);
424        let node_1 = node_sum(&leaf_0, &leaf_1);
425        let node_5 = node_sum(&leaf_2, &leaf_3);
426        let node_3 = node_sum(&node_1, &node_5);
427        let node_9 = node_sum(&leaf_4, &leaf_5);
428
429        let s_leaf_0 = storage_map.get(&0).unwrap().unwrap();
430        let s_leaf_1 = storage_map.get(&2).unwrap().unwrap();
431        let s_leaf_2 = storage_map.get(&4).unwrap().unwrap();
432        let s_leaf_3 = storage_map.get(&6).unwrap().unwrap();
433        let s_leaf_4 = storage_map.get(&8).unwrap().unwrap();
434        let s_leaf_5 = storage_map.get(&10).unwrap().unwrap();
435        let s_leaf_6 = storage_map.get(&12).unwrap().unwrap();
436        let s_node_1 = storage_map.get(&1).unwrap().unwrap();
437        let s_node_5 = storage_map.get(&5).unwrap().unwrap();
438        let s_node_9 = storage_map.get(&9).unwrap().unwrap();
439        let s_node_3 = storage_map.get(&3).unwrap().unwrap();
440
441        assert_eq!(*Node::from(s_leaf_0.into_owned()).hash(), leaf_0);
442        assert_eq!(*Node::from(s_leaf_1.into_owned()).hash(), leaf_1);
443        assert_eq!(*Node::from(s_leaf_2.into_owned()).hash(), leaf_2);
444        assert_eq!(*Node::from(s_leaf_3.into_owned()).hash(), leaf_3);
445        assert_eq!(*Node::from(s_leaf_4.into_owned()).hash(), leaf_4);
446        assert_eq!(*Node::from(s_leaf_5.into_owned()).hash(), leaf_5);
447        assert_eq!(*Node::from(s_leaf_6.into_owned()).hash(), leaf_6);
448        assert_eq!(*Node::from(s_node_1.into_owned()).hash(), node_1);
449        assert_eq!(*Node::from(s_node_5.into_owned()).hash(), node_5);
450        assert_eq!(*Node::from(s_node_9.into_owned()).hash(), node_9);
451        assert_eq!(*Node::from(s_node_3.into_owned()).hash(), node_3);
452    }
453
454    #[test]
455    fn load_returns_a_valid_tree() {
456        const LEAVES_COUNT: u64 = 2u64.pow(16) - 1;
457
458        let mut storage_map = StorageMap::<TestTable>::new();
459
460        let expected_root = {
461            let mut tree = MerkleTree::new(&mut storage_map);
462            let data = (0u64..LEAVES_COUNT)
463                .map(|i| i.to_be_bytes())
464                .collect::<Vec<_>>();
465            for datum in data.iter() {
466                let _ = tree.push(datum);
467            }
468            tree.root()
469        };
470
471        let root = {
472            let tree = MerkleTree::load(&mut storage_map, LEAVES_COUNT).unwrap();
473            tree.root()
474        };
475
476        assert_eq!(expected_root, root);
477    }
478
479    #[test]
480    fn load_returns_empty_tree_for_0_leaves() {
481        const LEAVES_COUNT: u64 = 0;
482
483        let expected_root = *MerkleTree::<(), ()>::empty_root();
484
485        let root = {
486            let mut storage_map = StorageMap::<TestTable>::new();
487            let tree = MerkleTree::load(&mut storage_map, LEAVES_COUNT).unwrap();
488            tree.root()
489        };
490
491        assert_eq!(expected_root, root);
492    }
493
494    #[test]
495    fn load_returns_a_load_error_if_the_storage_is_not_valid_for_the_leaves_count() {
496        const LEAVES_COUNT: u64 = 5;
497
498        let mut storage_map = StorageMap::<TestTable>::new();
499
500        let mut tree = MerkleTree::new(&mut storage_map);
501        let data = (0u64..LEAVES_COUNT)
502            .map(|i| i.to_be_bytes())
503            .collect::<Vec<_>>();
504        for datum in data.iter() {
505            let _ = tree.push(datum);
506        }
507
508        let err = MerkleTree::load(&mut storage_map, LEAVES_COUNT * 2)
509            .expect_err("Expected load() to return Error; got Ok");
510        assert!(matches!(err, MerkleTreeError::LoadError(_)));
511    }
512
513    #[test]
514    fn root_returns_the_empty_root_for_0_leaves() {
515        let mut storage_map = StorageMap::<TestTable>::new();
516        let tree = MerkleTree::new(&mut storage_map);
517
518        let root = tree.root();
519        assert_eq!(root, empty_sum().clone());
520    }
521
522    #[test]
523    fn root_returns_the_merkle_root_for_1_leaf() {
524        let mut storage_map = StorageMap::<TestTable>::new();
525        let mut tree = MerkleTree::new(&mut storage_map);
526
527        let data = &TEST_DATA[0..1]; // 1 leaf
528        for datum in data.iter() {
529            let _ = tree.push(datum);
530        }
531
532        let leaf_0 = leaf_sum(data[0]);
533
534        let root = tree.root();
535        assert_eq!(root, leaf_0);
536    }
537
538    #[test]
539    fn root_returns_the_merkle_root_for_7_leaves() {
540        let mut storage_map = StorageMap::<TestTable>::new();
541        let mut tree = MerkleTree::new(&mut storage_map);
542
543        let data = &TEST_DATA[0..7]; // 7 leaves
544        for datum in data.iter() {
545            let _ = tree.push(datum);
546        }
547
548        //               07
549        //              /  \
550        //             /    \
551        //            /      \
552        //           /        \
553        //          /          \
554        //         /            \
555        //       03              11
556        //      /  \            /  \
557        //     /    \          /    \
558        //   01      05      09      \
559        //  /  \    /  \    /  \      \
560        // 00  02  04  06  08  10     12
561        // 00  01  02  03  04  05     06
562
563        let leaf_0 = leaf_sum(data[0]);
564        let leaf_1 = leaf_sum(data[1]);
565        let leaf_2 = leaf_sum(data[2]);
566        let leaf_3 = leaf_sum(data[3]);
567        let leaf_4 = leaf_sum(data[4]);
568        let leaf_5 = leaf_sum(data[5]);
569        let leaf_6 = leaf_sum(data[6]);
570
571        let node_1 = node_sum(&leaf_0, &leaf_1);
572        let node_5 = node_sum(&leaf_2, &leaf_3);
573        let node_3 = node_sum(&node_1, &node_5);
574        let node_9 = node_sum(&leaf_4, &leaf_5);
575        let node_11 = node_sum(&node_9, &leaf_6);
576        let node_7 = node_sum(&node_3, &node_11);
577
578        let root = tree.root();
579        assert_eq!(root, node_7);
580    }
581
582    #[test]
583    fn prove_returns_invalid_proof_index_error_for_0_leaves() {
584        let mut storage_map = StorageMap::<TestTable>::new();
585        let tree = MerkleTree::new(&mut storage_map);
586
587        let err = tree
588            .prove(0)
589            .expect_err("Expected prove() to return Error; got Ok");
590        assert!(matches!(err, MerkleTreeError::InvalidProofIndex(0)));
591    }
592
593    #[test]
594    fn prove_returns_invalid_proof_index_error_when_index_is_greater_than_number_of_leaves(
595    ) {
596        let mut storage_map = StorageMap::<TestTable>::new();
597        let mut tree = MerkleTree::new(&mut storage_map);
598
599        let data = &TEST_DATA[0..5]; // 5 leaves
600        for datum in data.iter() {
601            let _ = tree.push(datum);
602        }
603
604        let err = tree
605            .prove(10)
606            .expect_err("Expected prove() to return Error; got Ok");
607        assert!(matches!(err, MerkleTreeError::InvalidProofIndex(10)))
608    }
609
610    #[test]
611    fn prove_returns_the_merkle_root_and_proof_set_for_1_leaf() {
612        let mut storage_map = StorageMap::<TestTable>::new();
613        let mut tree = MerkleTree::new(&mut storage_map);
614
615        let data = &TEST_DATA[0..1]; // 1 leaf
616        for datum in data.iter() {
617            let _ = tree.push(datum);
618        }
619
620        let leaf_0 = leaf_sum(data[0]);
621
622        {
623            let (root, proof_set) = tree.prove(0).unwrap();
624            assert_eq!(root, leaf_0);
625            assert!(proof_set.is_empty());
626        }
627    }
628
629    #[test]
630    fn prove_returns_the_merkle_root_and_proof_set_for_4_leaves() {
631        let mut storage_map = StorageMap::<TestTable>::new();
632        let mut tree = MerkleTree::new(&mut storage_map);
633
634        let data = &TEST_DATA[0..4]; // 4 leaves
635        for datum in data.iter() {
636            let _ = tree.push(datum);
637        }
638
639        //       03
640        //      /  \
641        //     /    \
642        //   01      05
643        //  /  \    /  \
644        // 00  02  04  06
645        // 00  01  02  03
646
647        let leaf_0 = leaf_sum(data[0]);
648        let leaf_1 = leaf_sum(data[1]);
649        let leaf_2 = leaf_sum(data[2]);
650        let leaf_3 = leaf_sum(data[3]);
651
652        let node_1 = node_sum(&leaf_0, &leaf_1);
653        let node_5 = node_sum(&leaf_2, &leaf_3);
654        let node_3 = node_sum(&node_1, &node_5);
655
656        {
657            let (root, proof_set) = tree.prove(0).unwrap();
658            assert_eq!(root, node_3);
659            assert_eq!(proof_set[0], leaf_1);
660            assert_eq!(proof_set[1], node_5);
661        }
662        {
663            let (root, proof_set) = tree.prove(1).unwrap();
664            assert_eq!(root, node_3);
665            assert_eq!(proof_set[0], leaf_0);
666            assert_eq!(proof_set[1], node_5);
667        }
668        {
669            let (root, proof_set) = tree.prove(2).unwrap();
670            assert_eq!(root, node_3);
671            assert_eq!(proof_set[0], leaf_3);
672            assert_eq!(proof_set[1], node_1);
673        }
674        {
675            let (root, proof_set) = tree.prove(3).unwrap();
676            assert_eq!(root, node_3);
677            assert_eq!(proof_set[0], leaf_2);
678            assert_eq!(proof_set[1], node_1);
679        }
680    }
681
682    #[test]
683    fn prove_returns_the_merkle_root_and_proof_set_for_5_leaves() {
684        let mut storage_map = StorageMap::<TestTable>::new();
685        let mut tree = MerkleTree::new(&mut storage_map);
686
687        let data = &TEST_DATA[0..5]; // 5 leaves
688        for datum in data.iter() {
689            let _ = tree.push(datum);
690        }
691
692        //          07
693        //          /\
694        //         /  \
695        //       03    \
696        //      /  \    \
697        //     /    \    \
698        //   01      05   \
699        //  /  \    /  \   \
700        // 00  02  04  06  08
701        // 00  01  02  03  04
702
703        let leaf_0 = leaf_sum(data[0]);
704        let leaf_1 = leaf_sum(data[1]);
705        let leaf_2 = leaf_sum(data[2]);
706        let leaf_3 = leaf_sum(data[3]);
707        let leaf_4 = leaf_sum(data[4]);
708
709        let node_1 = node_sum(&leaf_0, &leaf_1);
710        let node_5 = node_sum(&leaf_2, &leaf_3);
711        let node_3 = node_sum(&node_1, &node_5);
712        let node_7 = node_sum(&node_3, &leaf_4);
713
714        {
715            let (root, proof_set) = tree.prove(0).unwrap();
716            assert_eq!(root, node_7);
717            assert_eq!(proof_set[0], leaf_1);
718            assert_eq!(proof_set[1], node_5);
719            assert_eq!(proof_set[2], leaf_4);
720        }
721        {
722            let (root, proof_set) = tree.prove(1).unwrap();
723            assert_eq!(root, node_7);
724            assert_eq!(proof_set[0], leaf_0);
725            assert_eq!(proof_set[1], node_5);
726            assert_eq!(proof_set[2], leaf_4);
727        }
728        {
729            let (root, proof_set) = tree.prove(2).unwrap();
730            assert_eq!(root, node_7);
731            assert_eq!(proof_set[0], leaf_3);
732            assert_eq!(proof_set[1], node_1);
733            assert_eq!(proof_set[2], leaf_4);
734        }
735        {
736            let (root, proof_set) = tree.prove(3).unwrap();
737            assert_eq!(root, node_7);
738            assert_eq!(proof_set[0], leaf_2);
739            assert_eq!(proof_set[1], node_1);
740            assert_eq!(proof_set[2], leaf_4);
741        }
742        {
743            let (root, proof_set) = tree.prove(4).unwrap();
744            assert_eq!(root, node_7);
745            assert_eq!(proof_set[0], node_3);
746        }
747    }
748
749    #[test]
750    fn prove_returns_the_merkle_root_and_proof_set_for_7_leaves() {
751        let mut storage_map = StorageMap::<TestTable>::new();
752        let mut tree = MerkleTree::new(&mut storage_map);
753
754        let data = &TEST_DATA[0..7]; // 7 leaves
755        for datum in data.iter() {
756            let _ = tree.push(datum);
757        }
758
759        //               07
760        //              /  \
761        //             /    \
762        //            /      \
763        //           /        \
764        //          /          \
765        //         /            \
766        //       03              11
767        //      /  \            /  \
768        //     /    \          /    \
769        //   01      05      09      \
770        //  /  \    /  \    /  \      \
771        // 00  02  04  06  08  10     12
772        // 00  01  02  03  04  05     06
773
774        let leaf_0 = leaf_sum(data[0]);
775        let leaf_1 = leaf_sum(data[1]);
776        let leaf_2 = leaf_sum(data[2]);
777        let leaf_3 = leaf_sum(data[3]);
778        let leaf_4 = leaf_sum(data[4]);
779        let leaf_5 = leaf_sum(data[5]);
780        let leaf_6 = leaf_sum(data[6]);
781
782        let node_1 = node_sum(&leaf_0, &leaf_1);
783        let node_5 = node_sum(&leaf_2, &leaf_3);
784        let node_3 = node_sum(&node_1, &node_5);
785        let node_9 = node_sum(&leaf_4, &leaf_5);
786        let node_11 = node_sum(&node_9, &leaf_6);
787        let node_7 = node_sum(&node_3, &node_11);
788
789        {
790            let (root, proof_set) = tree.prove(0).unwrap();
791            assert_eq!(root, node_7);
792            assert_eq!(proof_set[0], leaf_1);
793            assert_eq!(proof_set[1], node_5);
794            assert_eq!(proof_set[2], node_11);
795        }
796        {
797            let (root, proof_set) = tree.prove(1).unwrap();
798            assert_eq!(root, node_7);
799            assert_eq!(proof_set[0], leaf_0);
800            assert_eq!(proof_set[1], node_5);
801            assert_eq!(proof_set[2], node_11);
802        }
803        {
804            let (root, proof_set) = tree.prove(2).unwrap();
805            assert_eq!(root, node_7);
806            assert_eq!(proof_set[0], leaf_3);
807            assert_eq!(proof_set[1], node_1);
808            assert_eq!(proof_set[2], node_11);
809        }
810        {
811            let (root, proof_set) = tree.prove(3).unwrap();
812            assert_eq!(root, node_7);
813            assert_eq!(proof_set[0], leaf_2);
814            assert_eq!(proof_set[1], node_1);
815            assert_eq!(proof_set[2], node_11);
816        }
817        {
818            let (root, proof_set) = tree.prove(4).unwrap();
819            assert_eq!(root, node_7);
820            assert_eq!(proof_set[0], leaf_5);
821            assert_eq!(proof_set[1], leaf_6);
822            assert_eq!(proof_set[2], node_3);
823        }
824        {
825            let (root, proof_set) = tree.prove(5).unwrap();
826            assert_eq!(root, node_7);
827            assert_eq!(proof_set[0], leaf_4);
828            assert_eq!(proof_set[1], leaf_6);
829            assert_eq!(proof_set[2], node_3);
830        }
831        {
832            let (root, proof_set) = tree.prove(6).unwrap();
833            assert_eq!(root, node_7);
834            assert_eq!(proof_set[0], node_9);
835            assert_eq!(proof_set[1], node_3);
836        }
837    }
838
839    #[test]
840    fn reset_reverts_tree_to_empty_state() {
841        let mut storage_map = StorageMap::<TestTable>::new();
842        let mut tree = MerkleTree::new(&mut storage_map);
843
844        let data = &TEST_DATA[0..4]; // 4 leaves
845        for datum in data.iter() {
846            let _ = tree.push(datum);
847        }
848
849        tree.reset();
850
851        let root = tree.root();
852        let expected_root = *MerkleTree::<(), ()>::empty_root();
853        assert_eq!(root, expected_root);
854
855        let data = &TEST_DATA[0..4]; // 4 leaves
856        for datum in data.iter() {
857            let _ = tree.push(datum);
858        }
859
860        let leaf_0 = leaf_sum(data[0]);
861        let leaf_1 = leaf_sum(data[1]);
862        let leaf_2 = leaf_sum(data[2]);
863        let leaf_3 = leaf_sum(data[3]);
864
865        let node_1 = node_sum(&leaf_0, &leaf_1);
866        let node_5 = node_sum(&leaf_2, &leaf_3);
867        let node_3 = node_sum(&node_1, &node_5);
868
869        let root = tree.root();
870        let expected_root = node_3;
871        assert_eq!(root, expected_root);
872    }
873
874    #[test]
875    fn load_overflows() {
876        // Given
877        let storage_map = StorageMap::<TestTable>::new();
878        const LEAVES_COUNT: u64 = u64::MAX;
879
880        // When
881        let result = MerkleTree::load(storage_map, LEAVES_COUNT).map(|_| ());
882
883        // Then
884        assert_eq!(result, Err(MerkleTreeError::TooLarge));
885    }
886
887    #[test]
888    fn push_overflows() {
889        // Given
890        let mut storage_map = StorageMap::<TestTable>::new();
891        const LEAVES_COUNT: u64 = u64::MAX / 2;
892        loop {
893            let result = MerkleTree::load(&mut storage_map, LEAVES_COUNT).map(|_| ());
894
895            if let Err(MerkleTreeError::LoadError(index)) = result {
896                storage_map.insert(&index, &Primitive::default()).unwrap();
897            } else {
898                break;
899            }
900        }
901
902        // When
903        let mut tree = MerkleTree::load(storage_map, LEAVES_COUNT)
904            .expect("Expected `load()` to succeed");
905        let _ = tree.push(&[]);
906        let result = tree.push(&[]);
907
908        // Then
909        assert_eq!(result, Err(MerkleTreeError::TooLarge));
910    }
911}