incrementalmerkletree/
frontier.rs

1use alloc::vec::Vec;
2use core::mem::size_of;
3
4use crate::{Address, Hashable, Level, MerklePath, Position, Source};
5
6#[cfg(feature = "legacy-api")]
7use {alloc::boxed::Box, alloc::collections::VecDeque, core::iter::repeat};
8
9#[cfg(any(test, feature = "test-dependencies"))]
10use {
11    core::num::{NonZeroU64, NonZeroU8},
12    rand::{
13        distributions::{Distribution, Standard},
14        Rng, RngCore,
15    },
16};
17
18/// Validation errors that can occur during reconstruction of a Merkle frontier from
19/// its constituent parts.
20#[derive(Clone, Debug, PartialEq, Eq)]
21pub enum FrontierError {
22    /// An error representing that the number of ommers provided in frontier construction does not
23    /// the expected length of the ommers list given the position.
24    PositionMismatch { expected_ommers: u8 },
25    /// An error representing that the position and/or list of ommers provided to frontier
26    /// construction would result in a frontier that exceeds the maximum statically allowed depth
27    /// of the tree. `depth` is the minimum tree depth that would be required in order for that
28    /// tree to contain the position in question.
29    MaxDepthExceeded { depth: u8 },
30}
31
32/// A [`NonEmptyFrontier`] is a reduced representation of a Merkle tree, containing a single leaf
33/// value, along with the vector of hashes produced by the reduction of previously appended leaf
34/// values that will be required when producing a witness for the current leaf.
35#[derive(Clone, Debug, PartialEq, Eq)]
36pub struct NonEmptyFrontier<H> {
37    position: Position,
38    leaf: H,
39    ommers: Vec<H>,
40}
41
42impl<H> NonEmptyFrontier<H> {
43    /// Constructs a new frontier with the specified value at position 0.
44    pub fn new(leaf: H) -> Self {
45        Self {
46            position: 0.into(),
47            leaf,
48            ommers: vec![],
49        }
50    }
51
52    /// Constructs a new frontier from its constituent parts.
53    pub fn from_parts(position: Position, leaf: H, ommers: Vec<H>) -> Result<Self, FrontierError> {
54        let expected_ommers = position.past_ommer_count();
55        if ommers.len() == expected_ommers.into() {
56            Ok(Self {
57                position,
58                leaf,
59                ommers,
60            })
61        } else {
62            Err(FrontierError::PositionMismatch { expected_ommers })
63        }
64    }
65
66    /// Decomposes the frontier into its constituent parts
67    pub fn into_parts(self) -> (Position, H, Vec<H>) {
68        (self.position, self.leaf, self.ommers)
69    }
70
71    /// Returns the position of the most recently appended leaf.
72    pub fn position(&self) -> Position {
73        self.position
74    }
75
76    /// Returns the leaf most recently appended to the frontier.
77    pub fn leaf(&self) -> &H {
78        &self.leaf
79    }
80
81    /// Returns the list of past hashes required to construct a witness for the
82    /// leaf most recently appended to the frontier.
83    pub fn ommers(&self) -> &[H] {
84        &self.ommers
85    }
86}
87
88impl<H: Hashable + Clone> NonEmptyFrontier<H> {
89    /// Append a new leaf to the frontier, and recompute ommers by hashing together full subtrees
90    /// until an empty ommer slot is found.
91    pub fn append(&mut self, leaf: H) {
92        let prior_position = self.position;
93        let prior_leaf = self.leaf.clone();
94        self.position += 1;
95        self.leaf = leaf;
96        if self.position.is_right_child() {
97            // if the new position is a right-hand leaf, the current leaf will directly become an
98            // ommer at level 0, and there is no other mutation made to the tree.
99            self.ommers.insert(0, prior_leaf);
100        } else {
101            // if the new position is even, then the current leaf will be hashed
102            // with the first ommer, and so forth up the tree.
103            let new_root_level = self.position.root_level();
104
105            let mut carry = Some((prior_leaf, 0.into()));
106            let mut new_ommers = Vec::with_capacity(self.position.past_ommer_count().into());
107            for (addr, source) in prior_position.witness_addrs(new_root_level) {
108                if let Source::Past(i) = source {
109                    if let Some((carry_ommer, carry_lvl)) = carry.as_ref() {
110                        if *carry_lvl == addr.level() {
111                            carry = Some((
112                                H::combine(addr.level(), &self.ommers[usize::from(i)], carry_ommer),
113                                addr.level() + 1,
114                            ))
115                        } else {
116                            // insert the carry at the first empty slot; then the rest of the
117                            // ommers will remain unchanged
118                            new_ommers.push(carry_ommer.clone());
119                            new_ommers.push(self.ommers[usize::from(i)].clone());
120                            carry = None;
121                        }
122                    } else {
123                        // when there's no carry, just push on the ommer value
124                        new_ommers.push(self.ommers[usize::from(i)].clone());
125                    }
126                }
127            }
128
129            // we carried value out, so we need to push on one more ommer.
130            if let Some((carry_ommer, _)) = carry {
131                new_ommers.push(carry_ommer);
132            }
133
134            self.ommers = new_ommers;
135        }
136    }
137
138    /// Generate the root of the Merkle tree by hashing against empty subtree roots.
139    pub fn root(&self, root_level: Option<Level>) -> H {
140        let max_level = root_level.unwrap_or_else(|| self.position.root_level());
141        self.position
142            .witness_addrs(max_level)
143            .fold(
144                (self.leaf.clone(), Level::from(0)),
145                |(digest, complete_lvl), (addr, source)| {
146                    // fold up from complete_lvl to addr.level() pairing with empty roots; if
147                    // complete_lvl == addr.level() this is just the complete digest to this point
148                    let digest = complete_lvl
149                        .iter_to(addr.level())
150                        .fold(digest, |d, l| H::combine(l, &d, &H::empty_root(l)));
151
152                    let res_digest = match source {
153                        Source::Past(i) => {
154                            H::combine(addr.level(), &self.ommers[usize::from(i)], &digest)
155                        }
156                        Source::Future => {
157                            H::combine(addr.level(), &digest, &H::empty_root(addr.level()))
158                        }
159                    };
160
161                    (res_digest, addr.level() + 1)
162                },
163            )
164            .0
165    }
166
167    /// Constructs a witness for the leaf at the tip of this frontier, given a source of node
168    /// values that complement this frontier.
169    ///
170    /// If the `complement_nodes` function returns `None` when the value is requested at a given
171    /// tree address, the address at which the failure occurs will be returned as an error.
172    pub fn witness<F>(&self, depth: u8, complement_nodes: F) -> Result<Vec<H>, Address>
173    where
174        F: Fn(Address) -> Option<H>,
175    {
176        // construct a complete trailing edge that includes the data from
177        // the following frontier not yet included in the trailing edge.
178        self.position()
179            .witness_addrs(depth.into())
180            .map(|(addr, source)| match source {
181                Source::Past(i) => Ok(self.ommers[usize::from(i)].clone()),
182                Source::Future => complement_nodes(addr).ok_or(addr),
183            })
184            .collect::<Result<Vec<_>, _>>()
185    }
186}
187
188#[cfg(any(test, feature = "test-dependencies"))]
189impl<H: Hashable + Clone> NonEmptyFrontier<H>
190where
191    Standard: Distribution<H>,
192{
193    /// Generates a random frontier of a Merkle tree having the specified nonzero size.
194    pub fn random_of_size<R: RngCore>(rng: &mut R, tree_size: NonZeroU64) -> Self {
195        let position = (u64::from(tree_size) - 1).into();
196        NonEmptyFrontier::from_parts(
197            position,
198            rng.gen(),
199            core::iter::repeat_with(|| rng.gen())
200                .take(position.past_ommer_count().into())
201                .collect(),
202        )
203        .unwrap()
204    }
205
206    pub fn random_with_prior_subtree_roots<R: RngCore>(
207        rng: &mut R,
208        tree_size: NonZeroU64,
209        subtree_depth: NonZeroU8,
210    ) -> (Vec<H>, Self) {
211        let prior_subtree_count: u64 = u64::from(tree_size) >> u8::from(subtree_depth);
212        if prior_subtree_count > 0 {
213            let prior_roots: Vec<H> = core::iter::repeat_with(|| rng.gen())
214                .take(prior_subtree_count as usize)
215                .collect();
216
217            let subtree_root_level = Level::from(u8::from(subtree_depth));
218
219            // Generate replacement ommers for the random frontier from the prior subtree roots.
220            let mut replacement_ommers: Vec<(Level, H)> = vec![];
221            let mut roots_iter = prior_roots.iter();
222            loop {
223                if let Some(top) = replacement_ommers.pop() {
224                    if let Some(prev) = replacement_ommers.pop() {
225                        if top.0 == prev.0 {
226                            // Combine, then continue the outer loop so that we eagerly combine as
227                            // many values from the stack as we can before pushing more on.
228                            replacement_ommers
229                                .push((top.0 + 1, H::combine(top.0, &prev.1, &top.1)));
230                            continue;
231                        } else {
232                            // We can't combine yet, so push `prev` back on. `top` will get pushed
233                            // back on or consumed below.
234                            replacement_ommers.push(prev);
235                        }
236                    }
237
238                    if let Some(root) = roots_iter.next() {
239                        if top.0 == subtree_root_level {
240                            replacement_ommers.push((
241                                subtree_root_level + 1,
242                                H::combine(subtree_root_level, &top.1, root),
243                            ));
244                        } else {
245                            replacement_ommers.push(top);
246                            replacement_ommers.push((subtree_root_level, root.clone()));
247                        }
248                    } else {
249                        // No more roots, so we just push `top` back on and break.
250                        replacement_ommers.push(top);
251                        break;
252                    }
253                } else if let Some(root) = roots_iter.next() {
254                    replacement_ommers.push((subtree_root_level, root.clone()));
255                } else {
256                    break;
257                }
258            }
259
260            let mut result = Self::random_of_size(rng, tree_size);
261            let olen = result.ommers.len();
262            for (idx, (_, ommer)) in replacement_ommers.into_iter().enumerate() {
263                result.ommers[olen - (idx + 1)] = ommer;
264            }
265
266            (prior_roots, result)
267        } else {
268            (vec![], Self::random_of_size(rng, tree_size))
269        }
270    }
271}
272
273/// A possibly-empty Merkle frontier.
274#[derive(Debug, Clone, PartialEq, Eq)]
275pub struct Frontier<H, const DEPTH: u8> {
276    frontier: Option<NonEmptyFrontier<H>>,
277}
278
279impl<H, const DEPTH: u8> TryFrom<NonEmptyFrontier<H>> for Frontier<H, DEPTH> {
280    type Error = FrontierError;
281    fn try_from(f: NonEmptyFrontier<H>) -> Result<Self, FrontierError> {
282        if f.position.root_level() <= Level::from(DEPTH) {
283            Ok(Frontier { frontier: Some(f) })
284        } else {
285            Err(FrontierError::MaxDepthExceeded {
286                depth: f.position.root_level().into(),
287            })
288        }
289    }
290}
291
292impl<H, const DEPTH: u8> Frontier<H, DEPTH> {
293    /// Constructs a new empty frontier.
294    pub fn empty() -> Self {
295        Self { frontier: None }
296    }
297
298    /// Constructs a new frontier from its constituent parts.
299    ///
300    /// Returns an error if the new frontier would exceed the maximum allowed depth or if the list
301    /// of ommers provided is not consistent with the position of the leaf.
302    pub fn from_parts(position: Position, leaf: H, ommers: Vec<H>) -> Result<Self, FrontierError> {
303        NonEmptyFrontier::from_parts(position, leaf, ommers).and_then(Self::try_from)
304    }
305
306    /// Return the wrapped NonEmptyFrontier reference, or None if the frontier is empty.
307    pub fn value(&self) -> Option<&NonEmptyFrontier<H>> {
308        self.frontier.as_ref()
309    }
310
311    /// Consumes this wrapper and returns the underlying `Option<NonEmptyFrontier>`
312    pub fn take(self) -> Option<NonEmptyFrontier<H>> {
313        self.frontier
314    }
315
316    /// Returns the amount of memory dynamically allocated for ommer values within the frontier.
317    pub fn dynamic_memory_usage(&self) -> usize {
318        self.frontier.as_ref().map_or(0, |f| {
319            size_of::<usize>() + (f.ommers.capacity() + 1) * size_of::<H>()
320        })
321    }
322
323    /// Returns the size of the Merkle tree that this frontier corresponds to.
324    pub fn tree_size(&self) -> u64 {
325        self.frontier
326            .as_ref()
327            .map_or(0, |f| u64::from(f.position()) + 1)
328    }
329}
330
331impl<H: Hashable + Clone, const DEPTH: u8> Frontier<H, DEPTH> {
332    /// Appends a new value to the frontier at the next available slot.
333    /// Returns true if successful and false if the frontier would exceed
334    /// the maximum allowed depth.
335    pub fn append(&mut self, value: H) -> bool {
336        if let Some(frontier) = self.frontier.as_mut() {
337            if frontier.position().is_complete_subtree(DEPTH.into()) {
338                false
339            } else {
340                frontier.append(value);
341                true
342            }
343        } else {
344            self.frontier = Some(NonEmptyFrontier::new(value));
345            true
346        }
347    }
348
349    /// Obtains the current root of this Merkle frontier by hashing
350    /// against empty nodes up to the maximum height of the pruned
351    /// tree that the frontier represents.
352    pub fn root(&self) -> H {
353        self.frontier
354            .as_ref()
355            .map_or(H::empty_root(DEPTH.into()), |frontier| {
356                frontier.root(Some(DEPTH.into()))
357            })
358    }
359
360    /// Constructs a [`MerklePath`] to the leaf at the tip of this frontier, given a source of node
361    /// values that complement this frontier.
362    ///
363    /// If the `complement_nodes` function returns `None` when the value is requested at a given
364    /// tree address, the address at which the failure occurs will be returned as an error.
365    ///
366    /// Returns `Ok(Some(MerklePath))` if successful, `Ok(None)` if the frontier is empty,
367    /// or an error containing the address of the failure.
368    pub fn witness<F>(&self, complement_nodes: F) -> Result<Option<MerklePath<H, DEPTH>>, Address>
369    where
370        F: Fn(Address) -> Option<H>,
371    {
372        self.frontier
373            .as_ref()
374            .map(|f| {
375                f.witness(DEPTH, complement_nodes).map(|path_elems| {
376                    MerklePath::from_parts(path_elems, f.position())
377                        .expect("Path length should be equal to frontier depth.")
378                })
379            })
380            .transpose()
381    }
382}
383
384#[cfg(any(test, feature = "test-dependencies"))]
385impl<H: Hashable + Clone, const DEPTH: u8> Frontier<H, DEPTH>
386where
387    Standard: Distribution<H>,
388{
389    /// Generates a random frontier of a Merkle tree having the specified size.
390    pub fn random_of_size<R: RngCore>(rng: &mut R, tree_size: u64) -> Self {
391        assert!(tree_size <= 2u64.checked_pow(DEPTH.into()).unwrap());
392        Frontier {
393            frontier: NonZeroU64::new(tree_size)
394                .map(|sz| NonEmptyFrontier::random_of_size(rng, sz)),
395        }
396    }
397
398    pub fn random_with_prior_subtree_roots<R: RngCore>(
399        rng: &mut R,
400        tree_size: u64,
401        subtree_depth: NonZeroU8,
402    ) -> (Vec<H>, Self) {
403        assert!(tree_size <= 2u64.checked_pow(DEPTH.into()).unwrap());
404        NonZeroU64::new(tree_size).map_or((vec![], Frontier::empty()), |tree_size| {
405            let (prior_roots, frontier) =
406                NonEmptyFrontier::random_with_prior_subtree_roots(rng, tree_size, subtree_depth);
407            (
408                prior_roots,
409                Frontier {
410                    frontier: Some(frontier),
411                },
412            )
413        })
414    }
415}
416
417#[cfg(feature = "legacy-api")]
418#[cfg_attr(docsrs, doc(cfg(feature = "legacy-api")))]
419pub struct PathFiller<H> {
420    queue: VecDeque<H>,
421}
422
423#[cfg(feature = "legacy-api")]
424impl<H: Hashable> PathFiller<H> {
425    pub fn empty() -> Self {
426        PathFiller {
427            queue: VecDeque::new(),
428        }
429    }
430
431    pub fn new(queue: VecDeque<H>) -> Self {
432        Self { queue }
433    }
434
435    pub fn next(&mut self, level: Level) -> H {
436        self.queue
437            .pop_front()
438            .unwrap_or_else(|| H::empty_root(level))
439    }
440}
441
442/// A Merkle tree of note commitments.
443#[derive(Clone, Debug, PartialEq, Eq)]
444#[cfg(feature = "legacy-api")]
445#[cfg_attr(docsrs, doc(cfg(feature = "legacy-api")))]
446pub struct CommitmentTree<H, const DEPTH: u8> {
447    pub(crate) left: Option<H>,
448    pub(crate) right: Option<H>,
449    pub(crate) parents: Vec<Option<H>>,
450}
451
452#[cfg(feature = "legacy-api")]
453impl<H, const DEPTH: u8> CommitmentTree<H, DEPTH> {
454    /// Creates an empty tree.
455    pub fn empty() -> Self {
456        CommitmentTree {
457            left: None,
458            right: None,
459            parents: vec![],
460        }
461    }
462
463    #[allow(clippy::result_unit_err)]
464    pub fn from_parts(
465        left: Option<H>,
466        right: Option<H>,
467        parents: Vec<Option<H>>,
468    ) -> Result<Self, ()> {
469        if parents.len() < usize::from(DEPTH) {
470            Ok(CommitmentTree {
471                left,
472                right,
473                parents,
474            })
475        } else {
476            Err(())
477        }
478    }
479
480    pub fn is_empty(&self) -> bool {
481        self.left.is_none() && self.right.is_none()
482    }
483
484    pub fn left(&self) -> &Option<H> {
485        &self.left
486    }
487
488    pub fn right(&self) -> &Option<H> {
489        &self.right
490    }
491
492    pub fn parents(&self) -> &Vec<Option<H>> {
493        &self.parents
494    }
495
496    pub fn leaf(&self) -> Option<&H> {
497        self.right.as_ref().or(self.left.as_ref())
498    }
499
500    pub fn ommers_iter(&self) -> Box<dyn Iterator<Item = &'_ H> + '_> {
501        if self.right.is_some() {
502            Box::new(
503                self.left
504                    .iter()
505                    .chain(self.parents.iter().filter_map(|v| v.as_ref())),
506            )
507        } else {
508            Box::new(self.parents.iter().filter_map(|v| v.as_ref()))
509        }
510    }
511
512    /// Returns the number of leaf nodes in the tree.
513    pub fn size(&self) -> usize {
514        self.parents.iter().enumerate().fold(
515            match (self.left.as_ref(), self.right.as_ref()) {
516                (None, None) => 0,
517                (Some(_), None) => 1,
518                (Some(_), Some(_)) => 2,
519                (None, Some(_)) => unreachable!(),
520            },
521            |acc, (i, p)| {
522                // Treat occupation of parents array as a binary number
523                // (right-shifted by 1)
524                acc + if p.is_some() { 1 << (i + 1) } else { 0 }
525            },
526        )
527    }
528
529    pub(crate) fn is_complete(&self, depth: u8) -> bool {
530        if depth == 0 {
531            self.left.is_some() && self.right.is_none() && self.parents.is_empty()
532        } else {
533            self.left.is_some()
534                && self.right.is_some()
535                && self
536                    .parents
537                    .iter()
538                    .chain(repeat(&None))
539                    .take((depth - 1).into())
540                    .all(|p| p.is_some())
541        }
542    }
543}
544
545#[cfg(feature = "legacy-api")]
546impl<H: Hashable + Clone, const DEPTH: u8> CommitmentTree<H, DEPTH> {
547    pub fn from_frontier(frontier: &Frontier<H, DEPTH>) -> Self {
548        frontier.value().map_or_else(Self::empty, |f| {
549            let mut ommers_iter = f.ommers().iter().cloned();
550            let (left, right) = if f.position().is_right_child() {
551                (
552                    ommers_iter
553                        .next()
554                        .expect("An ommer must exist if the frontier position is odd"),
555                    Some(f.leaf().clone()),
556                )
557            } else {
558                (f.leaf().clone(), None)
559            };
560
561            Self {
562                left: Some(left),
563                right,
564                parents: (1u8..DEPTH)
565                    .map(|i| {
566                        if u64::from(f.position()) & (1 << i) == 0 {
567                            None
568                        } else {
569                            ommers_iter.next()
570                        }
571                    })
572                    .collect(),
573            }
574        })
575    }
576
577    pub fn to_frontier(&self) -> Frontier<H, DEPTH> {
578        if self.size() == 0 {
579            Frontier::empty()
580        } else {
581            let ommers_iter = self.parents.iter().filter_map(|v| v.as_ref()).cloned();
582            let (leaf, ommers) = match (self.left.as_ref(), self.right.as_ref()) {
583                (Some(a), None) => (a.clone(), ommers_iter.collect()),
584                (Some(a), Some(b)) => (
585                    b.clone(),
586                    Some(a.clone()).into_iter().chain(ommers_iter).collect(),
587                ),
588                _ => unreachable!(),
589            };
590
591            // If a frontier cannot be successfully constructed from the
592            // parts of a commitment tree, it is a programming error.
593            Frontier::from_parts((self.size() - 1).try_into().unwrap(), leaf, ommers)
594                .expect("Frontier should be constructable from CommitmentTree.")
595        }
596    }
597
598    /// Adds a leaf node to the tree.
599    ///
600    /// Returns an error if the tree is full.
601    #[allow(clippy::result_unit_err)]
602    pub fn append(&mut self, node: H) -> Result<(), ()> {
603        if self.is_complete(DEPTH) {
604            // Tree is full
605            return Err(());
606        }
607
608        match (&self.left, &self.right) {
609            (None, _) => self.left = Some(node),
610            (_, None) => self.right = Some(node),
611            (Some(l), Some(r)) => {
612                let mut combined = H::combine(0.into(), l, r);
613                self.left = Some(node);
614                self.right = None;
615
616                for i in 0..DEPTH {
617                    let i_usize = usize::from(i);
618                    if i_usize < self.parents.len() {
619                        if let Some(p) = &self.parents[i_usize] {
620                            combined = H::combine((i + 1).into(), p, &combined);
621                            self.parents[i_usize] = None;
622                        } else {
623                            self.parents[i_usize] = Some(combined);
624                            break;
625                        }
626                    } else {
627                        self.parents.push(Some(combined));
628                        break;
629                    }
630                }
631            }
632        }
633
634        Ok(())
635    }
636
637    /// Returns the current root of the tree.
638    pub fn root(&self) -> H {
639        self.root_at_depth(DEPTH, PathFiller::empty())
640    }
641
642    pub fn root_at_depth(&self, depth: u8, mut filler: PathFiller<H>) -> H {
643        assert!(depth > 0);
644
645        // 1) Hash left and right leaves together.
646        //    - Empty leaves are used as needed.
647        //    - Note that `filler.next` is side-effecting and so cannot be factored out.
648        let leaf_root = H::combine(
649            0.into(),
650            &self
651                .left
652                .as_ref()
653                .map_or_else(|| filler.next(0.into()), |n| n.clone()),
654            &self
655                .right
656                .as_ref()
657                .map_or_else(|| filler.next(0.into()), |n| n.clone()),
658        );
659
660        // 2) Extend the parents to the desired depth with None values, then hash from leaf to
661        //    root. Roots of the empty subtrees are used as needed.
662        self.parents
663            .iter()
664            .chain(repeat(&None))
665            .take((depth - 1).into())
666            .zip(0u8..)
667            .fold(leaf_root, |root, (p, i)| {
668                let level = Level::from(i + 1);
669                match p {
670                    Some(node) => H::combine(level, node, &root),
671                    None => H::combine(level, &root, &filler.next(level)),
672                }
673            })
674    }
675}
676
677#[cfg(any(all(test, feature = "std"), feature = "test-dependencies"))]
678pub mod testing {
679    use core::fmt::Debug;
680    use proptest::collection::vec;
681    use proptest::prelude::*;
682    use rand::{distributions::Standard, prelude::Distribution};
683
684    #[cfg(feature = "std")]
685    use {core::hash::Hasher, std::collections::hash_map::DefaultHasher};
686
687    use crate::{frontier::Frontier, Hashable, Level};
688
689    impl<H: Hashable + Clone, const DEPTH: u8> crate::testing::Frontier<H>
690        for super::Frontier<H, DEPTH>
691    {
692        fn append(&mut self, value: H) -> bool {
693            super::Frontier::append(self, value)
694        }
695
696        fn root(&self) -> H {
697            super::Frontier::root(self)
698        }
699    }
700
701    #[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
702    pub struct TestNode(pub u64);
703
704    #[cfg(feature = "std")]
705    impl Hashable for TestNode {
706        fn empty_leaf() -> Self {
707            Self(0)
708        }
709
710        fn combine(level: Level, a: &Self, b: &Self) -> Self {
711            let mut hasher = DefaultHasher::new();
712            hasher.write_u8(level.into());
713            hasher.write_u64(a.0);
714            hasher.write_u64(b.0);
715            Self(hasher.finish())
716        }
717    }
718
719    impl Distribution<TestNode> for Standard {
720        fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> TestNode {
721            TestNode(rng.gen())
722        }
723    }
724
725    pub fn arb_test_node() -> impl Strategy<Value = TestNode> + Clone {
726        any::<u64>().prop_map(TestNode)
727    }
728
729    pub fn arb_frontier<H: Hashable + Clone + Debug, T: Strategy<Value = H>, const DEPTH: u8>(
730        min_size: usize,
731        arb_node: T,
732    ) -> impl Strategy<Value = Frontier<H, DEPTH>> {
733        assert!((1 << DEPTH) >= min_size + 100);
734        vec(arb_node, min_size..(min_size + 100)).prop_map(move |v| {
735            let mut frontier = Frontier::empty();
736            for node in v.into_iter() {
737                frontier.append(node);
738            }
739            frontier
740        })
741    }
742
743    #[cfg(feature = "legacy-api")]
744    use crate::frontier::CommitmentTree;
745
746    #[cfg(feature = "legacy-api")]
747    #[cfg_attr(docsrs, doc(cfg(feature = "legacy-api")))]
748    pub fn arb_commitment_tree<
749        H: Hashable + Clone + Debug,
750        T: Strategy<Value = H>,
751        const DEPTH: u8,
752    >(
753        min_size: usize,
754        arb_node: T,
755    ) -> impl Strategy<Value = CommitmentTree<H, DEPTH>> {
756        assert!((1 << DEPTH) >= min_size + 100);
757        vec(arb_node, min_size..(min_size + 100)).prop_map(move |v| {
758            let mut tree = CommitmentTree::empty();
759            for node in v.into_iter() {
760                tree.append(node).unwrap();
761            }
762            tree.parents.resize_with((DEPTH - 1).into(), || None);
763            tree
764        })
765    }
766}
767
768#[cfg(all(test, feature = "std"))]
769mod tests {
770    use alloc::string::{String, ToString};
771    use rand::SeedableRng;
772    use rand_chacha::ChaChaRng;
773
774    use super::{testing::TestNode, *};
775
776    #[cfg(feature = "legacy-api")]
777    use {
778        super::testing::{arb_commitment_tree, arb_test_node},
779        proptest::prelude::*,
780    };
781
782    #[test]
783    fn nonempty_frontier_root() {
784        let mut frontier = NonEmptyFrontier::new("a".to_string());
785        assert_eq!(frontier.root(None), "a");
786
787        frontier.append("b".to_string());
788        assert_eq!(frontier.root(None), "ab");
789
790        frontier.append("c".to_string());
791        assert_eq!(frontier.root(None), "abc_");
792    }
793
794    #[test]
795    fn frontier_from_parts() {
796        assert!(super::Frontier::<(), 1>::from_parts(0.into(), (), vec![]).is_ok());
797        assert!(super::Frontier::<(), 1>::from_parts(1.into(), (), vec![()]).is_ok());
798        assert!(super::Frontier::<(), 1>::from_parts(0.into(), (), vec![()]).is_err());
799    }
800
801    #[test]
802    fn frontier_root() {
803        let mut frontier: super::Frontier<String, 4> = super::Frontier::empty();
804        assert_eq!(frontier.root().len(), 16);
805        assert_eq!(frontier.root(), "________________");
806
807        frontier.append("a".to_string());
808        assert_eq!(frontier.root(), "a_______________");
809
810        frontier.append("b".to_string());
811        assert_eq!(frontier.root(), "ab______________");
812
813        frontier.append("c".to_string());
814        assert_eq!(frontier.root(), "abc_____________");
815    }
816
817    #[test]
818    fn nonempty_frontier_witness() {
819        let mut frontier = NonEmptyFrontier::<String>::new("a".to_string());
820        for c in 'b'..='g' {
821            frontier.append(c.to_string());
822        }
823        let bridge_value_at = |addr: Address| match <u8>::from(addr.level()) {
824            0 => Some("h".to_string()),
825            3 => Some("xxxxxxxx".to_string()),
826            _ => None,
827        };
828
829        assert_eq!(
830            Ok(["h", "ef", "abcd", "xxxxxxxx"]
831                .map(|v| v.to_string())
832                .to_vec()),
833            frontier.witness(4, bridge_value_at)
834        );
835    }
836
837    #[test]
838    fn frontier_witness() {
839        let mut frontier = Frontier::<String, 4>::empty();
840        for c in 'a'..='g' {
841            frontier.append(c.to_string());
842        }
843
844        assert_eq!(
845            frontier
846                .witness(|addr| Some(String::empty_root(addr.level())))
847                .map(|maybe_p| maybe_p.map(|p| p.path_elems().to_vec())),
848            Ok(Some(
849                ["_", "ef", "abcd", "________"]
850                    .map(|v| v.to_string())
851                    .to_vec()
852            )),
853        );
854    }
855
856    #[test]
857    #[cfg(feature = "legacy-api")]
858    fn test_commitment_tree_complete() {
859        let mut t: CommitmentTree<TestNode, 6> = CommitmentTree::empty();
860        for n in 1u64..=32 {
861            t.append(TestNode(n)).unwrap();
862            // every tree of a power-of-two height is complete
863            let is_complete = n.count_ones() == 1;
864            let level = usize::BITS - 1 - n.leading_zeros(); //log2
865            assert_eq!(
866                is_complete,
867                t.is_complete(level.try_into().unwrap()),
868                "Tree {:?} {} complete at height {}",
869                t,
870                if is_complete {
871                    "should be"
872                } else {
873                    "should not be"
874                },
875                n
876            );
877        }
878    }
879
880    #[test]
881    #[cfg(feature = "legacy-api")]
882    fn test_commitment_tree_roundtrip() {
883        let ct = CommitmentTree {
884            left: Some("a".to_string()),
885            right: Some("b".to_string()),
886            parents: vec![
887                Some("c".to_string()),
888                Some("d".to_string()),
889                Some("e".to_string()),
890                Some("f".to_string()),
891                None,
892                None,
893                None,
894            ],
895        };
896
897        let frontier: Frontier<String, 8> = ct.to_frontier();
898        let ct0 = CommitmentTree::from_frontier(&frontier);
899        assert_eq!(ct, ct0);
900        let frontier0: Frontier<String, 8> = ct0.to_frontier();
901        assert_eq!(frontier, frontier0);
902    }
903
904    #[test]
905    fn test_random_frontier_structure() {
906        let tree_size = (2u64.pow(4)) * 3 + 5;
907
908        let mut f: Frontier<TestNode, 8> = Frontier::empty();
909        for i in 0..tree_size {
910            f.append(TestNode(i));
911        }
912        let f = f.frontier.expect("Frontier should not be empty.");
913
914        let mut rng = ChaChaRng::seed_from_u64(0);
915        let (prior_roots, f0) = Frontier::<TestNode, 8>::random_with_prior_subtree_roots(
916            &mut rng,
917            tree_size,
918            NonZeroU8::new(4).unwrap(),
919        );
920        let f0 = f0.frontier.expect("Frontier should not be empty.");
921
922        assert_eq!(prior_roots.len(), 3);
923        assert_eq!(f.position, f0.position);
924        assert_eq!(f.ommers.len(), f0.ommers.len());
925
926        let expected_largest_ommer =
927            TestNode::combine(Level::from(4), &prior_roots[0], &prior_roots[1]);
928        assert_eq!(f0.ommers[f0.ommers.len() - 1], expected_largest_ommer);
929        assert_eq!(f0.ommers[f0.ommers.len() - 2], prior_roots[2]);
930    }
931
932    #[cfg(feature = "legacy-api")]
933    proptest! {
934        #[test]
935        fn prop_commitment_tree_roundtrip(ct in arb_commitment_tree(32, arb_test_node())) {
936            let frontier: Frontier<TestNode, 8> = ct.to_frontier();
937            let ct0 = CommitmentTree::from_frontier(&frontier);
938            assert_eq!(ct, ct0);
939            let frontier0: Frontier<TestNode, 8> = ct0.to_frontier();
940            assert_eq!(frontier, frontier0);
941        }
942
943        #[test]
944        fn prop_commitment_tree_roundtrip_str(ct in arb_commitment_tree::<_, _, 8>(32, any::<char>().prop_map(|c| c.to_string()))) {
945            let frontier: Frontier<String, 8> = ct.to_frontier();
946            let ct0 = CommitmentTree::from_frontier(&frontier);
947            assert_eq!(ct, ct0);
948            let frontier0: Frontier<String, 8> = ct0.to_frontier();
949            assert_eq!(frontier, frontier0);
950        }
951    }
952}