From 8d5cb014ce9fdb32ca7388ec6b7f06c995449dc5 Mon Sep 17 00:00:00 2001 From: Andrew Poelstra Date: Thu, 13 Jun 2024 14:34:06 +0000 Subject: [PATCH] merkle_tree: introduce MerkleNode trait to better-type merkle tree calculation --- bitcoin/src/blockdata/block.rs | 10 ++-- bitcoin/src/merkle_tree/block.rs | 4 +- bitcoin/src/merkle_tree/mod.rs | 88 ++++++++++++++++++-------------- 3 files changed, 57 insertions(+), 45 deletions(-) diff --git a/bitcoin/src/blockdata/block.rs b/bitcoin/src/blockdata/block.rs index 5f36ce0bb..d3337c86f 100644 --- a/bitcoin/src/blockdata/block.rs +++ b/bitcoin/src/blockdata/block.rs @@ -280,8 +280,8 @@ impl Block { /// Computes the transaction merkle root. pub fn compute_merkle_root(&self) -> Option { - let hashes = self.txdata.iter().map(|obj| obj.compute_txid().to_raw_hash()); - merkle_tree::calculate_root(hashes).map(|h| h.into()) + let hashes = self.txdata.iter().map(|obj| obj.compute_txid()); + merkle_tree::calculate_root(hashes) } /// Computes the witness commitment for the block's transaction list. @@ -300,12 +300,12 @@ impl Block { let hashes = self.txdata.iter().enumerate().map(|(i, t)| { if i == 0 { // Replace the first hash with zeroes. - Wtxid::all_zeros().to_raw_hash() + Wtxid::all_zeros() } else { - t.compute_wtxid().to_raw_hash() + t.compute_wtxid() } }); - merkle_tree::calculate_root(hashes).map(|h| h.into()) + merkle_tree::calculate_root(hashes) } /// Returns the weight of the block. diff --git a/bitcoin/src/merkle_tree/block.rs b/bitcoin/src/merkle_tree/block.rs index 8f090edef..288f26092 100644 --- a/bitcoin/src/merkle_tree/block.rs +++ b/bitcoin/src/merkle_tree/block.rs @@ -563,9 +563,9 @@ mod tests { .collect::>(); // Calculate the merkle root and height - let hashes = tx_ids.iter().map(|t| t.to_raw_hash()); + let hashes = tx_ids.iter().copied(); let merkle_root_1: TxMerkleNode = - merkle_tree::calculate_root(hashes).expect("hashes is not empty").into(); + merkle_tree::calculate_root(hashes).expect("hashes is not empty"); let mut height = 1; let mut ntx = tx_count; while ntx > 1 { diff --git a/bitcoin/src/merkle_tree/mod.rs b/bitcoin/src/merkle_tree/mod.rs index 4ef118d6d..8b5c7da1a 100644 --- a/bitcoin/src/merkle_tree/mod.rs +++ b/bitcoin/src/merkle_tree/mod.rs @@ -6,11 +6,12 @@ //! //! ``` //! # use bitcoin::{merkle_tree, Txid}; +//! # use bitcoin::merkle_tree::TxMerkleNode; //! # use bitcoin::hashes::Hash; //! # let tx1 = Txid::all_zeros(); // Dummy hash values. //! # let tx2 = Txid::all_zeros(); //! let tx_hashes = vec![tx1, tx2]; // All the hashes we wish to merkelize. -//! let root = merkle_tree::calculate_root(tx_hashes.into_iter()); +//! let root: Option = merkle_tree::calculate_root(tx_hashes.into_iter()); //! ``` mod block; @@ -18,10 +19,8 @@ mod block; use core::cmp::min; use core::iter; -use hashes::{sha256d, Hash}; -use io::Write; +use hashes::{sha256d, HashEngine as _}; -use crate::consensus::encode::Encodable; use crate::internal_macros::impl_hashencode; use crate::prelude::*; use crate::{Txid, Wtxid}; @@ -39,12 +38,42 @@ hashes::hash_newtype! { impl_hashencode!(TxMerkleNode); impl_hashencode!(WitnessMerkleNode); -impl From for TxMerkleNode { - fn from(txid: Txid) -> Self { Self::from_byte_array(txid.to_byte_array()) } +/// A node in a Merkle tree of transactions or witness data within a block. +pub trait MerkleNode: Copy { + /// The hash (TXID or WTXID) of a transaciton in the tree. + type Leaf; + + /// Convert a hash to a leaf node of the tree. + fn from_leaf(leaf: Self::Leaf) -> Self; + /// Combine two nodes to get a single node. The final node of a tree is called the "root". + fn combine(&self, other: &Self) -> Self; } -impl From for WitnessMerkleNode { - fn from(wtxid: Wtxid) -> Self { Self::from_byte_array(wtxid.to_byte_array()) } +// These two impl blocks are identical. FIXME once we have nailed down +// our hash traits, it should be possible to put bounds on `MerkleNode` +// and `MerkleNode::Leaf` which are sufficient to turn both methods into +// provided methods in the trait definition. +impl MerkleNode for TxMerkleNode { + type Leaf = Txid; + fn from_leaf(leaf: Self::Leaf) -> Self { Self::from_byte_array(leaf.to_byte_array()) } + + fn combine(&self, other: &Self) -> Self { + let mut encoder = sha256d::Hash::engine(); + encoder.input(self.as_byte_array()); + encoder.input(other.as_byte_array()); + Self(sha256d::Hash::from_engine(encoder)) + } +} +impl MerkleNode for WitnessMerkleNode { + type Leaf = Wtxid; + fn from_leaf(leaf: Self::Leaf) -> Self { Self::from_byte_array(leaf.to_byte_array()) } + + fn combine(&self, other: &Self) -> Self { + let mut encoder = sha256d::Hash::engine(); + encoder.input(self.as_byte_array()); + encoder.input(other.as_byte_array()); + Self(sha256d::Hash::from_engine(encoder)) + } } /// Calculates the merkle root of a list of *hashes*, inline (in place) in `hashes`. @@ -58,11 +87,7 @@ impl From for WitnessMerkleNode { /// - `None` if `hashes` is empty. The merkle root of an empty tree of hashes is undefined. /// - `Some(hash)` if `hashes` contains one element. A single hash is by definition the merkle root. /// - `Some(merkle_root)` if length of `hashes` is greater than one. -pub fn calculate_root_inline(hashes: &mut [T]) -> Option -where - T: Hash + Encodable, - ::Engine: Write, -{ +pub fn calculate_root_inline(hashes: &mut [T]) -> Option { match hashes.len() { 0 => None, 1 => Some(hashes[0]), @@ -79,14 +104,13 @@ where /// - `Some(merkle_root)` if length of `hashes` is greater than one. pub fn calculate_root(mut hashes: I) -> Option where - T: Hash + Encodable, - ::Engine: Write, - I: Iterator, + T: MerkleNode, + I: Iterator, { - let first = hashes.next()?; + let first: T::Leaf = hashes.next()?; let second = match hashes.next() { Some(second) => second, - None => return Some(first), + None => return Some(T::from_leaf(first)), }; let mut hashes = iter::once(first).chain(iter::once(second)).chain(hashes); @@ -96,24 +120,17 @@ where let (min, max) = hashes.size_hint(); let mut alloc = Vec::with_capacity(max.unwrap_or(min) / 2 + 1); - while let Some(hash1) = hashes.next() { + while let Some(hash1) = hashes.next().map(T::from_leaf) { // If the size is odd, use the last element twice. - let hash2 = hashes.next().unwrap_or(hash1); - let mut encoder = T::engine(); - hash1.consensus_encode(&mut encoder).expect("in-memory writers don't error"); - hash2.consensus_encode(&mut encoder).expect("in-memory writers don't error"); - alloc.push(T::from_engine(encoder)); + let hash2 = hashes.next().map(T::from_leaf).unwrap_or(hash1); + alloc.push(hash1.combine(&hash2)); } Some(merkle_root_r(&mut alloc)) } // `hashes` must contain at least one hash. -fn merkle_root_r(hashes: &mut [T]) -> T -where - T: Hash + Encodable, - ::Engine: Write, -{ +fn merkle_root_r(hashes: &mut [T]) -> T { if hashes.len() == 1 { return hashes[0]; } @@ -121,10 +138,7 @@ where for idx in 0..((hashes.len() + 1) / 2) { let idx1 = 2 * idx; let idx2 = min(idx1 + 1, hashes.len() - 1); - let mut encoder = T::engine(); - hashes[idx1].consensus_encode(&mut encoder).expect("in-memory writers don't error"); - hashes[idx2].consensus_encode(&mut encoder).expect("in-memory writers don't error"); - hashes[idx] = T::from_engine(encoder); + hashes[idx] = hashes[idx1].combine(&hashes[idx2]); } let half_len = hashes.len() / 2 + hashes.len() % 2; @@ -133,8 +147,6 @@ where #[cfg(test)] mod tests { - use hashes::sha256d; - use super::*; use crate::blockdata::block::Block; use crate::consensus::encode::deserialize; @@ -146,11 +158,11 @@ mod tests { let block: Block = deserialize(&segwit_block[..]).expect("Failed to deserialize block"); assert!(block.check_merkle_root()); // Sanity check. - let hashes_iter = block.txdata.iter().map(|obj| obj.compute_txid().to_raw_hash()); + let hashes_iter = block.txdata.iter().map(|obj| obj.compute_txid()); - let mut hashes_array = [sha256d::Hash::all_zeros(); 15]; + let mut hashes_array = [TxMerkleNode::all_zeros(); 15]; for (i, hash) in hashes_iter.clone().enumerate() { - hashes_array[i] = hash; + hashes_array[i] = TxMerkleNode::from_leaf(hash); } let from_iter = calculate_root(hashes_iter);