Kpp/merkle tests

crates/bitcoin-da/Cargo.toml

@@ -7,7 +7,7 @@ homepage = { workspace = true }
 license = "MIT OR Apache-2.0"
 publish = false
 repository = { workspace = true }
-rust-version = "1.66"
+rust-version = "1.67"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 

crates/bitcoin-da/src/helpers/merkle_tree.rs

@@ -5,7 +5,6 @@ use super::calculate_double_sha256;
 
 #[derive(Debug, Clone)]
 pub struct BitcoinMerkleTree {
-    depth: u32,
     nodes: Vec<Vec<[u8; 32]>>,
 }
 
@@ -14,14 +13,11 @@ impl BitcoinMerkleTree {
         if transactions.len() == 1 {
             // root is the coinbase txid
             return BitcoinMerkleTree {
-                depth: 1,
                 nodes: vec![transactions],
             };
         }
 
-        let depth = (transactions.len() - 1).ilog(2) + 1;
         let mut tree = BitcoinMerkleTree {
-            depth,
             nodes: vec![transactions],
         };
 
@@ -34,7 +30,7 @@ impl BitcoinMerkleTree {
             tree.nodes.push(vec![]);
             for i in 0..(prev_level_size / 2) {
                 preimage[..32].copy_from_slice(
-                    &tree.nodes[curr_level_offset - 1 as usize][prev_level_index_offset + i * 2],
+                    &tree.nodes[curr_level_offset - 1][prev_level_index_offset + i * 2],
                 );
                 preimage[32..].copy_from_slice(
                     &tree.nodes[curr_level_offset - 1][prev_level_index_offset + i * 2 + 1],
@@ -64,35 +60,26 @@ impl BitcoinMerkleTree {
 
     // Returns the Merkle root
     pub fn root(&self) -> [u8; 32] {
-        return self.nodes[self.nodes.len() - 1][0];
-    }
-
-    pub fn get_element(&self, level: u32, index: u32) -> [u8; 32] {
-        return self.nodes[level as usize][index as usize];
+        self.nodes[self.nodes.len() - 1][0]
     }
 
     pub fn get_idx_path(&self, index: u32) -> Vec<[u8; 32]> {
-        assert!(
-            index <= self.nodes[0].len() as u32 - 1,
-            "Index out of bounds"
-        );
+        assert!(index < self.nodes[0].len() as u32, "Index out of bounds");
         let mut path = vec![];
         let mut level = 0;
         let mut i = index;
         while level < self.nodes.len() as u32 - 1 {
             if i % 2 == 1 {
                 path.push(self.nodes[level as usize][i as usize - 1]);
+            } else if (self.nodes[level as usize].len() - 1) as u32 == i {
+                path.push(self.nodes[level as usize][i as usize]);
             } else {
-                if (self.nodes[level as usize].len() - 1) as u32 == i {
-                    path.push(self.nodes[level as usize][i as usize]);
-                } else {
-                    path.push(self.nodes[level as usize][(i + 1) as usize]);
-                }
+                path.push(self.nodes[level as usize][(i + 1) as usize]);
             }
             level += 1;
-            i = i / 2;
+            i /= 2;
         }
-        return path;
+        path
     }
 
     pub fn calculate_root_with_merkle_proof(
@@ -101,7 +88,7 @@ impl BitcoinMerkleTree {
         merkle_proof: Vec<[u8; 32]>,
     ) -> [u8; 32] {
         let mut preimage: [u8; 64] = [0; 64];
-        let mut combined_hash: [u8; 32] = txid.clone();
+        let mut combined_hash: [u8; 32] = txid;
         let mut index = idx;
         let mut level: u32 = 0;
         while level < merkle_proof.len() as u32 {
@@ -115,18 +102,21 @@ impl BitcoinMerkleTree {
                 combined_hash = calculate_double_sha256(&preimage);
             }
             level += 1;
-            index = index / 2;
+            index /= 2;
         }
         combined_hash
     }
 }
 
 #[cfg(test)]
 mod tests {
+    use bitcoin::hashes::Hash;
+
     use super::*;
+    use crate::helpers::test_utils::get_mock_txs;
 
     #[test]
-    fn test_merkle_tree() {
+    fn test_merkle_root_with_proof() {
         let mut transactions: Vec<[u8; 32]> = vec![];
         for i in 0u8..100u8 {
             let tx = [i; 32];
@@ -139,4 +129,35 @@ mod tests {
             BitcoinMerkleTree::calculate_root_with_merkle_proof(transactions[0], 0, idx_path);
         assert_eq!(root, calculated_root);
     }
+
+    #[test]
+    fn test_merkle_tree_single_tx() {
+        let tx = [5; 32];
+        assert_eq!(BitcoinMerkleTree::new(vec![tx]).root(), tx);
+    }
+
+    #[test]
+    fn test_merkle_tree_against_bitcoin_impl() {
+        compare_merkle_tree_against_bitcoin_impl(vec![[0; 32]; 100]);
+        compare_merkle_tree_against_bitcoin_impl(vec![[5; 32]; 10]);
+        compare_merkle_tree_against_bitcoin_impl(vec![[255; 32]; 33]);
+        compare_merkle_tree_against_bitcoin_impl(vec![[200; 32]; 2]);
+        compare_merkle_tree_against_bitcoin_impl(vec![[99; 32]; 1]);
+
+        let txs = get_mock_txs()
+            .iter()
+            .map(|tx| tx.compute_wtxid().to_byte_array())
+            .collect();
+        compare_merkle_tree_against_bitcoin_impl(txs);
+    }
+
+    fn compare_merkle_tree_against_bitcoin_impl(transactions: Vec<[u8; 32]>) {
+        let hashes = transactions
+            .iter()
+            .map(|tx| bitcoin::hash_types::Wtxid::from_slice(tx).unwrap());
+        let bitcoin_root = bitcoin::merkle_tree::calculate_root(hashes).unwrap();
+
+        let custom_root = BitcoinMerkleTree::new(transactions).root();
+        assert_eq!(bitcoin_root.to_byte_array(), custom_root);
+    }
 }

crates/bitcoin-da/src/helpers/mod.rs

@@ -113,5 +113,5 @@ pub fn calculate_double_sha256(input: &[u8]) -> [u8; 32] {
     hasher.update(input);
     let result = hasher.finalize_reset();
     hasher.update(result);
-    hasher.finalize().try_into().unwrap()
+    hasher.finalize().into()
 }