feat: add necessary trait bounds to balanced merkle tree

base_layer/mmr/Cargo.toml

@@ -17,6 +17,7 @@ tari_utilities = { git = "https://github.com/tari-project/tari_utilities.git", t
 tari_crypto = { git = "https://github.com/tari-project/tari-crypto.git", tag = "v0.16.8" }
 tari_common = {path = "../../common"}
 thiserror = "1.0.26"
+borsh = "0.9.3"
 digest = "0.9.0"
 log = "0.4"
 serde = { version = "1.0.97", features = ["derive"] }

base_layer/mmr/src/balanced_binary_merkle_proof.rs

@@ -20,19 +20,25 @@
 // WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
 // USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-use std::{collections::HashMap, convert::TryInto, marker::PhantomData};
+use std::{
+    collections::HashMap,
+    convert::{TryFrom, TryInto},
+    marker::PhantomData,
+};
 
+use borsh::{BorshDeserialize, BorshSerialize};
 use digest::Digest;
+use serde::{Deserialize, Serialize};
 use tari_common::DomainDigest;
 use tari_utilities::ByteArray;
 use thiserror::Error;
 
 use crate::{common::hash_together, BalancedBinaryMerkleTree, Hash};
 
-#[derive(Debug)]
+#[derive(BorshDeserialize, BorshSerialize, Deserialize, Serialize, Clone, Debug, Default, PartialEq, Eq)]
 pub struct BalancedBinaryMerkleProof<D> {
     pub path: Vec<Hash>,
-    pub node_index: usize,
+    pub node_index: u32,
     _phantom: PhantomData<D>,
 }
 
@@ -55,7 +61,10 @@ where D: Digest + DomainDigest
         &computed_root == root
     }
 
-    pub fn generate_proof(tree: &BalancedBinaryMerkleTree<D>, leaf_index: usize) -> Self {
+    pub fn generate_proof(
+        tree: &BalancedBinaryMerkleTree<D>,
+        leaf_index: usize,
+    ) -> Result<Self, MergedBalancedBinaryMerkleProofError> {
         let mut node_index = tree.get_node_index(leaf_index);
         let mut proof = Vec::new();
         while node_index > 0 {
@@ -67,11 +76,12 @@ where D: Digest + DomainDigest
             // Traverse to parent
             node_index = parent;
         }
-        Self {
+        Ok(Self {
             path: proof,
-            node_index: tree.get_node_index(leaf_index),
+            node_index: u32::try_from(tree.get_node_index(leaf_index))
+                .map_err(|_| MergedBalancedBinaryMerkleProofError::MathOverflow)?,
             _phantom: PhantomData,
-        }
+        })
     }
 }
 
@@ -81,6 +91,8 @@ pub enum MergedBalancedBinaryMerkleProofError {
     CantMergeZeroProofs,
     #[error("Bad proof semantics")]
     BadProofSemantics,
+    #[error("Math overflow")]
+    MathOverflow,
 }
 
 /// Flag to indicate if proof data represents an index or a node hash
