Added proof narrowing feature

Cargo.toml

@@ -22,6 +22,7 @@ borsh = { version = "1" }
 serde_json = "1.0.96"
 postcard = { version = "1.0.4", features = ["use-std"] }
 tendermint = { version = "0.35.0" }
+paste = "1.0.15"
 
 [features]
 default = ["std"]

src/lib.rs

@@ -429,6 +429,7 @@ mod tests {
         simple_merkle::db::MemDb,
         NamespaceMerkleTree, NamespacedHash, RangeProofType, CELESTIA_NS_ID_SIZE,
     };
+    use paste::paste;
 
     type DefaultNmt<const NS_ID_SIZE: usize> = NamespaceMerkleTree<
         MemDb<NamespacedHash<NS_ID_SIZE>>,
@@ -587,6 +588,85 @@ mod tests {
         }
     }
 
+    /// 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);
+                let leaf_hashes: Vec<_> = tree.leaves()[j..i]
+                    .iter()
+                    .map(|l| l.hash().clone())
+                    .collect();
+                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 = proof
+                            .narrow_range_with_hasher(
+                                &left_hashes,
+                                &right_hashes,
+                                NamespacedSha2Hasher::with_ignore_max_ns(tree.ignore_max_ns),
+                            )
+                            .unwrap();
+                        let new_leaves: Vec<_> = tree.leaves()[l..k]
+                            .iter()
+                            .map(|l| l.hash().clone())
+                            .collect();
+                        let res = tree.check_range_proof(
+                            &root,
+                            &new_leaves,
+                            narrowed_proof.siblings(),
+                            l,
+                        );
+                        if l != k {
+                            assert!(res.is_ok());
+                            assert_eq!(res.unwrap(), RangeProofType::Complete)
+                        } else {
+                            // Cannot prove the empty range!
+                            assert!(res.is_err())
+                        }
+                    }
+                }
+                let res = tree.check_range_proof(&root, &leaf_hashes, proof.siblings(), j);
+                if i != j {
+                    assert!(res.is_ok());
+                    assert_eq!(res.unwrap(), RangeProofType::Complete)
+                } else {
+                    // Cannot prove the empty range!
+                    assert!(res.is_err())
+                }
+            }
+        }
+        test_min_and_max_ns_against(&mut tree)
+    }
+
+    macro_rules! make_narrowing_test {
+        ($($x:literal),*) => {
+            $(
+            paste! {
+                #[test]
+                fn [<test_range_proof_narrowing_ $x>]() {
+                    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);
+                }
+            })*
+        };
+    }
+
+    make_narrowing_test!(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 40);
+
     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

@@ -3,6 +3,7 @@
 //! - A range of leaves forms a complete namespace
 //! - A range of leaves all exists in the same namespace
 use crate::maybestd::{mem, vec::Vec};
+use crate::simple_merkle::db::MemDb;
 use crate::{
     namespaced_hash::{NamespaceId, NamespaceMerkleHasher, NamespacedHash},
     simple_merkle::{
@@ -98,6 +99,59 @@ where
         )
     }
 
+    /// Narrows the proof range: uses an existing proof to create
+    /// a new proof for a subrange of the original 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 new_leaf_len = left_extra_raw_leaves.len() + right_extra_raw_leaves.len();
+        if new_leaf_len >= self.range_len() {
+            return Err(RangeProofError::WrongAmountOfLeavesProvided);
+        }
+
+        // TODO: make this more concise
+        let left_extra_hashes: Vec<_> = left_extra_raw_leaves
+            .iter()
+            .map(|data| {
+                M::with_ignore_max_ns(self.ignores_max_ns())
+                    .hash_leaf_with_namespace(data.as_ref(), leaf_namespace)
+            })
+            .collect();
+        let right_extra_hashes: Vec<_> = right_extra_raw_leaves
+            .iter()
+            .map(|data| {
+                M::with_ignore_max_ns(self.ignores_max_ns())
+                    .hash_leaf_with_namespace(data.as_ref(), leaf_namespace)
+            })
+            .collect();
+
+        let mut tree = NamespaceMerkleTree::<MemDb<M::Output>, M, NS_ID_SIZE>::with_hasher(
+            M::with_ignore_max_ns(self.ignores_max_ns()),
+        );
+
+        let proof = tree.inner.narrow_range_proof(
+            &left_extra_hashes,
+            self.start_idx() as usize..(self.range_len() - new_leaf_len),
+            &right_extra_hashes,
+            &mut self.siblings(),
+            self.start_idx() as usize,
+        )?;
+
+        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,7 +1,7 @@
 use core::ops::Range;
 
 use super::{
-    db::NoopDb,
+    db::{MemDb, NoopDb},
     error::RangeProofError,
     tree::{MerkleHash, MerkleTree},
     utils::compute_num_left_siblings,
@@ -82,6 +82,36 @@ where
         )
     }
 
+    /// Narrows the proof range: uses an existing proof to create
+    /// a new proof for a subrange of the original 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)?;
+        if new_leaf_len >= self.range_len() {
+            return Err(RangeProofError::WrongAmountOfLeavesProvided);
+        }
+        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 + self.range_len() - new_leaf_len as usize; // TODO safe arithmetic
+
+        let mut tree = MerkleTree::<MemDb<M::Output>, 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