semaphore_rs_trees/lazy/
mod.rs

1use std::fs::OpenOptions;
2use std::io::Write;
3use std::iter::{once, repeat_n, successors};
4use std::ops::{Deref, DerefMut};
5use std::path::PathBuf;
6use std::str::FromStr;
7use std::sync::{Arc, Mutex};
8
9use mmap_rs::{MmapFlags, MmapMut, MmapOptions};
10use rayon::prelude::*;
11use semaphore_rs_hasher::{Hash, Hasher};
12use thiserror::Error;
13
14use crate::{Branch, InclusionProof};
15
16pub trait VersionMarker {}
17#[derive(Debug)]
18pub struct Canonical;
19impl VersionMarker for Canonical {}
20#[derive(Debug)]
21pub struct Derived;
22impl VersionMarker for Derived {}
23
24/// A storage-optimized merkle tree.
25///
26/// It has a certain linear-buffer represented
27/// prefix subtree and the rest of the tree is represented using lazy,
28/// pointer-based structures. This makes it possible to hold even large trees in
29/// memory, assuming only a relatively small subset is ever modified.
30///
31/// It exposes an immutable API, so that multiple versions can be kept in memory
32/// while reusing as much structure as possible.
33///
34/// The update method also allows the specification of a mutability hint, which
35/// can be used to vastly improve storage characteristics, but also requires the
36/// caller to ensure certain additional invariants hold. See
37/// [`LazyMerkleTree::update_with_mutation`] for details.
38pub struct LazyMerkleTree<H: Hasher, V: VersionMarker = Derived> {
39    tree: AnyTree<H>,
40    _version: V,
41}
42
43impl<H, Version> LazyMerkleTree<H, Version>
44where
45    H: Hasher,
46    <H as Hasher>::Hash: Hash,
47    Version: VersionMarker,
48{
49    /// Creates a new, fully lazy (without any dense prefix) tree.
50    #[must_use]
51    pub fn new(depth: usize, empty_value: H::Hash) -> LazyMerkleTree<H, Canonical> {
52        LazyMerkleTree {
53            tree: AnyTree::new(depth, empty_value),
54            _version: Canonical,
55        }
56    }
57
58    /// Creates a new tree with a dense prefix of the given depth.
59    #[must_use]
60    pub fn new_with_dense_prefix(
61        depth: usize,
62        prefix_depth: usize,
63        empty_value: &H::Hash,
64    ) -> LazyMerkleTree<H, Canonical> {
65        LazyMerkleTree {
66            tree: AnyTree::new_with_dense_prefix(depth, prefix_depth, empty_value),
67            _version: Canonical,
68        }
69    }
70
71    /// Creates a new tree with a dense prefix of the given depth, and with
72    /// initial leaves populated from the given slice.
73    #[must_use]
74    pub fn new_with_dense_prefix_with_initial_values(
75        depth: usize,
76        prefix_depth: usize,
77        empty_value: &H::Hash,
78        initial_values: &[H::Hash],
79    ) -> LazyMerkleTree<H, Canonical> {
80        LazyMerkleTree {
81            tree: AnyTree::new_with_dense_prefix_with_initial_values(
82                depth,
83                prefix_depth,
84                empty_value,
85                initial_values,
86            ),
87            _version: Canonical,
88        }
89    }
90
91    /// Creates a new memory mapped file specified by path and creates a tree
92    /// with dense prefix of the given depth with initial values
93    pub fn new_mmapped_with_dense_prefix_with_init_values(
94        depth: usize,
95        prefix_depth: usize,
96        empty_value: &H::Hash,
97        initial_values: &[H::Hash],
98        file_path: &str,
99    ) -> Result<LazyMerkleTree<H, Canonical>, DenseMMapError> {
100        Ok(LazyMerkleTree {
101            tree: AnyTree::new_mmapped_with_dense_prefix_with_init_values(
102                depth,
103                prefix_depth,
104                empty_value,
105                initial_values,
106                file_path,
107            )?,
108            _version: Canonical,
109        })
110    }
111
112    /// Attempts to restore previous tree state from memory mapped file
113    ///
114    /// # Errors
115    /// - dense mmap tree restore failed
116    pub fn attempt_dense_mmap_restore(
117        depth: usize,
118        prefix_depth: usize,
119        empty_leaf: &H::Hash,
120        file_path: &str,
121    ) -> Result<LazyMerkleTree<H, Canonical>, DenseMMapError> {
122        Ok(LazyMerkleTree {
123            tree: AnyTree::try_restore_dense_mmap_tree_state(
124                depth,
125                prefix_depth,
126                empty_leaf,
127                file_path,
128            )?,
129            _version: Canonical,
130        })
131    }
132
133    /// Returns the depth of the tree.
134    #[must_use]
135    pub const fn depth(&self) -> usize {
136        self.tree.depth()
137    }
138
139    /// Returns the root of the tree.
140    #[must_use]
141    pub fn root(&self) -> H::Hash {
142        self.tree.root()
143    }
144
145    /// Sets the value at the given index to the given value. This is fully
146    /// immutable, returning a new tree and leaving the old one unchanged.
147    /// Reuses as much memory as possible, allocating only `depth` nodes.
148    #[must_use]
149    pub fn update(&self, index: usize, value: &H::Hash) -> LazyMerkleTree<H, Derived> {
150        LazyMerkleTree {
151            tree: self
152                .tree
153                .update_with_mutation_condition(index, value, false),
154            _version: Derived,
155        }
156    }
157
158    /// Returns the Merkle proof for the given index.
159    #[must_use]
160    pub fn proof(&self, index: usize) -> InclusionProof<H> {
161        self.tree.proof(index)
162    }
163
164    /// Verifies the given proof for the given value.
165    #[must_use]
166    pub fn verify(&self, value: H::Hash, proof: &InclusionProof<H>) -> bool {
167        proof.root(value) == self.root()
168    }
169
170    /// Returns the value at the given index.
171    #[must_use]
172    pub fn get_leaf(&self, index: usize) -> H::Hash {
173        self.tree.get_leaf(index)
174    }
175
176    /// Returns an iterator over all leaves.
177    pub fn leaves(&self) -> impl Iterator<Item = H::Hash> + '_ {
178        // TODO this could be made faster by a custom iterator
179        (0..(1 << self.depth())).map(|i| self.get_leaf(i))
180    }
181}
182
183impl<H> LazyMerkleTree<H, Canonical>
184where
185    H: Hasher,
186    <H as Hasher>::Hash: Hash,
187{
188    /// Sets the value at the given index to the given value. This is a mutable
189    /// operation, that will modify any dense subtrees in place.
190    ///
191    /// This has potential consequences for the soundness of the whole
192    /// structure:
193    /// it has the potential to invalidate some trees that share nodes with
194    /// this one, so if many versions are kept at once, special care must be
195    /// taken when calling this. The only trees that are guaranteed to still be
196    /// valid after this operation, are those that already specify the same
197    /// value at the given index. For example, if a linear history of updates is
198    /// kept in memory, this operation is a good way to "flatten" updates into
199    /// the oldest kept version.
200    ///
201    /// This operation is useful for storage optimizations, as it avoids
202    /// allocating any new memory in dense subtrees.
203    #[must_use]
204    pub fn update_with_mutation(self, index: usize, value: &H::Hash) -> Self {
205        Self {
206            tree: self.tree.update_with_mutation_condition(index, value, true),
207            _version: Canonical,
208        }
209    }
210
211    /// Gives a `Derived` version of this tree. Useful for initializing
212    /// versioned trees.
213    #[must_use]
214    pub fn derived(&self) -> LazyMerkleTree<H, Derived> {
215        LazyMerkleTree {
216            tree: self.tree.clone(),
217            _version: Derived,
218        }
219    }
220}
221
222impl<H> Clone for LazyMerkleTree<H, Derived>
223where
224    H: Hasher,
225    <H as Hasher>::Hash: Hash,
226{
227    fn clone(&self) -> Self {
228        Self {
229            tree: self.tree.clone(),
230            _version: Derived,
231        }
232    }
233}
234
235enum AnyTree<H: Hasher> {
236    Empty(EmptyTree<H>),
237    Sparse(SparseTree<H>),
238    Dense(DenseTree<H>),
239    DenseMMap(DenseMMapTree<H>),
240}
241
242impl<H> AnyTree<H>
243where
244    H: Hasher,
245    <H as Hasher>::Hash: Hash,
246{
247    fn new(depth: usize, empty_value: H::Hash) -> Self {
248        Self::Empty(EmptyTree::new(depth, empty_value))
249    }
250
251    fn new_with_dense_prefix_with_initial_values(
252        depth: usize,
253        prefix_depth: usize,
254        empty_value: &H::Hash,
255        initial_values: &[H::Hash],
256    ) -> Self {
257        assert!(depth >= prefix_depth);
258        let dense = DenseTree::new_with_values(initial_values, empty_value, prefix_depth);
259        let mut result: Self = dense.into();
260        let mut current_depth = prefix_depth;
261        while current_depth < depth {
262            result =
263                SparseTree::new(result, EmptyTree::new(current_depth, *empty_value).into()).into();
264            current_depth += 1;
265        }
266        result
267    }
268
269    fn new_with_dense_prefix(depth: usize, prefix_depth: usize, empty_value: &H::Hash) -> Self {
270        assert!(depth >= prefix_depth);
271        let mut result: Self = EmptyTree::new(prefix_depth, *empty_value)
272            .alloc_dense()
273            .into();
274        let mut current_depth = prefix_depth;
275        while current_depth < depth {
276            result =
277                SparseTree::new(result, EmptyTree::new(current_depth, *empty_value).into()).into();
278            current_depth += 1;
279        }
280        result
281    }
282
283    fn new_mmapped_with_dense_prefix_with_init_values(
284        depth: usize,
285        prefix_depth: usize,
286        empty_value: &H::Hash,
287        initial_values: &[H::Hash],
288        file_path: &str,
289    ) -> Result<Self, DenseMMapError> {
290        assert!(depth >= prefix_depth);
291        let dense =
292            DenseMMapTree::new_with_values(initial_values, empty_value, prefix_depth, file_path)?;
293        let mut result: Self = dense.into();
294        let mut current_depth = prefix_depth;
295        while current_depth < depth {
296            result =
297                SparseTree::new(result, EmptyTree::new(current_depth, *empty_value).into()).into();
298            current_depth += 1;
299        }
300        Ok(result)
301    }
302
303    fn try_restore_dense_mmap_tree_state(
304        depth: usize,
305        prefix_depth: usize,
306        empty_leaf: &H::Hash,
307        file_path: &str,
308    ) -> Result<Self, DenseMMapError> {
309        let dense_mmap = DenseMMapTree::attempt_restore(empty_leaf, prefix_depth, file_path)?;
310
311        let mut result: Self = dense_mmap.into();
312
313        let mut current_depth = prefix_depth;
314        while current_depth < depth {
315            result =
316                SparseTree::new(result, EmptyTree::new(current_depth, *empty_leaf).into()).into();
317            current_depth += 1;
318        }
319
320        Ok(result)
321    }
322
323    const fn depth(&self) -> usize {
324        match self {
325            Self::Empty(tree) => tree.depth,
326            Self::Sparse(tree) => tree.depth,
327            Self::Dense(tree) => tree.depth,
328            Self::DenseMMap(tree) => tree.depth,
329        }
330    }
331
332    fn root(&self) -> H::Hash {
333        match self {
334            Self::Empty(tree) => tree.root(),
335            Self::Sparse(tree) => tree.root(),
336            Self::Dense(tree) => tree.root(),
337            Self::DenseMMap(tree) => tree.root(),
338        }
339    }
340
341    fn proof(&self, index: usize) -> InclusionProof<H> {
342        assert!(index < (1 << self.depth()));
343        let mut path = Vec::with_capacity(self.depth());
344        match self {
345            Self::Empty(tree) => tree.write_proof(index, &mut path),
346            Self::Sparse(tree) => tree.write_proof(index, &mut path),
347            Self::Dense(tree) => tree.write_proof(index, &mut path),
348            Self::DenseMMap(tree) => tree.write_proof(index, &mut path),
349        }
350        path.reverse();
351        InclusionProof(path)
352    }
353
354    fn write_proof(&self, index: usize, path: &mut Vec<Branch<H::Hash>>) {
355        match self {
356            Self::Empty(tree) => tree.write_proof(index, path),
357            Self::Sparse(tree) => tree.write_proof(index, path),
358            Self::Dense(tree) => tree.write_proof(index, path),
359            Self::DenseMMap(tree) => tree.write_proof(index, path),
360        }
361    }
362
363    fn update_with_mutation_condition(
364        &self,
365        index: usize,
366        value: &H::Hash,
367        is_mutation_allowed: bool,
368    ) -> Self {
369        match self {
370            Self::Empty(tree) => tree
371                .update_with_mutation_condition(index, value, is_mutation_allowed)
372                .into(),
373            Self::Sparse(tree) => tree
374                .update_with_mutation_condition(index, value, is_mutation_allowed)
375                .into(),
376            Self::Dense(tree) => {
377                tree.update_with_mutation_condition(index, value, is_mutation_allowed)
378            }
379            Self::DenseMMap(tree) => {
380                tree.update_with_mutation_condition(index, value, is_mutation_allowed)
381            }
382        }
383    }
384
385    fn get_leaf(&self, index: usize) -> H::Hash {
386        match self {
387            Self::Empty(tree) => tree.get_leaf(),
388            Self::Sparse(tree) => tree.get_leaf(index),
389            Self::Dense(tree) => tree.get_leaf(index),
390            Self::DenseMMap(tree) => tree.get_leaf(index),
391        }
392    }
393}
394
395impl<H> Clone for AnyTree<H>
396where
397    H: Hasher,
398    <H as Hasher>::Hash: Hash,
399{
400    fn clone(&self) -> Self {
401        match self {
402            Self::Empty(t) => t.clone().into(),
403            Self::Sparse(t) => t.clone().into(),
404            Self::Dense(t) => t.clone().into(),
405            Self::DenseMMap(t) => t.clone().into(),
406        }
407    }
408}
409
410impl<H: Hasher> From<EmptyTree<H>> for AnyTree<H> {
411    fn from(tree: EmptyTree<H>) -> Self {
412        Self::Empty(tree)
413    }
414}
415
416impl<H: Hasher> From<SparseTree<H>> for AnyTree<H> {
417    fn from(tree: SparseTree<H>) -> Self {
418        Self::Sparse(tree)
419    }
420}
421
422impl<H: Hasher> From<DenseTree<H>> for AnyTree<H> {
423    fn from(tree: DenseTree<H>) -> Self {
424        Self::Dense(tree)
425    }
426}
427
428impl<H: Hasher> From<DenseMMapTree<H>> for AnyTree<H> {
429    fn from(tree: DenseMMapTree<H>) -> Self {
430        Self::DenseMMap(tree)
431    }
432}
433
434struct EmptyTree<H: Hasher> {
435    depth: usize,
436    empty_tree_values: Arc<Vec<H::Hash>>,
437}
438
439impl<H: Hasher> Clone for EmptyTree<H> {
440    fn clone(&self) -> Self {
441        Self {
442            depth: self.depth,
443            empty_tree_values: self.empty_tree_values.clone(),
444        }
445    }
446}
447
448impl<H> EmptyTree<H>
449where
450    H: Hasher,
451    <H as Hasher>::Hash: Hash,
452{
453    #[must_use]
454    fn new(depth: usize, empty_value: H::Hash) -> Self {
455        let empty_tree_values = {
456            let values = successors(Some(empty_value), |value| Some(H::hash_node(value, value)))
457                .take(depth + 1)
458                .collect();
459            Arc::new(values)
460        };
461        Self {
462            depth,
463            empty_tree_values,
464        }
465    }
466
467    fn write_proof(&self, index: usize, path: &mut Vec<Branch<H::Hash>>) {
468        for depth in (1..=self.depth).rev() {
469            let val = self.empty_tree_values[depth - 1];
470            let branch = if get_turn_at_depth(index, depth) == Turn::Left {
471                Branch::Left(val)
472            } else {
473                Branch::Right(val)
474            };
475            path.push(branch);
476        }
477    }
478
479    #[must_use]
480    fn update_with_mutation_condition(
481        &self,
482        index: usize,
483        value: &H::Hash,
484        is_mutation_allowed: bool,
485    ) -> SparseTree<H> {
486        self.alloc_sparse()
487            .update_with_mutation_condition(index, value, is_mutation_allowed)
488    }
489
490    #[must_use]
491    fn alloc_sparse(&self) -> SparseTree<H> {
492        if self.depth == 0 {
493            SparseTree::new_leaf(self.root())
494        } else {
495            let next_child: Self = Self {
496                depth: self.depth - 1,
497                empty_tree_values: self.empty_tree_values.clone(),
498            };
499            SparseTree::new(next_child.clone().into(), next_child.into())
500        }
501    }
502
503    #[must_use]
504    fn alloc_dense(&self) -> DenseTree<H> {
505        let values = self
506            .empty_tree_values
507            .iter()
508            .rev()
509            .enumerate()
510            .flat_map(|(depth, value)| repeat_n(value, 1 << depth));
511        let padded_values = once(&self.empty_tree_values[0])
512            .chain(values)
513            .cloned()
514            .collect();
515        DenseTree {
516            depth: self.depth,
517            root_index: 1,
518            storage: Arc::new(Mutex::new(padded_values)),
519        }
520    }
521
522    #[must_use]
523    fn root(&self) -> H::Hash {
524        self.empty_tree_values[self.depth]
525    }
526
527    fn get_leaf(&self) -> H::Hash {
528        self.empty_tree_values[0]
529    }
530}
531
532struct Children<H: Hasher> {
533    left: Arc<AnyTree<H>>,
534    right: Arc<AnyTree<H>>,
535}
536
537impl<H: Hasher> Clone for Children<H> {
538    fn clone(&self) -> Self {
539        Self {
540            left: self.left.clone(),
541            right: self.right.clone(),
542        }
543    }
544}
545
546struct SparseTree<H: Hasher> {
547    depth: usize,
548    root: H::Hash,
549    children: Option<Children<H>>,
550}
551
552#[derive(Debug, PartialEq, Eq)]
553enum Turn {
554    Left,
555    Right,
556}
557
558const fn get_turn_at_depth(index: usize, depth: usize) -> Turn {
559    if index & (1 << (depth - 1)) == 0 {
560        Turn::Left
561    } else {
562        Turn::Right
563    }
564}
565
566const fn clear_turn_at_depth(index: usize, depth: usize) -> usize {
567    index & !(1 << (depth - 1))
568}
569
570impl<H> From<Children<H>> for SparseTree<H>
571where
572    H: Hasher,
573    <H as Hasher>::Hash: Hash,
574{
575    fn from(children: Children<H>) -> Self {
576        assert_eq!(children.left.depth(), children.right.depth());
577        let (depth, root) = {
578            let left = children.left.clone();
579            let right = children.right.clone();
580            let depth = left.depth() + 1;
581            let root = H::hash_node(&left.root(), &right.root());
582            (depth, root)
583        };
584        Self {
585            depth,
586            root,
587            children: Some(children),
588        }
589    }
590}
591
592impl<H> Clone for SparseTree<H>
593where
594    H: Hasher,
595    <H as Hasher>::Hash: Hash,
596{
597    fn clone(&self) -> Self {
598        Self {
599            depth: self.depth,
600            root: self.root,
601            children: self.children.clone(),
602        }
603    }
604}
605
606impl<H> SparseTree<H>
607where
608    H: Hasher,
609    <H as Hasher>::Hash: Hash,
610{
611    fn new(left: AnyTree<H>, right: AnyTree<H>) -> Self {
612        assert_eq!(left.depth(), right.depth());
613        let children = Children {
614            left: Arc::new(left),
615            right: Arc::new(right),
616        };
617        children.into()
618    }
619
620    const fn new_leaf(value: H::Hash) -> Self {
621        Self {
622            depth: 0,
623            root: value,
624            children: None,
625        }
626    }
627
628    fn write_proof(&self, index: usize, path: &mut Vec<Branch<H::Hash>>) {
629        if let Some(children) = &self.children {
630            let next_index = clear_turn_at_depth(index, self.depth);
631            if get_turn_at_depth(index, self.depth) == Turn::Left {
632                path.push(Branch::Left(children.right.root()));
633                children.left.write_proof(next_index, path);
634            } else {
635                path.push(Branch::Right(children.left.root()));
636                children.right.write_proof(next_index, path);
637            }
638        }
639    }
640
641    #[must_use]
642    fn update_with_mutation_condition(
643        &self,
644        index: usize,
645        value: &H::Hash,
646        is_mutation_allowed: bool,
647    ) -> Self {
648        let Some(children) = &self.children else {
649            // no children – this is a leaf
650            return Self::new_leaf(*value);
651        };
652
653        let next_index = clear_turn_at_depth(index, self.depth);
654        let children = if get_turn_at_depth(index, self.depth) == Turn::Left {
655            let left = &children.left;
656            let new_left =
657                left.update_with_mutation_condition(next_index, value, is_mutation_allowed);
658            Children {
659                left: Arc::new(new_left),
660                right: children.right.clone(),
661            }
662        } else {
663            let right = &children.right;
664            let new_right =
665                right.update_with_mutation_condition(next_index, value, is_mutation_allowed);
666            Children {
667                left: children.left.clone(),
668                right: Arc::new(new_right),
669            }
670        };
671
672        children.into()
673    }
674
675    fn root(&self) -> H::Hash {
676        self.root
677    }
678
679    fn get_leaf(&self, index: usize) -> H::Hash {
680        self.children.as_ref().map_or_else(
681            || self.root,
682            |children| {
683                let next_index = clear_turn_at_depth(index, self.depth);
684                if get_turn_at_depth(index, self.depth) == Turn::Left {
685                    children.left.get_leaf(next_index)
686                } else {
687                    children.right.get_leaf(next_index)
688                }
689            },
690        )
691    }
692}
693
694#[derive(Debug)]
695struct DenseTree<H: Hasher> {
696    depth: usize,
697    root_index: usize,
698    storage: Arc<Mutex<Vec<H::Hash>>>,
699}
700
701impl<H: Hasher> Clone for DenseTree<H> {
702    fn clone(&self) -> Self {
703        Self {
704            depth: self.depth,
705            root_index: self.root_index,
706            storage: self.storage.clone(),
707        }
708    }
709}
710
711impl<H> DenseTree<H>
712where
713    H: Hasher,
714    <H as Hasher>::Hash: Hash,
715{
716    fn vec_from_values(values: &[H::Hash], empty_value: &H::Hash, depth: usize) -> Vec<H::Hash> {
717        let leaf_count = 1 << depth;
718        let storage_size = 1 << (depth + 1);
719        let mut storage = Vec::with_capacity(storage_size);
720
721        let empties = repeat_n(empty_value, leaf_count);
722        storage.extend(empties);
723        storage.extend_from_slice(values);
724        if values.len() < leaf_count {
725            let empties = repeat_n(empty_value, leaf_count - values.len());
726            storage.extend(empties);
727        }
728
729        // We iterate over mutable layers of the tree
730        for current_depth in (1..=depth).rev() {
731            let (top, child_layer) = storage.split_at_mut(1 << current_depth);
732            let parent_layer = &mut top[(1 << (current_depth - 1))..];
733
734            parent_layer
735                .par_iter_mut()
736                .enumerate()
737                .for_each(|(i, value)| {
738                    let left = &child_layer[2 * i];
739                    let right = &child_layer[2 * i + 1];
740                    *value = H::hash_node(left, right);
741                });
742        }
743
744        storage
745    }
746
747    fn new_with_values(values: &[H::Hash], empty_value: &H::Hash, depth: usize) -> Self {
748        let storage = Self::vec_from_values(values, empty_value, depth);
749
750        Self {
751            depth,
752            root_index: 1,
753            storage: Arc::new(Mutex::new(storage)),
754        }
755    }
756
757    fn with_ref<F, R>(&self, fun: F) -> R
758    where
759        F: FnOnce(DenseTreeRef<H>) -> R,
760    {
761        let guard = self.storage.lock().expect("lock poisoned, terminating");
762        let r = DenseTreeRef {
763            depth: self.depth,
764            root_index: self.root_index,
765            storage: &guard,
766            locked_storage: &self.storage,
767        };
768        fun(r)
769    }
770
771    fn write_proof(&self, index: usize, path: &mut Vec<Branch<H::Hash>>) {
772        self.with_ref(|r| r.write_proof(index, path));
773    }
774
775    fn get_leaf(&self, index: usize) -> H::Hash {
776        self.with_ref(|r| {
777            let leaf_index_in_dense_tree = index + (self.root_index << self.depth);
778            r.storage[leaf_index_in_dense_tree]
779        })
780    }
781
782    fn update_with_mutation_condition(
783        &self,
784        index: usize,
785        value: &H::Hash,
786        is_mutation_allowed: bool,
787    ) -> AnyTree<H> {
788        if is_mutation_allowed {
789            self.update_with_mutation(index, value);
790            self.clone().into()
791        } else {
792            self.with_ref(|r| r.update(index, value)).into()
793        }
794    }
795
796    fn update_with_mutation(&self, index: usize, value: &H::Hash) {
797        let mut storage = self.storage.lock().expect("lock poisoned, terminating");
798        let leaf_index_in_dense_tree = index + (self.root_index << self.depth);
799        storage[leaf_index_in_dense_tree] = *value;
800        let mut current = leaf_index_in_dense_tree / 2;
801        while current > 0 {
802            let left = &storage[2 * current];
803            let right = &storage[2 * current + 1];
804            storage[current] = H::hash_node(left, right);
805            current /= 2;
806        }
807    }
808
809    fn root(&self) -> H::Hash {
810        self.storage.lock().unwrap()[self.root_index]
811    }
812}
813
814struct DenseTreeRef<'a, H: Hasher> {
815    depth: usize,
816    root_index: usize,
817    storage: &'a Vec<H::Hash>,
818    locked_storage: &'a Arc<Mutex<Vec<H::Hash>>>,
819}
820
821impl<H: Hasher> From<DenseTreeRef<'_, H>> for DenseTree<H> {
822    fn from(value: DenseTreeRef<H>) -> Self {
823        Self {
824            depth: value.depth,
825            root_index: value.root_index,
826            storage: value.locked_storage.clone(),
827        }
828    }
829}
830
831impl<H: Hasher> From<DenseTreeRef<'_, H>> for AnyTree<H> {
832    fn from(value: DenseTreeRef<H>) -> Self {
833        Self::Dense(value.into())
834    }
835}
836
837impl<H> DenseTreeRef<'_, H>
838where
839    H: Hasher,
840    <H as Hasher>::Hash: Hash,
841{
842    fn root(&self) -> H::Hash {
843        self.storage[self.root_index]
844    }
845
846    const fn left(&self) -> DenseTreeRef<H> {
847        Self {
848            depth: self.depth - 1,
849            root_index: 2 * self.root_index,
850            storage: self.storage,
851            locked_storage: self.locked_storage,
852        }
853    }
854
855    const fn right(&self) -> DenseTreeRef<H> {
856        Self {
857            depth: self.depth - 1,
858            root_index: 2 * self.root_index + 1,
859            storage: self.storage,
860            locked_storage: self.locked_storage,
861        }
862    }
863
864    fn write_proof(&self, index: usize, path: &mut Vec<Branch<H::Hash>>) {
865        if self.depth == 0 {
866            return;
867        }
868        let next_index = clear_turn_at_depth(index, self.depth);
869        if get_turn_at_depth(index, self.depth) == Turn::Left {
870            path.push(Branch::Left(self.right().root()));
871            self.left().write_proof(next_index, path);
872        } else {
873            path.push(Branch::Right(self.left().root()));
874            self.right().write_proof(next_index, path);
875        }
876    }
877
878    fn update(&self, index: usize, hash: &H::Hash) -> SparseTree<H> {
879        if self.depth == 0 {
880            return SparseTree::new_leaf(*hash);
881        }
882        let next_index = clear_turn_at_depth(index, self.depth);
883        if get_turn_at_depth(index, self.depth) == Turn::Left {
884            let left = self.left();
885            let new_left = left.update(next_index, hash);
886            let right = self.right();
887            let new_root = H::hash_node(&new_left.root(), &right.root());
888            SparseTree {
889                children: Some(Children {
890                    left: Arc::new(new_left.into()),
891                    right: Arc::new(self.right().into()),
892                }),
893                root: new_root,
894                depth: self.depth,
895            }
896        } else {
897            let right = self.right();
898            let new_right = right.update(next_index, hash);
899            let left = self.left();
900            let new_root = H::hash_node(&left.root(), &new_right.root());
901            SparseTree {
902                children: Some(Children {
903                    left: Arc::new(self.left().into()),
904                    right: Arc::new(new_right.into()),
905                }),
906                root: new_root,
907                depth: self.depth,
908            }
909        }
910    }
911}
912
913struct DenseMMapTree<H: Hasher> {
914    depth: usize,
915    root_index: usize,
916    storage: Arc<Mutex<MmapMutWrapper<H>>>,
917}
918
919impl<H: Hasher> Clone for DenseMMapTree<H> {
920    fn clone(&self) -> Self {
921        Self {
922            depth: self.depth,
923            root_index: self.root_index,
924            storage: self.storage.clone(),
925        }
926    }
927}
928
929impl<H> DenseMMapTree<H>
930where
931    H: Hasher,
932    <H as Hasher>::Hash: Hash,
933{
934    /// Creates a new `DenseMMapTree` with initial values and depth
935    ///
936    /// # Errors
937    ///
938    /// - returns Err if path buf failed to be created with provided string
939    /// - returns Err if mmap creation fails
940    fn new_with_values(
941        values: &[H::Hash],
942        empty_value: &H::Hash,
943        depth: usize,
944        mmap_file_path: &str,
945    ) -> Result<Self, DenseMMapError> {
946        let path_buf = match PathBuf::from_str(mmap_file_path) {
947            Ok(pb) => pb,
948            Err(_e) => return Err(DenseMMapError::FailedToCreatePathBuf),
949        };
950
951        let storage = DenseTree::<H>::vec_from_values(values, empty_value, depth);
952
953        let mmap = MmapMutWrapper::new_from_storage(path_buf, &storage)?;
954
955        Ok(Self {
956            depth,
957            root_index: 1,
958            storage: Arc::new(Mutex::new(mmap)),
959        })
960    }
961
962    /// Given the file path and tree depth,
963    /// it attempts to restore the memory map
964    ///
965    /// # Errors
966    ///
967    /// - returns Err if path buf creation fails
968    /// - Derives errors from `MmapMutWrapper`
969    ///
970    /// # Panics
971    ///
972    /// - mutex lock is poisoned
973    fn attempt_restore(
974        empty_leaf: &H::Hash,
975        depth: usize,
976        mmap_file_path: &str,
977    ) -> Result<Self, DenseMMapError> {
978        let path_buf = match PathBuf::from_str(mmap_file_path) {
979            Ok(pb) => pb,
980            Err(_e) => return Err(DenseMMapError::FailedToCreatePathBuf),
981        };
982
983        let mmap = MmapMutWrapper::attempt_restore(empty_leaf, depth, path_buf)?;
984
985        Ok(Self {
986            depth,
987            root_index: 1,
988            storage: Arc::new(Mutex::new(mmap)),
989        })
990    }
991
992    fn with_ref<F, R>(&self, fun: F) -> R
993    where
994        F: FnOnce(DenseTreeMMapRef<H>) -> R,
995    {
996        let guard = self.storage.lock().expect("lock poisoned, terminating");
997        let r = DenseTreeMMapRef {
998            depth: self.depth,
999            root_index: self.root_index,
1000            storage: &guard,
1001            locked_storage: &self.storage,
1002        };
1003        fun(r)
1004    }
1005
1006    fn write_proof(&self, index: usize, path: &mut Vec<Branch<H::Hash>>) {
1007        self.with_ref(|r| r.write_proof(index, path));
1008    }
1009
1010    fn get_leaf(&self, index: usize) -> H::Hash {
1011        self.with_ref(|r| {
1012            let leaf_index_in_dense_tree = index + (self.root_index << self.depth);
1013            r.storage[leaf_index_in_dense_tree]
1014        })
1015    }
1016
1017    fn update_with_mutation_condition(
1018        &self,
1019        index: usize,
1020        value: &H::Hash,
1021        is_mutation_allowed: bool,
1022    ) -> AnyTree<H> {
1023        if is_mutation_allowed {
1024            self.update_with_mutation(index, value);
1025            self.clone().into()
1026        } else {
1027            self.with_ref(|r| r.update(index, value)).into()
1028        }
1029    }
1030
1031    fn update_with_mutation(&self, index: usize, value: &H::Hash) {
1032        let mut storage = self.storage.lock().expect("lock poisoned, terminating");
1033        let leaf_index_in_dense_tree = index + (self.root_index << self.depth);
1034        storage[leaf_index_in_dense_tree] = *value;
1035        let mut current = leaf_index_in_dense_tree / 2;
1036        while current > 0 {
1037            let left = &storage[2 * current];
1038            let right = &storage[2 * current + 1];
1039            storage[current] = H::hash_node(left, right);
1040            current /= 2;
1041        }
1042    }
1043
1044    fn root(&self) -> H::Hash {
1045        self.storage.lock().expect("lock poisoned")[self.root_index]
1046    }
1047}
1048
1049struct DenseTreeMMapRef<'a, H: Hasher> {
1050    depth: usize,
1051    root_index: usize,
1052    storage: &'a MmapMutWrapper<H>,
1053    locked_storage: &'a Arc<Mutex<MmapMutWrapper<H>>>,
1054}
1055
1056impl<H: Hasher> From<DenseTreeMMapRef<'_, H>> for DenseMMapTree<H> {
1057    fn from(value: DenseTreeMMapRef<H>) -> Self {
1058        Self {
1059            depth: value.depth,
1060            root_index: value.root_index,
1061            storage: value.locked_storage.clone(),
1062        }
1063    }
1064}
1065
1066impl<H: Hasher> From<DenseTreeMMapRef<'_, H>> for AnyTree<H> {
1067    fn from(value: DenseTreeMMapRef<H>) -> Self {
1068        Self::DenseMMap(value.into())
1069    }
1070}
1071
1072impl<H> DenseTreeMMapRef<'_, H>
1073where
1074    H: Hasher,
1075    <H as Hasher>::Hash: Hash,
1076{
1077    fn root(&self) -> H::Hash {
1078        self.storage[self.root_index]
1079    }
1080
1081    const fn left(&self) -> DenseTreeMMapRef<H> {
1082        Self {
1083            depth: self.depth - 1,
1084            root_index: 2 * self.root_index,
1085            storage: self.storage,
1086            locked_storage: self.locked_storage,
1087        }
1088    }
1089
1090    const fn right(&self) -> DenseTreeMMapRef<H> {
1091        Self {
1092            depth: self.depth - 1,
1093            root_index: 2 * self.root_index + 1,
1094            storage: self.storage,
1095            locked_storage: self.locked_storage,
1096        }
1097    }
1098
1099    fn write_proof(&self, index: usize, path: &mut Vec<Branch<H::Hash>>) {
1100        if self.depth == 0 {
1101            return;
1102        }
1103        let next_index = clear_turn_at_depth(index, self.depth);
1104        if get_turn_at_depth(index, self.depth) == Turn::Left {
1105            path.push(Branch::Left(self.right().root()));
1106            self.left().write_proof(next_index, path);
1107        } else {
1108            path.push(Branch::Right(self.left().root()));
1109            self.right().write_proof(next_index, path);
1110        }
1111    }
1112
1113    fn update(&self, index: usize, hash: &H::Hash) -> SparseTree<H> {
1114        if self.depth == 0 {
1115            return SparseTree::new_leaf(*hash);
1116        }
1117        let next_index = clear_turn_at_depth(index, self.depth);
1118        if get_turn_at_depth(index, self.depth) == Turn::Left {
1119            let left = self.left();
1120            let new_left = left.update(next_index, hash);
1121            let right = self.right();
1122            let new_root = H::hash_node(&new_left.root(), &right.root());
1123            SparseTree {
1124                children: Some(Children {
1125                    left: Arc::new(new_left.into()),
1126                    right: Arc::new(self.right().into()),
1127                }),
1128                root: new_root,
1129                depth: self.depth,
1130            }
1131        } else {
1132            let right = self.right();
1133            let new_right = right.update(next_index, hash);
1134            let left = self.left();
1135            let new_root = H::hash_node(&left.root(), &new_right.root());
1136            SparseTree {
1137                children: Some(Children {
1138                    left: Arc::new(self.left().into()),
1139                    right: Arc::new(new_right.into()),
1140                }),
1141                root: new_root,
1142                depth: self.depth,
1143            }
1144        }
1145    }
1146}
1147
1148pub struct MmapMutWrapper<H: Hasher> {
1149    mmap: MmapMut,
1150    phantom: std::marker::PhantomData<H>,
1151}
1152
1153impl<H> MmapMutWrapper<H>
1154where
1155    H: Hasher,
1156    <H as Hasher>::Hash: Hash,
1157{
1158    /// Creates a new memory map backed with file with provided size
1159    /// and fills the entire map with initial value
1160    ///
1161    /// # Errors
1162    ///
1163    /// - returns Err if file creation has failed
1164    /// - returns Err if bytes couldn't be written to file
1165    ///
1166    /// # Panics
1167    ///
1168    /// - empty hash value serialization failed
1169    /// - file size cannot be set
1170    /// - file is too large, possible truncation can occur
1171    /// - cannot build memory map
1172    pub fn new_from_storage(
1173        file_path: PathBuf,
1174        storage: &[H::Hash],
1175    ) -> Result<Self, DenseMMapError> {
1176        // Safety: potential uninitialized padding from `H::Hash` is safe to use if
1177        // we're casting back to the same type.
1178        let buf = bytemuck::cast_slice(storage);
1179        let buf_len = buf.len();
1180
1181        let mut file = match OpenOptions::new()
1182            .read(true)
1183            .write(true)
1184            .create(true)
1185            .truncate(true)
1186            .open(file_path)
1187        {
1188            Ok(file) => file,
1189            Err(_e) => return Err(DenseMMapError::FileCreationFailed),
1190        };
1191
1192        file.set_len(buf_len as u64).expect("cannot set file size");
1193        if file.write_all(buf).is_err() {
1194            return Err(DenseMMapError::FileCannotWriteBytes);
1195        }
1196
1197        let mmap = unsafe {
1198            MmapOptions::new(usize::try_from(buf_len as u64).expect("file size truncated"))
1199                .expect("cannot create memory map")
1200                .with_file(&file, 0)
1201                .with_flags(MmapFlags::SHARED)
1202                .map_mut()
1203                .expect("cannot build memory map")
1204        };
1205
1206        Ok(Self {
1207            mmap,
1208            phantom: std::marker::PhantomData,
1209        })
1210    }
1211
1212    /// Given the file path and tree depth,
1213    /// it attempts to restore the memory map
1214    ///
1215    /// # Errors
1216    ///
1217    /// - returns Err if file doesn't exist
1218    /// - returns Err if file size doesn't match the expected tree size
1219    ///
1220    /// # Panics
1221    ///
1222    /// - cannot get file metadata to check for file length
1223    /// - truncated file size when attempting to build memory map
1224    /// - cannot build memory map
1225    pub fn attempt_restore(
1226        empty_leaf: &H::Hash,
1227        depth: usize,
1228        file_path: PathBuf,
1229    ) -> Result<Self, DenseMMapError> {
1230        let file = match OpenOptions::new().read(true).write(true).open(file_path) {
1231            Ok(file) => file,
1232            Err(_e) => return Err(DenseMMapError::FileDoesntExist),
1233        };
1234
1235        let size_of_empty_leaf = std::mem::size_of_val(empty_leaf);
1236        let expected_file_size = (1 << (depth + 1)) * size_of_empty_leaf as u64;
1237
1238        if expected_file_size != file.metadata().expect("cannot get file metadata").len() {
1239            return Err(DenseMMapError::FileSizeShouldMatchTree);
1240        }
1241
1242        let mmap = unsafe {
1243            MmapOptions::new(
1244                usize::try_from(expected_file_size).expect("expected file size truncated"),
1245            )
1246            .expect("cannot create memory map")
1247            .with_file(&file, 0)
1248            .with_flags(MmapFlags::SHARED)
1249            .map_mut()
1250            .expect("cannot build memory map")
1251        };
1252
1253        Ok(Self {
1254            mmap,
1255            phantom: std::marker::PhantomData,
1256        })
1257    }
1258}
1259
1260impl<H> Deref for MmapMutWrapper<H>
1261where
1262    H: Hasher,
1263    <H as Hasher>::Hash: Hash,
1264{
1265    type Target = [H::Hash];
1266
1267    fn deref(&self) -> &Self::Target {
1268        bytemuck::cast_slice(self.mmap.as_slice())
1269    }
1270}
1271
1272impl<H> DerefMut for MmapMutWrapper<H>
1273where
1274    H: Hasher,
1275    <H as Hasher>::Hash: Hash,
1276{
1277    fn deref_mut(&mut self) -> &mut Self::Target {
1278        bytemuck::cast_slice_mut(self.mmap.as_mut_slice())
1279    }
1280}
1281
1282#[derive(Error, Debug)]
1283pub enum DenseMMapError {
1284    #[error("file size should match expected tree size")]
1285    FileSizeShouldMatchTree,
1286    #[error("file doesn't exist")]
1287    FileDoesntExist,
1288    #[error("failed to create a file")]
1289    FileCreationFailed,
1290    #[error("cannot write bytes to file")]
1291    FileCannotWriteBytes,
1292    #[error("failed to create pathbuf")]
1293    FailedToCreatePathBuf,
1294}
1295
1296#[cfg(test)]
1297mod tests {
1298    use hex_literal::hex;
1299    use semaphore_rs_hasher::Hasher;
1300    use semaphore_rs_keccak::keccak::Keccak256;
1301
1302    use super::*;
1303
1304    struct TestHasher;
1305
1306    impl Hasher for TestHasher {
1307        type Hash = u64;
1308
1309        fn hash_node(left: &Self::Hash, right: &Self::Hash) -> Self::Hash {
1310            left + 2 * right + 1
1311        }
1312    }
1313
1314    #[test]
1315    fn test_updates_in_sparse() {
1316        let tree_1 = LazyMerkleTree::<TestHasher>::new(2, 0);
1317        assert_eq!(tree_1.root(), 4);
1318        let tree_2 = tree_1.update(0, &1);
1319        assert_eq!(tree_1.root(), 4);
1320        assert_eq!(tree_2.root(), 5);
1321        let tree_3 = tree_2.update(2, &2);
1322        assert_eq!(tree_1.root(), 4);
1323        assert_eq!(tree_2.root(), 5);
1324        assert_eq!(tree_3.root(), 9);
1325    }
1326
1327    #[test]
1328    fn test_updates_in_dense() {
1329        let tree_1 = LazyMerkleTree::<TestHasher>::new_with_dense_prefix(2, 2, &0);
1330        assert_eq!(tree_1.root(), 4);
1331        let tree_2 = tree_1.update(0, &1);
1332        assert_eq!(tree_1.root(), 4);
1333        assert_eq!(tree_2.root(), 5);
1334        let tree_3 = tree_2.update(2, &2);
1335        assert_eq!(tree_1.root(), 4);
1336        assert_eq!(tree_2.root(), 5);
1337        assert_eq!(tree_3.root(), 9);
1338    }
1339
1340    #[test]
1341    fn test_mutable_updates_in_dense() {
1342        let tree = LazyMerkleTree::<Keccak256>::new_with_dense_prefix(2, 2, &[0; 32]);
1343        let original_tree = LazyMerkleTree {
1344            tree: tree.tree.clone(),
1345            _version: Derived,
1346        };
1347        assert_eq!(
1348            original_tree.root(),
1349            hex!("b4c11951957c6f8f642c4af61cd6b24640fec6dc7fc607ee8206a99e92410d30")
1350        );
1351        let tree = tree.update_with_mutation(
1352            0,
1353            &hex!("0000000000000000000000000000000000000000000000000000000000000001"),
1354        );
1355        assert_eq!(
1356            original_tree.root(),
1357            hex!("c1ba1812ff680ce84c1d5b4f1087eeb08147a4d510f3496b2849df3a73f5af95")
1358        );
1359        let tree = tree.update_with_mutation(
1360            1,
1361            &hex!("0000000000000000000000000000000000000000000000000000000000000002"),
1362        );
1363        assert_eq!(
1364            original_tree.root(),
1365            hex!("893760ec5b5bee236f29e85aef64f17139c3c1b7ff24ce64eb6315fca0f2485b")
1366        );
1367        let tree = tree.update_with_mutation(
1368            2,
1369            &hex!("0000000000000000000000000000000000000000000000000000000000000003"),
1370        );
1371        assert_eq!(
1372            original_tree.root(),
1373            hex!("222ff5e0b5877792c2bc1670e2ccd0c2c97cd7bb1672a57d598db05092d3d72c")
1374        );
1375        let _tree = tree.update_with_mutation(
1376            3,
1377            &hex!("0000000000000000000000000000000000000000000000000000000000000004"),
1378        );
1379        assert_eq!(
1380            original_tree.root(),
1381            hex!("a9bb8c3f1f12e9aa903a50c47f314b57610a3ab32f2d463293f58836def38d36")
1382        );
1383    }
1384
1385    #[test]
1386    fn test_mutable_updates_in_dense_with_dense_prefix() {
1387        let h0 = [0; 32];
1388        let h1 = hex!("0000000000000000000000000000000000000000000000000000000000000001");
1389        let h2 = hex!("0000000000000000000000000000000000000000000000000000000000000002");
1390        let h3 = hex!("0000000000000000000000000000000000000000000000000000000000000003");
1391        let h4 = hex!("0000000000000000000000000000000000000000000000000000000000000004");
1392        let tree = LazyMerkleTree::<Keccak256>::new_with_dense_prefix(2, 1, &[0; 32]);
1393        let original_tree = LazyMerkleTree {
1394            tree: tree.tree.clone(),
1395            _version: Derived,
1396        };
1397        assert_eq!(
1398            tree.root(),
1399            hex!("b4c11951957c6f8f642c4af61cd6b24640fec6dc7fc607ee8206a99e92410d30")
1400        );
1401        let t1 = tree.update_with_mutation(0, &h1);
1402        assert_eq!(
1403            t1.root(),
1404            hex!("c1ba1812ff680ce84c1d5b4f1087eeb08147a4d510f3496b2849df3a73f5af95")
1405        );
1406        let t2 = t1.update_with_mutation(1, &h2);
1407        assert_eq!(
1408            t2.root(),
1409            hex!("893760ec5b5bee236f29e85aef64f17139c3c1b7ff24ce64eb6315fca0f2485b")
1410        );
1411        let t3 = t2.update_with_mutation(2, &h3);
1412        assert_eq!(
1413            t3.root(),
1414            hex!("222ff5e0b5877792c2bc1670e2ccd0c2c97cd7bb1672a57d598db05092d3d72c")
1415        );
1416        let t4 = t3.update_with_mutation(3, &h4);
1417        assert_eq!(
1418            t4.root(),
1419            hex!("a9bb8c3f1f12e9aa903a50c47f314b57610a3ab32f2d463293f58836def38d36")
1420        );
1421        // first two leaves are in the dense subtree, the rest is sparse, therefore only
1422        // first 2 get updated inplace.
1423        assert_eq!(
1424            original_tree.leaves().collect::<Vec<_>>(),
1425            vec![h1, h2, h0, h0]
1426        );
1427        // all leaves are updated in the properly tracked tree
1428        assert_eq!(t4.leaves().collect::<Vec<_>>(), vec![h1, h2, h3, h4]);
1429    }
1430
1431    #[test]
1432    fn test_proof() {
1433        let tree = LazyMerkleTree::<Keccak256>::new_with_dense_prefix(2, 1, &[0; 32]);
1434        let tree = tree.update_with_mutation(
1435            0,
1436            &hex!("0000000000000000000000000000000000000000000000000000000000000001"),
1437        );
1438        let tree = tree.update_with_mutation(
1439            1,
1440            &hex!("0000000000000000000000000000000000000000000000000000000000000002"),
1441        );
1442        let tree = tree.update_with_mutation(
1443            2,
1444            &hex!("0000000000000000000000000000000000000000000000000000000000000003"),
1445        );
1446        let tree = tree.update_with_mutation(
1447            3,
1448            &hex!("0000000000000000000000000000000000000000000000000000000000000004"),
1449        );
1450
1451        let proof = tree.proof(2);
1452        assert_eq!(proof.leaf_index(), 2);
1453        assert!(tree.verify(
1454            hex!("0000000000000000000000000000000000000000000000000000000000000003"),
1455            &proof
1456        ));
1457        assert!(!tree.verify(
1458            hex!("0000000000000000000000000000000000000000000000000000000000000001"),
1459            &proof
1460        ));
1461    }
1462
1463    #[test]
1464    fn test_giant_tree_with_initial_vals() {
1465        let h0 = [0; 32];
1466        let h1 = hex!("0000000000000000000000000000000000000000000000000000000000000001");
1467        let h2 = hex!("0000000000000000000000000000000000000000000000000000000000000002");
1468        let h3 = hex!("0000000000000000000000000000000000000000000000000000000000000003");
1469        let h4 = hex!("0000000000000000000000000000000000000000000000000000000000000004");
1470        let updates: Vec<(usize, _)> = vec![(0, h1), (1, h2), (2, h3), (3, h4)];
1471        let mut from_empty =
1472            LazyMerkleTree::<Keccak256>::new_with_dense_prefix(63, 10, &h0).derived();
1473        for (ix, hash) in &updates {
1474            from_empty = from_empty.update(*ix, hash);
1475        }
1476        let from_initial_vals =
1477            LazyMerkleTree::<Keccak256>::new_with_dense_prefix_with_initial_values(
1478                63,
1479                10,
1480                &h0,
1481                &[h1, h2, h3, h4],
1482            )
1483            .derived();
1484        assert_eq!(from_empty.root(), from_initial_vals.root());
1485    }
1486
1487    #[test]
1488    fn test_giant_trees() {
1489        let h0 = [0; 32];
1490        let h1 = hex!("0000000000000000000000000000000000000000000000000000000000000001");
1491        let h2 = hex!("0000000000000000000000000000000000000000000000000000000000000002");
1492        let h3 = hex!("0000000000000000000000000000000000000000000000000000000000000003");
1493        let h4 = hex!("0000000000000000000000000000000000000000000000000000000000000004");
1494        let updates: Vec<(usize, _)> = vec![
1495            (1, h1),
1496            (2, h2),
1497            (1_000_000_000, h3),
1498            (1_000_000_000_000, h4),
1499        ];
1500        let mut tree = LazyMerkleTree::<Keccak256>::new_with_dense_prefix(63, 10, &h0).derived();
1501        for (ix, hash) in &updates {
1502            tree = tree.update(*ix, hash);
1503        }
1504        for (ix, hash) in &updates {
1505            let proof = tree.proof(*ix);
1506            assert_eq!(proof.root(*hash), tree.root());
1507        }
1508        let first_three_leaves = tree.leaves().take(3).collect::<Vec<_>>();
1509        assert_eq!(first_three_leaves, vec![h0, h1, h2]);
1510
1511        let mut tree = LazyMerkleTree::<Keccak256>::new_with_dense_prefix(63, 10, &h0);
1512        let original_tree = tree.derived();
1513        for (ix, hash) in &updates {
1514            tree = tree.update_with_mutation(*ix, hash);
1515        }
1516        for (ix, hash) in &updates {
1517            let proof = tree.proof(*ix);
1518            assert_eq!(proof.root(*hash), tree.root());
1519        }
1520        let first_three_leaves = original_tree.leaves().take(3).collect::<Vec<_>>();
1521        assert_eq!(first_three_leaves, vec![h0, h1, h2]);
1522        let first_three_leaves = tree.leaves().take(3).collect::<Vec<_>>();
1523        assert_eq!(first_three_leaves, vec![h0, h1, h2]);
1524    }
1525
1526    #[test]
1527    fn test_dense_mmap_tree() {
1528        let h0 = [0; 32];
1529        let h1 = hex!("0000000000000000000000000000000000000000000000000000000000000001");
1530        let h2 = hex!("0000000000000000000000000000000000000000000000000000000000000002");
1531        let h3 = hex!("0000000000000000000000000000000000000000000000000000000000000003");
1532        let h4 = hex!("0000000000000000000000000000000000000000000000000000000000000004");
1533        let h5 = hex!("0000000000000000000000000000000000000000000000000000000000000005");
1534        let h6 = hex!("0000000000000000000000000000000000000000000000000000000000000006");
1535        let h7 = hex!("0000000000000000000000000000000000000000000000000000000000000007");
1536        let h8 = hex!("0000000000000000000000000000000000000000000000000000000000000008");
1537
1538        let initial_values = vec![h1, h2, h3, h4, h5, h6, h7, h8];
1539
1540        let tree: LazyMerkleTree<Keccak256, Canonical> =
1541            LazyMerkleTree::<Keccak256>::new_mmapped_with_dense_prefix_with_init_values(
1542                3,
1543                3,
1544                &h0,
1545                &initial_values,
1546                "./testfile",
1547            )
1548            .unwrap();
1549        let tree_leaves = tree.leaves().collect::<Vec<_>>();
1550
1551        assert_eq!(tree_leaves, initial_values);
1552
1553        let proof_h1 = tree.proof(0);
1554        assert!(tree.verify(h1, &proof_h1));
1555
1556        let proof_h2 = tree.proof(1);
1557        assert!(tree.verify(h2, &proof_h2));
1558
1559        // drop a tree, the mmap file should still be there
1560        drop(tree);
1561
1562        let tree: LazyMerkleTree<Keccak256, Canonical> =
1563            LazyMerkleTree::<Keccak256>::attempt_dense_mmap_restore(3, 3, &h0, "./testfile")
1564                .unwrap();
1565
1566        // repeat asserts again
1567        let tree_leaves = tree.leaves().collect::<Vec<_>>();
1568
1569        assert_eq!(tree_leaves, initial_values);
1570
1571        let proof_h1 = tree.proof(0);
1572        assert!(tree.verify(h1, &proof_h1));
1573
1574        let proof_h2 = tree.proof(1);
1575        assert!(tree.verify(h2, &proof_h2));
1576
1577        // remove mmap file at the end
1578        std::fs::remove_file("./testfile").unwrap();
1579    }
1580}
semaphore_rs_trees/lazy/mod.rs

semaphore_rs_trees/lazy/
mod.rs