@@ -94,8 +106,8 @@ pub enum MergedBalancedBinaryMerkleDataType {
 #[derive(Debug)]
 pub struct MergedBalancedBinaryMerkleProof<D> {
     pub paths: Vec<Vec<(MergedBalancedBinaryMerkleDataType, Vec<u8>)>>, // these tuples can contain indexes or hashes!
-    pub node_indices: Vec<usize>,
-    pub heights: Vec<usize>,
+    pub node_indices: Vec<u32>,
+    pub heights: Vec<u32>,
     _phantom: PhantomData<D>,
 }
 
@@ -105,7 +117,16 @@ where D: Digest + DomainDigest
     pub fn create_from_proofs(
         proofs: Vec<BalancedBinaryMerkleProof<D>>,
     ) -> Result<Self, MergedBalancedBinaryMerkleProofError> {
-        let heights = proofs.iter().map(|proof| proof.path.len()).collect::<Vec<_>>();
+        let heights = proofs
+            .iter()
+            .map(|proof| {
+                proof
+                    .path
+                    .len()
+                    .try_into()
+                    .map_err(|_| MergedBalancedBinaryMerkleProofError::MathOverflow)
+            })
+            .collect::<Result<Vec<_>, _>>()?;
         let max_height = heights
             .iter()
             .max()
@@ -117,7 +138,7 @@ where D: Digest + DomainDigest
             let mut hash_map = HashMap::new();
             for (index, proof) in proofs.iter().enumerate() {
                 // If this path was already joined ignore it.
-                if join_indices[index].is_none() && proof.path.len() > height {
+                if join_indices[index].is_none() && proof.path.len() > height as usize {
                     let parent = (indices[index] - 1) >> 1;
                     if let Some(other_proof) = hash_map.insert(parent, index) {
                         join_indices[index] = Some(other_proof);
@@ -129,7 +150,7 @@ where D: Digest + DomainDigest
                             0,
                             (
                                 MergedBalancedBinaryMerkleDataType::Hash,
-                                proof.path[proof.path.len() - 1 - height].clone(),
+                                proof.path[proof.path.len() - 1 - height as usize].clone(),
                             ),
                         );
                     }
@@ -223,11 +244,11 @@ mod test {
             let hash_last = leaves[n - 1].clone();
             let bmt = BalancedBinaryMerkleTree::<DomainSeparatedHasher<Blake256, TestDomain>>::create(leaves);
             let root = bmt.get_merkle_root();
-            let proof = BalancedBinaryMerkleProof::generate_proof(&bmt, 0);
+            let proof = BalancedBinaryMerkleProof::generate_proof(&bmt, 0).unwrap();
             assert!(proof.verify(&root, hash_0));
-            let proof = BalancedBinaryMerkleProof::generate_proof(&bmt, n / 2);
+            let proof = BalancedBinaryMerkleProof::generate_proof(&bmt, n / 2).unwrap();
             assert!(proof.verify(&root, hash_n_half));
-            let proof = BalancedBinaryMerkleProof::generate_proof(&bmt, n - 1);
+            let proof = BalancedBinaryMerkleProof::generate_proof(&bmt, n - 1).unwrap();
             assert!(proof.verify(&root, hash_last));
         }
     }
@@ -241,7 +262,8 @@ mod test {
         let proofs = indices
             .iter()
             .map(|i| BalancedBinaryMerkleProof::generate_proof(&bmt, *i))
-            .collect::<Vec<_>>();
+            .collect::<Result<Vec<_>, _>>()
+            .unwrap();
         let merged_proof = MergedBalancedBinaryMerkleProof::create_from_proofs(proofs).unwrap();
         assert!(merged_proof
             .verify_consume(&root, indices.iter().map(|i| leaves[*i].clone()).collect::<Vec<_>>())
@@ -255,7 +277,8 @@ mod test {
         let root = bmt.get_merkle_root();
         let proofs = (0..255)
             .map(|i| BalancedBinaryMerkleProof::generate_proof(&bmt, i))
-            .collect::<Vec<_>>();
+            .collect::<Result<Vec<_>, _>>()
+            .unwrap();
         let merged_proof = MergedBalancedBinaryMerkleProof::create_from_proofs(proofs).unwrap();
         assert!(merged_proof.verify_consume(&root, leaves).unwrap());
     }

base_layer/mmr/src/balanced_binary_merkle_tree.rs

@@ -92,7 +92,7 @@ where D: Digest + DomainDigest
         leaf_index + (self.hashes.len() >> 1)
     }
 
-    pub fn find_leaf_index_for_hash(&self, hash: &Hash) -> Result<usize, BalancedBinaryMerkleTreeError> {
+    pub fn find_leaf_index_for_hash(&self, hash: &Hash) -> Result<u32, BalancedBinaryMerkleTreeError> {
         let pos = self
             .hashes
             .position(hash)
@@ -102,7 +102,7 @@ where D: Digest + DomainDigest
             // The hash provided was not for leaf, but for node.
             Err(BalancedBinaryMerkleTreeError::LeafNotFound)
         } else {
-            Ok(pos - (self.hashes.len() >> 1))
+            Ok(pos as u32 - (self.hashes.len() as u32 >> 1))
         }
     }
 }