nmt_rs/
lib.rs

1#![cfg_attr(not(feature = "std"), no_std)]
2#![deny(missing_docs)]
3//! This crate implements a Namespaced Merkle Tree compatible with <https://github.com/celestiaorg/nmt>. To quote from their documentation:
4//!
5//! > A Namespaced Merkle Tree is an ordered Merkle tree that uses a modified hash function so that each node in the tree
6//! > includes the range of namespaces of the messages in all of the descendants of each node. The leafs in the tree are
7//! > ordered by the namespace identifiers of the messages. In a namespaced Merkle tree, each non-leaf node in the tree contains
8//! > the lowest and highest namespace identifiers found in all the leaf nodes that are descendants of the non-leaf node, in addition
9//! > to the hash of the concatenation of the children of the node. This enables Merkle inclusion proofs to be created that prove to
10//! > a verifier that all the elements of the tree for a specific namespace have been included in a Merkle inclusion proof.
11//! >
12//! > The concept was first introduced by [@musalbas](https://github.com/musalbas) in the [LazyLedger academic paper](https://arxiv.org/abs/1905.09274).
13//!
14//! This implementation was developed independently by [Sovereign Labs](https://www.sovereign.xyz/), and is not endorsed by the Celestia foundation.
15
16#[cfg(not(feature = "std"))]
17extern crate alloc;
18
19mod maybestd {
20    #[cfg(not(feature = "std"))]
21    pub use alloc::{boxed, vec};
22    #[cfg(all(not(feature = "std"), feature = "serde"))]
23    pub use alloc::{format, string};
24    #[cfg(not(feature = "std"))]
25    pub use core::{cmp, fmt, hash, marker, mem, ops};
26    #[cfg(feature = "std")]
27    pub use std::{boxed, cmp, fmt, hash, marker, mem, ops, vec};
28    #[cfg(all(feature = "std", feature = "serde"))]
29    pub use std::{format, string};
30
31    pub mod hash_or_btree_map {
32        #[cfg(not(feature = "std"))]
33        pub use alloc::collections::btree_map::{BTreeMap as Map, Entry};
34        #[cfg(feature = "std")]
35        pub use std::collections::hash_map::{Entry, HashMap as Map};
36    }
37}
38
39use crate::maybestd::{hash_or_btree_map, ops::Range, vec::Vec};
40
41pub use nmt_proof::NamespaceProof;
42use simple_merkle::{
43    db::{LeafWithHash, MemDb, PreimageDb},
44    error::RangeProofError,
45    proof::Proof,
46    tree::{MerkleHash, MerkleTree},
47    utils::compute_num_left_siblings,
48};
49
50mod namespaced_hash;
51pub use namespaced_hash::*;
52
53mod tendermint_hash;
54pub use tendermint_hash::*;
55
56// pub mod db;
57pub mod nmt_proof;
58pub mod simple_merkle;
59
60const CELESTIA_NS_ID_SIZE: usize = 29;
61/// A namespaced merkle tree as used in Celestia. Uses a sha256 hasher and 29 byte namespace IDs.
62pub type CelestiaNmt = NamespaceMerkleTree<
63    MemDb<NamespacedHash<CELESTIA_NS_ID_SIZE>>,
64    NamespacedSha2Hasher<CELESTIA_NS_ID_SIZE>,
65    CELESTIA_NS_ID_SIZE,
66>;
67
68/// Checks if a proof contains any partial namespaces
69fn check_proof_completeness<const NS_ID_SIZE: usize>(
70    leaves: &[NamespacedHash<NS_ID_SIZE>],
71    proof: &[NamespacedHash<NS_ID_SIZE>],
72    num_left_siblings: usize,
73) -> RangeProofType {
74    // Check if the proof is complete
75    let mut proof_type = RangeProofType::Complete;
76
77    if num_left_siblings != 0 {
78        let rightmost_left_sibling = &proof[num_left_siblings - 1];
79        if rightmost_left_sibling.max_namespace() >= leaves[0].min_namespace() {
80            proof_type = RangeProofType::Partial
81        }
82    }
83
84    let num_right_siblings = proof.len() - num_left_siblings;
85    if num_right_siblings != 0 {
86        let leftmost_right_sibling = &proof[num_left_siblings];
87        if leftmost_right_sibling.min_namespace()
88            <= leaves
89                .last()
90                .expect("leaves has already been checked to be non-empty")
91                .max_namespace()
92        {
93            proof_type = RangeProofType::Partial
94        }
95    }
96
97    proof_type
98}
99
100/// A namespaced merkle tree, implemented as a wrapper around a simple merkle tree.
101pub struct NamespaceMerkleTree<Db, M: MerkleHash, const NS_ID_SIZE: usize> {
102    namespace_ranges: hash_or_btree_map::Map<NamespaceId<NS_ID_SIZE>, Range<usize>>,
103    highest_ns: NamespaceId<NS_ID_SIZE>,
104    ignore_max_ns: bool,
105    inner: MerkleTree<Db, M>,
106}
107
108impl<Db, M, const NS_ID_SIZE: usize> NamespaceMerkleTree<Db, M, NS_ID_SIZE>
109where
110    Db: PreimageDb<M::Output>,
111    M: NamespaceMerkleHasher<NS_ID_SIZE, Output = NamespacedHash<NS_ID_SIZE>> + Default,
112{
113    /// Creates a new tree with the default hasher
114    pub fn new() -> Self {
115        Default::default()
116    }
117}
118
119impl<Db, M, const NS_ID_SIZE: usize> NamespaceMerkleTree<Db, M, NS_ID_SIZE>
120where
121    Db: PreimageDb<M::Output>,
122    M: NamespaceMerkleHasher<NS_ID_SIZE, Output = NamespacedHash<NS_ID_SIZE>>,
123{
124    /// Creates a new nmt with the provided hasher
125    pub fn with_hasher(hasher: M) -> Self {
126        Self {
127            namespace_ranges: Default::default(),
128            highest_ns: NamespaceId([0u8; NS_ID_SIZE]),
129            ignore_max_ns: hasher.ignores_max_ns(),
130            inner: MerkleTree::<Db, M>::with_hasher(hasher),
131        }
132    }
133
134    /// Adds a leaf to the namespaced merkle tree. Leaves must be pushed in namespace order.
135    pub fn push_leaf(
136        &mut self,
137        raw_data: &[u8],
138        namespace: NamespaceId<NS_ID_SIZE>,
139    ) -> Result<(), &'static str> {
140        // Force leaves to be pushed in order
141        if namespace < self.highest_ns {
142            return Err("Leaves' namespaces should be inserted in ascending order");
143        }
144        let leaf =
145            LeafWithHash::new_with_namespace(raw_data.to_vec(), namespace, self.ignore_max_ns);
146        self.highest_ns = namespace;
147        self.inner.push_leaf_with_hash(leaf);
148
149        let leaves_len = self.leaves().len();
150        match self.namespace_ranges.entry(namespace) {
151            hash_or_btree_map::Entry::Occupied(entry) => {
152                entry.into_mut().end = leaves_len;
153            }
154            hash_or_btree_map::Entry::Vacant(entry) => {
155                entry.insert(leaves_len - 1..leaves_len);
156            }
157        }
158        Ok(())
159    }
160
161    /// Returns the root of the tree, computing it if necessary. Repeated calls return a cached root.
162    pub fn root(&mut self) -> NamespacedHash<NS_ID_SIZE> {
163        self.inner.root()
164    }
165
166    /// Checks a given range proof
167    fn check_range_proof(
168        &self,
169        root: &NamespacedHash<NS_ID_SIZE>,
170        leaves: &[NamespacedHash<NS_ID_SIZE>],
171        proof: &[NamespacedHash<NS_ID_SIZE>],
172        leaves_start_idx: usize,
173    ) -> Result<RangeProofType, RangeProofError> {
174        // As an optimization, the internal call doesn't recurse into subtrees of size smaller than 2
175        // so we need to ensure that the root has size 2 or greater.
176        match leaves.len() {
177            0 => {
178                if root == &M::EMPTY_ROOT && proof.is_empty() {
179                    return Ok(RangeProofType::Complete);
180                }
181                return Err(RangeProofError::NoLeavesProvided);
182            }
183            1 => {
184                if proof.is_empty() {
185                    if &leaves[0] == root && leaves_start_idx == 0 {
186                        return Ok(RangeProofType::Complete);
187                    }
188                    return Err(RangeProofError::TreeDoesNotContainLeaf);
189                }
190            }
191            _ => {}
192        };
193
194        // Check that the leaf hashes are well-formed
195        for hash in proof.iter() {
196            if hash.min_namespace() > hash.max_namespace() {
197                return Err(RangeProofError::MalformedTree);
198            }
199        }
200
201        let num_left_siblings = compute_num_left_siblings(leaves_start_idx);
202
203        let proof_completeness = check_proof_completeness(leaves, proof, num_left_siblings);
204
205        self.inner
206            .check_range_proof(root, leaves, proof, leaves_start_idx)?;
207
208        Ok(proof_completeness)
209    }
210
211    /// Creates a range proof providing the sibling hashes required to show that a set of values really does occur in
212    /// the merkle tree at some half-open range of indices. Intermediate hashes are identified by an in-order traversal
213    /// and are returned in that same order. Panics if the range to prove is larger than the tree's leaf array.
214    ///
215    /// Example: consider the following merkle tree with leaves [C, D, E, F]
216    /// ```ascii
217    ///          root
218    ///        /      \
219    ///       A        B
220    ///      / \      /  \
221    ///     C   D    E    F
222    ///
223    /// ```
224    ///
225    /// A range proof of build_range_proof(1..3) would return the vector [C, F], since those two hashes, together
226    /// with the two leaves in the range, are sufficient to reconstruct the tree
227    pub fn build_range_proof(&mut self, leaf_range: Range<usize>) -> Proof<M> {
228        self.inner.build_range_proof(leaf_range)
229    }
230
231    /// Fetch a range of leaves from the tree, along with a proof of their inclusion.
232    pub fn get_range_with_proof(
233        &mut self,
234        leaf_range: Range<usize>,
235    ) -> (Vec<Vec<u8>>, NamespaceProof<M, NS_ID_SIZE>) {
236        let (leaves, proof) = self.inner.get_range_with_proof(leaf_range);
237        (
238            leaves,
239            NamespaceProof::PresenceProof {
240                proof,
241                ignore_max_ns: self.ignore_max_ns,
242            },
243        )
244    }
245
246    /// Get the leaf at a given index in the tree, along with a proof of its inclusion.
247    pub fn get_index_with_proof(&mut self, idx: usize) -> (Vec<u8>, Proof<M>) {
248        self.inner.get_index_with_proof(idx)
249    }
250
251    /// Get an entire namespace from the tree, along with an inclusion proof for that range.
252    pub fn get_namespace_with_proof(
253        &mut self,
254        namespace: NamespaceId<NS_ID_SIZE>,
255    ) -> (Vec<Vec<u8>>, NamespaceProof<M, NS_ID_SIZE>) {
256        let leaf_range = if let Some(range) = self.namespace_ranges.get(&namespace) {
257            range.clone()
258        } else {
259            0..0
260        };
261        let leaves = self.inner.get_leaves(leaf_range);
262
263        (leaves, self.get_namespace_proof(namespace))
264    }
265
266    /// Return all the leaves from the tree.
267    pub fn leaves(&self) -> &[LeafWithHash<M>] {
268        self.inner.leaves()
269    }
270
271    /// Get a proof for the given namespace.
272    pub fn get_namespace_proof(
273        &mut self,
274        namespace: NamespaceId<NS_ID_SIZE>,
275    ) -> NamespaceProof<M, NS_ID_SIZE> {
276        // If the namespace is outside the range covered by the root, we're done
277        if !self.root().contains::<M>(namespace) {
278            return NamespaceProof::AbsenceProof {
279                proof: Default::default(),
280                ignore_max_ns: self.ignore_max_ns,
281                leaf: None,
282            };
283        }
284
285        // If the namespace has data, just look up that namespace range and prove it by index
286        if let Some(leaf_range) = self.namespace_ranges.get(&namespace) {
287            return NamespaceProof::PresenceProof {
288                proof: self.inner.build_range_proof(leaf_range.clone()),
289                ignore_max_ns: self.ignore_max_ns,
290            };
291        }
292
293        // Otherwise, the namespace is within the range covered by the tree, but doesn't actually exist.
294        // To prove this, we can prove that for some index `i`, leaves[i-1].namespace() < namespace < leaves[i].namespace().
295        // Since a range proof for the range [i, i+1) includes the namespaced hash of the left sibling
296        // of the leaf i, which must include the namespace of the leaf i-1, then
297        // proving this range is sufficient to establish that the namespace doesn't exist.
298        let namespace = self
299            .inner
300            .leaves()
301            .binary_search_by(|l| l.hash().min_namespace().cmp(&namespace));
302
303        // The builtin binary search method returns the index where the item could be inserted while maintaining sorted order,
304        // which is the index of the leaf we want to prove
305        let idx =
306            namespace.expect_err("tree cannot contain leaf with namespace that does not exist");
307
308        let proof = self.build_range_proof(idx..idx + 1);
309
310        let mut proof = NamespaceProof::PresenceProof {
311            proof,
312            ignore_max_ns: self.ignore_max_ns,
313        };
314        proof.convert_to_absence_proof(self.inner.leaves()[idx].hash().clone());
315        proof
316    }
317
318    fn verify_namespace(
319        &self,
320        root: &NamespacedHash<NS_ID_SIZE>,
321        raw_leaves: &[impl AsRef<[u8]>],
322        namespace: NamespaceId<NS_ID_SIZE>,
323        proof: &NamespaceProof<M, NS_ID_SIZE>,
324    ) -> Result<(), RangeProofError> {
325        if root.is_empty_root::<M>() && raw_leaves.is_empty() {
326            return Ok(());
327        }
328
329        match proof {
330            NamespaceProof::AbsenceProof { leaf, .. } => {
331                if !root.contains::<M>(namespace) {
332                    return Ok(());
333                }
334                let leaf = leaf.clone().ok_or(RangeProofError::MalformedProof(
335                    "Absence proof was inside tree range but did not contain a leaf",
336                ))?;
337                // Check that they haven't provided an absence proof for a non-empty namespace
338                if !raw_leaves.is_empty() {
339                    return Err(RangeProofError::MalformedProof(
340                        "provided an absence proof for a non-empty namespace",
341                    ));
342                }
343                // Check that the provided namespace actually precedes the leaf
344                if namespace >= leaf.min_namespace() {
345                    return Err(RangeProofError::MalformedProof(
346                        "provided leaf must have namespace greater than the namespace which is being proven absent",
347                    ));
348                }
349                let num_left_siblings = compute_num_left_siblings(proof.start_idx() as usize);
350
351                // Check that the namespace actually follows the closest sibling
352                let siblings = proof.siblings();
353                if num_left_siblings > 0 {
354                    let rightmost_left_sibling = &siblings[num_left_siblings - 1];
355                    if rightmost_left_sibling.max_namespace() >= namespace {
356                        return Err(RangeProofError::MalformedProof("proven namespace must be greater than the namespace of the rightmost left sibling"));
357                    }
358                }
359                // Then, check that the root is real
360                self.inner.check_range_proof(
361                    root,
362                    &[leaf],
363                    proof.siblings(),
364                    proof.start_idx() as usize,
365                )?;
366            }
367            NamespaceProof::PresenceProof { ignore_max_ns, .. } => {
368                if !root.contains::<M>(namespace) {
369                    return Err(RangeProofError::TreeDoesNotContainLeaf);
370                }
371                let leaf_hashes: Vec<NamespacedHash<NS_ID_SIZE>> = raw_leaves
372                    .iter()
373                    .map(|data| {
374                        M::with_ignore_max_ns(*ignore_max_ns)
375                            .hash_leaf_with_namespace(data.as_ref(), namespace)
376                    })
377                    .collect();
378                let proof_type = self.check_range_proof(
379                    root,
380                    &leaf_hashes,
381                    proof.siblings(),
382                    proof.start_idx() as usize,
383                )?;
384                if proof_type == RangeProofType::Partial {
385                    return Err(RangeProofError::MissingLeaf);
386                }
387            }
388        }
389        Ok(())
390    }
391}
392
393impl<Db, M, const NS_ID_SIZE: usize> Default for NamespaceMerkleTree<Db, M, NS_ID_SIZE>
394where
395    Db: PreimageDb<M::Output>,
396    M: MerkleHash + Default,
397{
398    fn default() -> Self {
399        Self {
400            namespace_ranges: Default::default(),
401            highest_ns: NamespaceId::default(),
402            ignore_max_ns: true,
403            inner: Default::default(),
404        }
405    }
406}
407
408/// Indicates whether the proof includes all leaves from every namespace it covers.
409#[derive(Debug, PartialEq, Clone, Copy)]
410pub enum RangeProofType {
411    /// A range proof over a single namespace is complete if it includes all the leaves
412    /// in that namespace. A range proof over several namespaces is complete if all
413    /// individual namespaces are complete.
414    Complete,
415    /// A range proof over a single namespace is partial if it omits at least one leaf from that namespace.
416    /// A range proof over several namespaces is partial if it includes at least one namespace that is partial.
417    ///
418    /// Note that (since ranges are contiguous) only the first or last namespace covered by a range proof may be partial.
419    Partial,
420}
421
422#[cfg(test)]
423mod tests {
424    use crate::maybestd::vec::Vec;
425    use crate::simple_merkle::error::RangeProofError;
426    use crate::NamespaceMerkleHasher;
427    use crate::{
428        namespaced_hash::{NamespaceId, NamespacedSha2Hasher},
429        nmt_proof::NamespaceProof,
430        simple_merkle::db::MemDb,
431        NamespaceMerkleTree, NamespacedHash, RangeProofType, CELESTIA_NS_ID_SIZE,
432    };
433
434    type DefaultNmt<const NS_ID_SIZE: usize> = NamespaceMerkleTree<
435        MemDb<NamespacedHash<NS_ID_SIZE>>,
436        NamespacedSha2Hasher<NS_ID_SIZE>,
437        NS_ID_SIZE,
438    >;
439
440    fn ns_id_from_u64<const NS_ID_SIZE: usize>(val: u64) -> NamespaceId<NS_ID_SIZE> {
441        // make sure the namespace id can hold the provided value
442        assert!(NS_ID_SIZE >= 8);
443        let mut namespace = NamespaceId::default();
444        namespace.0[NS_ID_SIZE - 8..].copy_from_slice(&val.to_be_bytes());
445        namespace
446    }
447
448    /// Builds a tree from provided namespace ids
449    fn tree_from_namespace_ids<const NS_ID_SIZE: usize>(
450        ns_ids: impl AsRef<[u64]>,
451    ) -> DefaultNmt<NS_ID_SIZE> {
452        let mut tree = DefaultNmt::new();
453        for (i, &ns_id) in ns_ids.as_ref().iter().enumerate() {
454            let data = format!("leaf_{i}");
455            let namespace = ns_id_from_u64(ns_id);
456            tree.push_leaf(data.as_bytes(), namespace)
457                .expect("Failed to push the leaf");
458        }
459        tree
460    }
461
462    fn tree_from_one_namespace<const NS_ID_SIZE: usize>(
463        leaves: u64,
464        namespace: u64,
465    ) -> DefaultNmt<NS_ID_SIZE> {
466        let mut tree = DefaultNmt::new();
467        let namespace = ns_id_from_u64(namespace);
468        for i in 0..leaves {
469            let data = format!("leaf_{i}");
470            tree.push_leaf(data.as_bytes(), namespace)
471                .expect("Failed to push the leaf");
472        }
473        tree
474    }
475
476    /// Builds a tree with N leaves
477    fn tree_with_n_leaves<const NS_ID_SIZE: usize>(n: usize) -> DefaultNmt<NS_ID_SIZE> {
478        tree_from_namespace_ids((0..n as u64).collect::<Vec<_>>())
479    }
480
481    #[test]
482    fn test_absence_proof_leaf_advances_the_namespace() {
483        let mut tree = tree_from_namespace_ids::<8>(&[1, 2, 3, 4, 6, 7, 8, 9]);
484        let namespace = ns_id_from_u64(5);
485        let proof = tree.get_namespace_proof(namespace);
486        let no_leaves: &[&[u8]] = &[];
487
488        proof
489            .verify_complete_namespace(&tree.root(), no_leaves, namespace)
490            .unwrap();
491
492        let NamespaceProof::AbsenceProof {
493            leaf: Some(leaf), ..
494        } = proof
495        else {
496            unreachable!();
497        };
498
499        // https://github.com/celestiaorg/nmt/blob/master/docs/spec/nmt.md#verification-of-nmt-absence-proof
500        assert!(leaf.min_namespace() > namespace);
501    }
502
503    #[test]
504    fn test_absence_proof_return_err_if_leaf_doesnt_follow_rightmost_left_sibling() {
505        let mut tree = tree_from_namespace_ids::<8>(&[1, 2, 3, 4, 6, 7, 8, 9]);
506        let namespace = ns_id_from_u64(5);
507        let proof = tree.get_namespace_proof(namespace);
508        let no_leaves: &[&[u8]] = &[];
509
510        for i in [3, 4, 6, 7] {
511            let mut proof = proof.clone();
512            let NamespaceProof::AbsenceProof { leaf, .. } = &mut proof else {
513                unreachable!();
514            };
515            let data = format!("leaf_{i}").as_bytes().to_vec();
516            *leaf = Some(
517                NamespacedSha2Hasher::default().hash_leaf_with_namespace(&data, ns_id_from_u64(i)),
518            );
519            proof
520                .verify_complete_namespace(&tree.root(), no_leaves, ns_id_from_u64(2))
521                .unwrap_err();
522        }
523    }
524
525    #[test]
526    fn test_absence_proof_doesnt_include_leaf_if_namespace_is_out_of_root_ns_range() {
527        let mut tree = tree_from_namespace_ids::<8>(&[2, 3, 4, 5]);
528        for namespace in [1, 6] {
529            let namespace = ns_id_from_u64(namespace);
530            let proof = tree.get_namespace_proof(namespace);
531
532            proof
533                .clone()
534                .verify_complete_namespace(&tree.root(), &Vec::<Vec<u8>>::new(), namespace)
535                .unwrap();
536
537            assert!(matches!(
538                proof,
539                NamespaceProof::AbsenceProof { leaf: None, .. }
540            ));
541        }
542    }
543
544    #[test]
545    fn test_wrong_amount_of_leaves() {
546        let mut tree = tree_from_namespace_ids::<8>(&[1, 2, 2, 2, 3, 4, 5, 6]);
547        let namespace = ns_id_from_u64(2);
548        let proof = tree.get_namespace_proof(namespace);
549
550        let leaves = [b"leaf_1", b"leaf_2", b"leaf_3", b"leaf_4"];
551
552        for leaves in [&leaves[..], &leaves[..2]] {
553            proof
554                .verify_complete_namespace(&tree.root(), leaves, namespace)
555                .unwrap_err();
556            proof
557                .verify_range(&tree.root(), leaves, namespace)
558                .unwrap_err();
559        }
560
561        proof
562            .verify_complete_namespace(&tree.root(), &leaves[..3], namespace)
563            .unwrap();
564        proof
565            .verify_range(&tree.root(), &leaves[..3], namespace)
566            .unwrap();
567    }
568
569    /// Builds a tree with n leaves, and then creates and checks proofs of all
570    /// valid ranges.
571    fn test_range_proof_roundtrip_with_n_leaves<const NS_ID_SIZE: usize>(n: usize) {
572        let mut tree = tree_with_n_leaves::<NS_ID_SIZE>(n);
573        let root = tree.root();
574        for i in 1..=n {
575            for j in 0..=i {
576                let proof = tree.build_range_proof(j..i);
577                let leaf_hashes: Vec<_> = tree.leaves()[j..i]
578                    .iter()
579                    .map(|l| l.hash().clone())
580                    .collect();
581                let res = tree.check_range_proof(&root, &leaf_hashes, proof.siblings(), j);
582                if i != j {
583                    assert!(res.is_ok());
584                    assert_eq!(res.unwrap(), RangeProofType::Complete)
585                } else {
586                    // Cannot prove the empty range!
587                    assert!(res.is_err())
588                }
589            }
590        }
591        test_min_and_max_ns_against(&mut tree)
592    }
593
594    #[test]
595    fn test_range_proof_roundtrip() {
596        for x in 0..20 {
597            test_range_proof_roundtrip_with_n_leaves::<8>(x);
598            test_range_proof_roundtrip_with_n_leaves::<17>(x);
599            test_range_proof_roundtrip_with_n_leaves::<24>(x);
600            test_range_proof_roundtrip_with_n_leaves::<CELESTIA_NS_ID_SIZE>(x);
601            test_range_proof_roundtrip_with_n_leaves::<32>(x);
602        }
603    }
604
605    fn test_range_proof_narrowing_within_namespace<const NS_ID_SIZE: usize>(n: usize) {
606        let ns_id = 4;
607        let mut tree = tree_from_one_namespace::<NS_ID_SIZE>(n as u64, ns_id); // since there's a single namespace, the actual ID shouldn't matter
608        let root = tree.root();
609        for i in 1..=n {
610            for j in 0..=i {
611                let proof_nmt = NamespaceProof::PresenceProof {
612                    proof: tree.build_range_proof(j..i),
613                    ignore_max_ns: tree.ignore_max_ns,
614                };
615                for k in (j + 1)..=i {
616                    for l in j..=k {
617                        let left_leaf_datas: Vec<_> =
618                            tree.leaves()[j..l].iter().map(|l| l.data()).collect();
619                        let right_leaf_datas: Vec<_> =
620                            tree.leaves()[k..i].iter().map(|l| l.data()).collect();
621                        let narrowed_proof_nmt = proof_nmt.narrow_range(
622                            &left_leaf_datas,
623                            &right_leaf_datas,
624                            ns_id_from_u64(ns_id),
625                        );
626                        if k == l {
627                            // Cannot prove the empty range!
628                            assert!(narrowed_proof_nmt.is_err());
629                            assert_eq!(
630                                narrowed_proof_nmt.unwrap_err(),
631                                RangeProofError::NoLeavesProvided
632                            );
633                            continue;
634                        } else {
635                            assert!(narrowed_proof_nmt.is_ok());
636                        }
637                        let narrowed_proof = narrowed_proof_nmt.unwrap();
638                        let new_leaves: Vec<_> = tree.leaves()[l..k]
639                            .iter()
640                            .map(|l| l.hash().clone())
641                            .collect();
642                        tree.check_range_proof(&root, &new_leaves, narrowed_proof.siblings(), l)
643                            .unwrap();
644                    }
645                }
646            }
647        }
648        test_min_and_max_ns_against(&mut tree)
649    }
650
651    #[test]
652    fn test_range_proof_narrowing_nmt() {
653        for x in 0..20 {
654            test_range_proof_narrowing_within_namespace::<8>(x);
655            test_range_proof_narrowing_within_namespace::<17>(x);
656            test_range_proof_narrowing_within_namespace::<24>(x);
657            test_range_proof_narrowing_within_namespace::<CELESTIA_NS_ID_SIZE>(x);
658            test_range_proof_narrowing_within_namespace::<32>(x);
659        }
660    }
661
662    /// Builds a tree with n leaves, and then creates and checks proofs of all valid
663    /// ranges, and attempts to narrow every proof and re-check it for the narrowed range
664    fn test_range_proof_narrowing_with_n_leaves<const NS_ID_SIZE: usize>(n: usize) {
665        let mut tree = tree_with_n_leaves::<NS_ID_SIZE>(n);
666        let root = tree.root();
667        for i in 1..=n {
668            for j in 0..=i {
669                let proof = tree.build_range_proof(j..i);
670                for k in (j + 1)..=i {
671                    for l in j..=k {
672                        let left_hashes: Vec<_> = tree.leaves()[j..l]
673                            .iter()
674                            .map(|l| l.hash().clone())
675                            .collect();
676                        let right_hashes: Vec<_> = tree.leaves()[k..i]
677                            .iter()
678                            .map(|l| l.hash().clone())
679                            .collect();
680                        let narrowed_proof_simple = proof.narrow_range_with_hasher(
681                            &left_hashes,
682                            &right_hashes,
683                            NamespacedSha2Hasher::with_ignore_max_ns(tree.ignore_max_ns),
684                        );
685                        if k == l {
686                            // Cannot prove the empty range!
687                            assert!(narrowed_proof_simple.is_err());
688                            assert_eq!(
689                                narrowed_proof_simple.unwrap_err(),
690                                RangeProofError::NoLeavesProvided
691                            );
692                            continue;
693                        } else {
694                            assert!(narrowed_proof_simple.is_ok());
695                        }
696                        let narrowed_proof = narrowed_proof_simple.unwrap();
697                        let new_leaves: Vec<_> = tree.leaves()[l..k]
698                            .iter()
699                            .map(|l| l.hash().clone())
700                            .collect();
701                        tree.check_range_proof(&root, &new_leaves, narrowed_proof.siblings(), l)
702                            .unwrap();
703                    }
704                }
705            }
706        }
707        test_min_and_max_ns_against(&mut tree)
708    }
709
710    #[test]
711    fn test_range_proof_narrowing_simple() {
712        for x in 0..20 {
713            test_range_proof_narrowing_with_n_leaves::<8>(x);
714            test_range_proof_narrowing_with_n_leaves::<17>(x);
715            test_range_proof_narrowing_with_n_leaves::<24>(x);
716            test_range_proof_narrowing_with_n_leaves::<CELESTIA_NS_ID_SIZE>(x);
717            test_range_proof_narrowing_with_n_leaves::<32>(x);
718        }
719    }
720
721    fn test_completeness_check_impl<const NS_ID_SIZE: usize>() {
722        // Build a tree with 32 leaves spread evenly across 8 namespaces
723        let mut tree = DefaultNmt::<NS_ID_SIZE>::new();
724        for x in 0..32 {
725            let namespace = ns_id_from_u64(x / 4);
726            let _ = tree.push_leaf(x.to_be_bytes().as_ref(), namespace);
727        }
728        let root = tree.root();
729        let leaf_hashes: Vec<_> = tree.leaves().iter().map(|x| x.hash().clone()).collect();
730
731        // For each potential range of size four, build and check a range proof
732        for i in 0..=28 {
733            let leaf_range = i..i + 4;
734            let proof = tree.build_range_proof(leaf_range.clone());
735
736            let result =
737                tree.check_range_proof(&root, &leaf_hashes[leaf_range], proof.siblings(), i);
738            assert!(result.is_ok());
739
740            // We've set up our tree to have four leaves in each namespace, so a
741            // range of leaves covers a complete namespace only if and only if the start index
742            // is divisible by four
743            let should_be_complete = (i % 4) == 0;
744            if should_be_complete {
745                assert_eq!(result, Ok(RangeProofType::Complete))
746            } else {
747                assert_eq!(result, Ok(RangeProofType::Partial))
748            }
749        }
750        for nid in 0..100u64 {
751            let namespace = ns_id_from_u64(nid);
752            let (leaves, proof) = tree.get_namespace_with_proof(namespace);
753
754            let pf = proof.verify_complete_namespace(&root, &leaves, namespace);
755            assert!(pf.is_ok());
756        }
757    }
758
759    #[test]
760    fn test_completeness_check() {
761        test_completeness_check_impl::<8>();
762        test_completeness_check_impl::<17>();
763        test_completeness_check_impl::<24>();
764        test_completeness_check_impl::<CELESTIA_NS_ID_SIZE>();
765        test_completeness_check_impl::<32>();
766    }
767
768    // Try building and checking a proof of the min namespace, and the max namespace.
769    // Then, add a node to the max namespace and check the max again.
770    fn test_min_and_max_ns_against<const NS_ID_SIZE: usize>(tree: &mut DefaultNmt<NS_ID_SIZE>) {
771        let root = tree.root();
772        let min_namespace = NamespaceId([0; NS_ID_SIZE]);
773        let max_namespace = NamespaceId([0xff; NS_ID_SIZE]);
774        let (leaves, proof) = tree.get_namespace_with_proof(min_namespace);
775        assert!(proof
776            .verify_complete_namespace(&root, &leaves, min_namespace)
777            .is_ok());
778
779        let (leaves, proof) = tree.get_namespace_with_proof(max_namespace);
780        assert!(proof
781            .verify_complete_namespace(&root, &leaves, max_namespace)
782            .is_ok());
783
784        tree.push_leaf(b"some_leaf", max_namespace)
785            .expect("can always push max namespace");
786
787        let root = tree.root();
788        let (leaves, proof) = tree.get_namespace_with_proof(max_namespace);
789        assert!(proof
790            .verify_complete_namespace(&root, &leaves, max_namespace)
791            .is_ok());
792    }
793
794    fn test_namespace_verification_impl<const NS_ID_SIZE: usize>() {
795        let mut tree = DefaultNmt::<NS_ID_SIZE>::new();
796        // Put a bunch of data in the tree
797        for x in 0..33 {
798            // Ensure that some namespaces are skipped, including the zero namespace
799            let namespace = ns_id_from_u64(((x / 5) * 3) + 1);
800            let _ = tree.push_leaf(x.to_be_bytes().as_ref(), namespace);
801        }
802        let root = tree.root();
803        let raw_leaves: Vec<Vec<u8>> = tree.leaves().iter().map(|x| x.data().to_vec()).collect();
804
805        // Build proofs for each range that's actually included, and check that the range can be retrieved correctly
806        for (namespace, range) in tree.namespace_ranges.clone().iter() {
807            let proof = tree.build_range_proof(range.clone());
808            assert!(!range.is_empty());
809
810            let proof = NamespaceProof::PresenceProof {
811                proof,
812                ignore_max_ns: true,
813            };
814
815            assert!(tree
816                .verify_namespace(&root, &raw_leaves[range.clone()], *namespace, &proof)
817                .is_ok());
818        }
819
820        // Build and check proofs for a bunch of namespaces, including some that are present and some that are absent.
821        for nid in 0..100u64 {
822            let namespace = ns_id_from_u64(nid);
823            let (leaves, proof) = tree.get_namespace_with_proof(namespace);
824            let pf = proof.verify_complete_namespace(&root, &leaves, namespace);
825            assert!(pf.is_ok());
826        }
827
828        test_min_and_max_ns_against(&mut tree)
829    }
830
831    #[test]
832    fn test_namespace_verification() {
833        test_namespace_verification_impl::<8>();
834        test_namespace_verification_impl::<17>();
835        test_namespace_verification_impl::<24>();
836        test_namespace_verification_impl::<CELESTIA_NS_ID_SIZE>();
837        test_namespace_verification_impl::<32>();
838    }
839
840    #[allow(unused)]
841    fn compilation_test_nmt_is_send() {
842        fn is_send<T: Send>(_t: T) {}
843
844        is_send(DefaultNmt::<1>::new());
845    }
846
847    #[allow(unused)]
848    fn compilation_test_nmt_is_sync() {
849        fn is_sync<T: Sync>(_t: T) {}
850
851        is_sync(DefaultNmt::<1>::new());
852    }
853}
nmt_rs/lib.rs

nmt_rs/
lib.rs
