Theodore/proofnarrowing (#34)

* Added proof narrowing feature

Given a proof for a range of leaves [n..m] (and no other information about the tree), allows supplying sub-ranges contiguous with each end of the original range - i.e. [n..k], [l..m], where k < l - to generate a new proof for the sub-range [k..l].

This is particularly useful in Celestia, for generating proofs of specific share ranges within a namespace (e.g. individual proofs for each blob, when there are multiple blobs in a namespace) using a namespace range proof.

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "nmt-rs"
-version = "0.2.1"
+version = "0.2.3"
 edition = "2021"
 description = "A namespaced merkle tree compatible with Celestia"
 license = "MIT OR Apache-2.0"

README.md

@@ -20,6 +20,7 @@ This code has not been audited, and may contain critical vulnerabilities. Do not
 
 - [x] Verify namespace range proofs
 
+- [x] Narrow namespace range proofs: supply part of the range to generate a proof for the remaining sub-range
 
 ## License
 

src/lib.rs

@@ -422,6 +422,7 @@ pub enum RangeProofType {
 #[cfg(test)]
 mod tests {
     use crate::maybestd::vec::Vec;
+    use crate::simple_merkle::error::RangeProofError;
     use crate::NamespaceMerkleHasher;
     use crate::{
         namespaced_hash::{NamespaceId, NamespacedSha2Hasher},
@@ -458,6 +459,20 @@ mod tests {
         tree
     }
 
+    fn tree_from_one_namespace<const NS_ID_SIZE: usize>(
+        leaves: u64,
+        namespace: u64,
+    ) -> DefaultNmt<NS_ID_SIZE> {
+        let mut tree = DefaultNmt::new();
+        let namespace = ns_id_from_u64(namespace);
+        for i in 0..leaves {
+            let data = format!("leaf_{i}");
+            tree.push_leaf(data.as_bytes(), namespace)
+                .expect("Failed to push the leaf");
+        }
+        tree
+    }
+
     /// Builds a tree with N leaves
     fn tree_with_n_leaves<const NS_ID_SIZE: usize>(n: usize) -> DefaultNmt<NS_ID_SIZE> {
         tree_from_namespace_ids((0..n as u64).collect::<Vec<_>>())
@@ -587,6 +602,122 @@ mod tests {
         }
     }
 
+    fn test_range_proof_narrowing_within_namespace<const NS_ID_SIZE: usize>(n: usize) {
+        let ns_id = 4;
+        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
+        let root = tree.root();
+        for i in 1..=n {
+            for j in 0..=i {
+                let proof_nmt = NamespaceProof::PresenceProof {
+                    proof: tree.build_range_proof(j..i),
+                    ignore_max_ns: tree.ignore_max_ns,
+                };
+                for k in (j + 1)..=i {
+                    for l in j..=k {
+                        let left_leaf_datas: Vec<_> =
+                            tree.leaves()[j..l].iter().map(|l| l.data()).collect();
+                        let right_leaf_datas: Vec<_> =
+                            tree.leaves()[k..i].iter().map(|l| l.data()).collect();
+                        let narrowed_proof_nmt = proof_nmt.narrow_range(
+                            &left_leaf_datas,
+                            &right_leaf_datas,
+                            ns_id_from_u64(ns_id),
+                        );
+                        if k == l {
+                            // Cannot prove the empty range!
+                            assert!(narrowed_proof_nmt.is_err());
+                            assert_eq!(
+                                narrowed_proof_nmt.unwrap_err(),
+                                RangeProofError::NoLeavesProvided
+                            );
+                            continue;
+                        } else {
+                            assert!(narrowed_proof_nmt.is_ok());
+                        }
+                        let narrowed_proof = narrowed_proof_nmt.unwrap();
+                        let new_leaves: Vec<_> = tree.leaves()[l..k]
+                            .iter()
+                            .map(|l| l.hash().clone())
+                            .collect();
+                        tree.check_range_proof(&root, &new_leaves, narrowed_proof.siblings(), l)
+                            .unwrap();
+                    }
+                }
+            }
+        }
+        test_min_and_max_ns_against(&mut tree)
+    }
+
+    #[test]
+    fn test_range_proof_narrowing_nmt() {
+        for x in 0..20 {
+            test_range_proof_narrowing_within_namespace::<8>(x);
+            test_range_proof_narrowing_within_namespace::<17>(x);
+            test_range_proof_narrowing_within_namespace::<24>(x);
+            test_range_proof_narrowing_within_namespace::<CELESTIA_NS_ID_SIZE>(x);
+            test_range_proof_narrowing_within_namespace::<32>(x);
+        }
+    }
+
+    /// Builds a tree with n leaves, and then creates and checks proofs of all valid
+    /// ranges, and attempts to narrow every proof and re-check it for the narrowed range
+    fn test_range_proof_narrowing_with_n_leaves<const NS_ID_SIZE: usize>(n: usize) {
+        let mut tree = tree_with_n_leaves::<NS_ID_SIZE>(n);
+        let root = tree.root();
+        for i in 1..=n {
+            for j in 0..=i {
+                let proof = tree.build_range_proof(j..i);
+                for k in (j + 1)..=i {
+                    for l in j..=k {
+                        let left_hashes: Vec<_> = tree.leaves()[j..l]
+                            .iter()
+                            .map(|l| l.hash().clone())
+                            .collect();
+                        let right_hashes: Vec<_> = tree.leaves()[k..i]
+                            .iter()
+                            .map(|l| l.hash().clone())
+                            .collect();
+                        let narrowed_proof_simple = proof.narrow_range_with_hasher(
+                            &left_hashes,
+                            &right_hashes,
+                            NamespacedSha2Hasher::with_ignore_max_ns(tree.ignore_max_ns),
+                        );
+                        if k == l {
+                            // Cannot prove the empty range!
+                            assert!(narrowed_proof_simple.is_err());
+                            assert_eq!(
+                                narrowed_proof_simple.unwrap_err(),
+                                RangeProofError::NoLeavesProvided
+                            );
+                            continue;
+                        } else {
+                            assert!(narrowed_proof_simple.is_ok());
+                        }
+                        let narrowed_proof = narrowed_proof_simple.unwrap();
+                        let new_leaves: Vec<_> = tree.leaves()[l..k]
+                            .iter()
+                            .map(|l| l.hash().clone())
+                            .collect();
+                        tree.check_range_proof(&root, &new_leaves, narrowed_proof.siblings(), l)
+                            .unwrap();
+                    }
+                }
+            }
+        }
+        test_min_and_max_ns_against(&mut tree)
+    }
+
+    #[test]
+    fn test_range_proof_narrowing_simple() {
+        for x in 0..20 {
+            test_range_proof_narrowing_with_n_leaves::<8>(x);
+            test_range_proof_narrowing_with_n_leaves::<17>(x);
+            test_range_proof_narrowing_with_n_leaves::<24>(x);
+            test_range_proof_narrowing_with_n_leaves::<CELESTIA_NS_ID_SIZE>(x);
+            test_range_proof_narrowing_with_n_leaves::<32>(x);
+        }
+    }
+
     fn test_completeness_check_impl<const NS_ID_SIZE: usize>() {
         // Build a tree with 32 leaves spread evenly across 8 namespaces
         let mut tree = DefaultNmt::<NS_ID_SIZE>::new();

src/nmt_proof.rs

@@ -98,6 +98,50 @@ where
         )
     }
 
+    /// Narrows the proof range: uses an existing proof to create
+    /// a new proof for a subrange of the original proof's range
+    ///
+    /// # Arguments
+    ///  - left_extra_raw_leaves: The data for the leaves that will narrow the range from the left
+    ///    side (i.e. all the leaves from the left edge of the currently proven range, to the left
+    ///    edge of the new desired shrunk range)
+    ///  - right_extra_raw_leaves: Analogously, data for all the leaves between the right edge of
+    ///    the desired shrunken range, and the right edge of the current proof's range
+    pub fn narrow_range<L: AsRef<[u8]>>(
+        &self,
+        left_extra_raw_leaves: &[L],
+        right_extra_raw_leaves: &[L],
+        leaf_namespace: NamespaceId<NS_ID_SIZE>,
+    ) -> Result<Self, RangeProofError> {
+        if self.is_of_absence() {
+            return Err(RangeProofError::MalformedProof(
+                "Cannot narrow the range of an absence proof",
+            ));
+        }
+
+        let leaves_to_hashes = |l: &[L]| -> Vec<NamespacedHash<NS_ID_SIZE>> {
+            l.iter()
+                .map(|data| {
+                    M::with_ignore_max_ns(self.ignores_max_ns())
+                        .hash_leaf_with_namespace(data.as_ref(), leaf_namespace)
+                })
+                .collect()
+        };
+        let left_extra_hashes = leaves_to_hashes(left_extra_raw_leaves);
+        let right_extra_hashes = leaves_to_hashes(right_extra_raw_leaves);
+
+        let proof = self.merkle_proof().narrow_range_with_hasher(
+            &left_extra_hashes,
+            &right_extra_hashes,
+            M::with_ignore_max_ns(self.ignores_max_ns()),
+        )?;
+
+        Ok(Self::PresenceProof {
+            proof,
+            ignore_max_ns: self.ignores_max_ns(),
+        })
+    }
+
     /// Convert a proof of the presence of some leaf to the proof of the absence of another leaf
     pub fn convert_to_absence_proof(&mut self, leaf: NamespacedHash<NS_ID_SIZE>) {
         match self {

src/simple_merkle/proof.rs

@@ -1,4 +1,4 @@
-use core::ops::Range;
+use core::{cmp::Ordering, ops::Range};
 
 use super::{
     db::NoopDb,
@@ -82,6 +82,53 @@ where
         )
     }
 
+    /// Narrows the proof range: uses an existing proof to create
+    /// a new proof for a subrange of the original proof's range
+    ///
+    /// # Arguments
+    ///  - left_extra_leaves: The hashes of the leaves that will narrow the range from the left
+    ///    side (i.e. all the leaves from the left edge of the currently proven range, to the left
+    ///    edge of the new desired shrunk range)
+    ///  - right_extra_leaves: Analogously, hashes of all the leaves between the right edge of
+    ///    the desired shrunken range, and the right edge of the current proof's range
+    pub fn narrow_range_with_hasher(
+        &self,
+        left_extra_leaves: &[M::Output],
+        right_extra_leaves: &[M::Output],
+        hasher: M,
+    ) -> Result<Self, RangeProofError> {
+        let new_leaf_len = left_extra_leaves
+            .len()
+            .checked_add(right_extra_leaves.len())
+            .ok_or(RangeProofError::TreeTooLarge)?;
+        match new_leaf_len.cmp(&self.range_len()) {
+            Ordering::Equal => {
+                // We cannot prove the empty range!
+                return Err(RangeProofError::NoLeavesProvided);
+            }
+            Ordering::Greater => return Err(RangeProofError::WrongAmountOfLeavesProvided),
+            Ordering::Less => { /* Ok! */ }
+        }
+
+        // Indices relative to the leaves of the entire tree
+        let new_start_idx = (self.start_idx() as usize)
+            .checked_add(left_extra_leaves.len())
+            .ok_or(RangeProofError::TreeTooLarge)?;
+        let new_end_idx = new_start_idx
+            .checked_add(self.range_len())
+            .and_then(|i| i.checked_sub(new_leaf_len))
+            .ok_or(RangeProofError::TreeTooLarge)?;
+
+        let mut tree = MerkleTree::<NoopDb, M>::with_hasher(hasher);
+        tree.narrow_range_proof(
+            left_extra_leaves,
+            new_start_idx..new_end_idx,
+            right_extra_leaves,
+            &mut self.siblings().as_slice(),
+            self.start_idx() as usize,
+        )
+    }
+
     /// Returns the siblings provided as part of the proof.
     pub fn siblings(&self) -> &Vec<M::Output> {
         &self.siblings