From 81f65c9c3adb948a8520336759f2a5b593c00348 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ahmet=20Yaz=C4=B1c=C4=B1?= <75089142+yaziciahmet@users.noreply.github.com> Date: Thu, 15 Aug 2024 13:46:38 +0300 Subject: [PATCH 1/2] Add merkle tree tests (#995) * Add merkle tree tests * Add 1 tx test to bitcoin compare --------- Co-authored-by: yaziciahmet --- crates/bitcoin-da/src/helpers/merkle_tree.rs | 52 ++++++++++++++++++-- 1 file changed, 48 insertions(+), 4 deletions(-) diff --git a/crates/bitcoin-da/src/helpers/merkle_tree.rs b/crates/bitcoin-da/src/helpers/merkle_tree.rs index 1c4fc9091..81e2f3dfc 100644 --- a/crates/bitcoin-da/src/helpers/merkle_tree.rs +++ b/crates/bitcoin-da/src/helpers/merkle_tree.rs @@ -11,6 +11,12 @@ pub struct BitcoinMerkleTree { impl BitcoinMerkleTree { pub fn new(transactions: Vec<[u8; 32]>) -> Self { + if transactions.len() == 0 { + return BitcoinMerkleTree { + depth: 0, + nodes: vec![], + }; + } if transactions.len() == 1 { // root is the coinbase txid return BitcoinMerkleTree { @@ -63,8 +69,11 @@ impl BitcoinMerkleTree { } // Returns the Merkle root - pub fn root(&self) -> [u8; 32] { - return self.nodes[self.nodes.len() - 1][0]; + pub fn root(&self) -> Option<[u8; 32]> { + if self.nodes.len() == 0 { + return None; + } + return Some(self.nodes[self.nodes.len() - 1][0]); } pub fn get_element(&self, level: u32, index: u32) -> [u8; 32] { @@ -123,20 +132,55 @@ impl BitcoinMerkleTree { #[cfg(test)] mod tests { + use bitcoin::hashes::Hash; + + use crate::helpers::test_utils::get_mock_txs; + use super::*; #[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]; transactions.push(tx); } let tree = BitcoinMerkleTree::new(transactions.clone()); - let root = tree.root(); + let root = tree.root().unwrap(); let idx_path = tree.get_idx_path(0); let calculated_root = BitcoinMerkleTree::calculate_root_with_merkle_proof(transactions[0], 0, idx_path); assert_eq!(root, calculated_root); } + + #[test] + fn test_empty_merkle_tree() { + assert_eq!(BitcoinMerkleTree::new(vec![]).root(), None); + } + + #[test] + fn test_merkle_tree_single_tx() { + let tx = [5; 32]; + assert_eq!(BitcoinMerkleTree::new(vec![tx]).root().unwrap(), 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().unwrap(); + assert_eq!(bitcoin_root.to_byte_array(), custom_root); + } } From 7a3f4e0f85f7b0c43a0b2c31121770e08e9c6ca1 Mon Sep 17 00:00:00 2001 From: Roman Proskuryakoff Date: Tue, 20 Aug 2024 01:00:35 +0400 Subject: [PATCH 2/2] Fix errors/lints --- crates/bitcoin-da/Cargo.toml | 2 +- crates/bitcoin-da/src/helpers/merkle_tree.rs | 67 +++++++------------- crates/bitcoin-da/src/helpers/mod.rs | 2 +- 3 files changed, 24 insertions(+), 47 deletions(-) diff --git a/crates/bitcoin-da/Cargo.toml b/crates/bitcoin-da/Cargo.toml index 0a4c9ff69..285da8c09 100644 --- a/crates/bitcoin-da/Cargo.toml +++ b/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 diff --git a/crates/bitcoin-da/src/helpers/merkle_tree.rs b/crates/bitcoin-da/src/helpers/merkle_tree.rs index 81e2f3dfc..ff56887ff 100644 --- a/crates/bitcoin-da/src/helpers/merkle_tree.rs +++ b/crates/bitcoin-da/src/helpers/merkle_tree.rs @@ -5,29 +5,19 @@ use super::calculate_double_sha256; #[derive(Debug, Clone)] pub struct BitcoinMerkleTree { - depth: u32, nodes: Vec>, } impl BitcoinMerkleTree { pub fn new(transactions: Vec<[u8; 32]>) -> Self { - if transactions.len() == 0 { - return BitcoinMerkleTree { - depth: 0, - nodes: vec![], - }; - } 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], }; @@ -40,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], @@ -69,39 +59,27 @@ impl BitcoinMerkleTree { } // Returns the Merkle root - pub fn root(&self) -> Option<[u8; 32]> { - if self.nodes.len() == 0 { - return None; - } - return Some(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]; + pub fn root(&self) -> [u8; 32] { + 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( @@ -110,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 { @@ -124,7 +102,7 @@ impl BitcoinMerkleTree { combined_hash = calculate_double_sha256(&preimage); } level += 1; - index = index / 2; + index /= 2; } combined_hash } @@ -134,9 +112,8 @@ impl BitcoinMerkleTree { mod tests { use bitcoin::hashes::Hash; - use crate::helpers::test_utils::get_mock_txs; - use super::*; + use crate::helpers::test_utils::get_mock_txs; #[test] fn test_merkle_root_with_proof() { @@ -146,22 +123,17 @@ mod tests { transactions.push(tx); } let tree = BitcoinMerkleTree::new(transactions.clone()); - let root = tree.root().unwrap(); + let root = tree.root(); let idx_path = tree.get_idx_path(0); let calculated_root = BitcoinMerkleTree::calculate_root_with_merkle_proof(transactions[0], 0, idx_path); assert_eq!(root, calculated_root); } - #[test] - fn test_empty_merkle_tree() { - assert_eq!(BitcoinMerkleTree::new(vec![]).root(), None); - } - #[test] fn test_merkle_tree_single_tx() { let tx = [5; 32]; - assert_eq!(BitcoinMerkleTree::new(vec![tx]).root().unwrap(), tx); + assert_eq!(BitcoinMerkleTree::new(vec![tx]).root(), tx); } #[test] @@ -172,15 +144,20 @@ mod tests { 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(); + 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 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().unwrap(); + let custom_root = BitcoinMerkleTree::new(transactions).root(); assert_eq!(bitcoin_root.to_byte_array(), custom_root); } } diff --git a/crates/bitcoin-da/src/helpers/mod.rs b/crates/bitcoin-da/src/helpers/mod.rs index 7411b1f1a..4ec27d493 100644 --- a/crates/bitcoin-da/src/helpers/mod.rs +++ b/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() }