#[cfg(feature = "alloc")]
use alloc::vec::Vec;
use hashes::{sha256d, HashEngine};
#[cfg(not(feature = "alloc"))]
use internals::array_vec::ArrayVec;
#[doc(inline)]
pub use crate::hash_types::{TxMerkleNode, TxMerkleNodeEncoder, WitnessMerkleNode};
use crate::hash_types::{Txid, Wtxid};
use crate::transaction::TxIdentifier;
pub(crate) trait MerkleNode: Copy + PartialEq {
type Leaf: TxIdentifier;
fn from_leaf(leaf: Self::Leaf) -> Self;
#[must_use]
fn combine(&self, other: &Self) -> Self;
fn calculate_root<I: Iterator<Item = Self::Leaf>>(iter: I) -> Option<Self> {
{
#[cfg(feature = "alloc")]
let mut stack = Vec::<(usize, Self)>::with_capacity(32);
#[cfg(not(feature = "alloc"))]
let mut stack = ArrayVec::<(usize, Self), 15>::new();
for (mut n, leaf) in iter.enumerate() {
#[cfg(not(feature = "alloc"))]
if stack.len() == 15 {
return None;
}
stack.push((0, Self::from_leaf(leaf)));
while n & 1 == 1 {
let right = stack.pop().unwrap();
let left = stack.pop().unwrap();
if left.1 == right.1 {
return None;
}
debug_assert_eq!(left.0, right.0);
stack.push((left.0 + 1, left.1.combine(&right.1)));
n >>= 1;
}
}
while stack.len() > 1 {
let mut right = stack.pop().unwrap();
let left = stack.pop().unwrap();
while right.0 != left.0 {
assert!(right.0 < left.0);
right = (right.0 + 1, right.1.combine(&right.1)); }
stack.push((left.0 + 1, left.1.combine(&right.1)));
}
stack.pop().map(|(_, h)| h)
}
}
}
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::from_byte_array(sha256d::Hash::from_engine(encoder).to_byte_array())
}
}
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::from_byte_array(sha256d::Hash::from_engine(encoder).to_byte_array())
}
}
#[cfg(test)]
mod tests {
use crate::hash_types::*;
fn make_leaf_node(byte: u8) -> (Txid, TxMerkleNode) {
let leaf = Txid::from_byte_array([byte; 32]);
let node = TxMerkleNode::from_leaf(leaf);
(leaf, node)
}
#[test]
fn tx_merkle_node_single_leaf() {
let (leaf, node) = make_leaf_node(1);
let root = TxMerkleNode::calculate_root([leaf].into_iter());
assert!(root.is_some(), "Root should exist for a single leaf");
assert_eq!(root.unwrap(), node, "Root should equal the leaf node");
}
#[test]
fn tx_merkle_node_two_leaves() {
let (leaf1, node1) = make_leaf_node(1);
let (leaf2, node2) = make_leaf_node(2);
let combined = node1.combine(&node2);
let root = TxMerkleNode::calculate_root([leaf1, leaf2].into_iter());
assert_eq!(
root.unwrap(),
combined,
"Root of two leaves should equal combine of the two leaf nodes"
);
}
#[test]
fn tx_merkle_node_duplicate_leaves() {
let leaf = Txid::from_byte_array([3; 32]);
let root = TxMerkleNode::calculate_root([leaf, leaf].into_iter());
assert!(root.is_none(), "Duplicate leaves should return None");
}
#[test]
fn tx_merkle_node_empty() {
assert!(
TxMerkleNode::calculate_root([].into_iter()).is_none(),
"Empty iterator should return None"
);
}
#[test]
fn tx_merkle_node_2n_minus_1_unbalanced_tree() {
let (leaf1, node1) = make_leaf_node(1);
let (leaf2, node2) = make_leaf_node(2);
let (leaf3, node3) = make_leaf_node(3);
let (leaf4, node4) = make_leaf_node(4);
let (leaf5, node5) = make_leaf_node(5);
let (leaf6, node6) = make_leaf_node(6);
let (leaf7, node7) = make_leaf_node(7);
let subtree_a = node1.combine(&node2);
let subtree_b = node3.combine(&node4);
let subtree_c = node5.combine(&node6);
let subtree_d = node7.combine(&node7);
let subtree_ab = subtree_a.combine(&subtree_b);
let subtree_cd = subtree_c.combine(&subtree_d);
let expected = subtree_ab.combine(&subtree_cd);
let root = TxMerkleNode::calculate_root(
[leaf1, leaf2, leaf3, leaf4, leaf5, leaf6, leaf7].into_iter(),
);
assert_eq!(root, Some(expected));
}
#[test]
#[cfg(feature = "alloc")]
fn tx_merkle_node_balanced_multi_level_tree() {
use alloc::vec::Vec;
let leaves: Vec<_> = (0..16).map(|i| Txid::from_byte_array([i; 32])).collect();
let mut level = leaves.iter().map(|l| TxMerkleNode::from_leaf(*l)).collect::<Vec<_>>();
while level.len() > 1 {
level = level.chunks(2).map(|chunk| chunk[0].combine(&chunk[1])).collect();
}
let expected = level.pop().unwrap();
let root = TxMerkleNode::calculate_root(leaves.into_iter());
assert_eq!(root, Some(expected));
}
#[test]
fn tx_merkle_node_oversize_tree() {
let root = TxMerkleNode::calculate_root((0..32768u32).map(|i| {
let mut buf = [0u8; 32];
buf[..4].copy_from_slice(&i.to_le_bytes());
Txid::from_byte_array(buf)
}));
#[cfg(feature = "alloc")]
assert_ne!(root, None);
#[cfg(not(feature = "alloc"))]
assert_eq!(root, None);
let root = TxMerkleNode::calculate_root((0..32767u32).map(|i| {
let mut buf = [0u8; 32];
buf[..4].copy_from_slice(&i.to_le_bytes());
Txid::from_byte_array(buf)
}));
assert_ne!(root, None);
}
#[test]
fn witness_merkle_node_single_leaf() {
let leaf = Wtxid::from_byte_array([1; 32]);
let root = WitnessMerkleNode::calculate_root([leaf].into_iter());
assert!(root.is_some(), "Root should exist for a single witness leaf");
let node = WitnessMerkleNode::from_leaf(leaf);
assert_eq!(root.unwrap(), node, "Root should equal the leaf node");
}
#[test]
fn witness_merkle_node_duplicate_leaves() {
let leaf = Wtxid::from_byte_array([2; 32]);
let root = WitnessMerkleNode::calculate_root([leaf, leaf].into_iter());
assert!(root.is_none(), "Duplicate witness leaves should return None");
}
}