Merge rust-bitcoin/rust-bitcoin#1389: Add `PartialMerkleTree::calc_tree_height` method

752adff9d1 Add method calc_height (Tobin C. Harding)
46f5588646 Add unit test for calc_tree_width (Tobin C. Harding)

Pull request description:

  Add a private `PartialMerkleTree::calc_tree_width` function and a unit test to test it.

ACKs for top commit:
  apoelstra:
    ACK 752adff9d1
  clarkmoody:
    ACK 752adff9d1

Tree-SHA512: 9c4ad9f6ff47d8faad1c7c1e977427f1528af2712ceffd05357d0c9117b5fdb7b2783afc00a75cb19b853bfbd7b3895baa3a3563bdc496593cc9b06ce80dbbf8
This commit is contained in:
Andrew Poelstra 2023-10-03 13:41:56 +00:00
commit fcd4ad51ef
No known key found for this signature in database
GPG Key ID: C588D63CE41B97C1
1 changed files with 64 additions and 10 deletions

View File

@ -250,11 +250,8 @@ impl PartialMerkleTree {
bits: Vec::with_capacity(txids.len()),
hashes: vec![],
};
// calculate height of tree
let mut height = 0;
while pmt.calc_tree_width(height) > 1 {
height += 1;
}
let height = pmt.calc_tree_height();
// traverse the partial tree
pmt.traverse_and_build(height, 0, txids, matches);
pmt
@ -286,11 +283,9 @@ impl PartialMerkleTree {
if self.bits.len() < self.hashes.len() {
return Err(NotEnoughBits);
};
// calculate height of tree
let mut height = 0;
while self.calc_tree_width(height) > 1 {
height += 1;
}
let height = self.calc_tree_height();
// traverse the partial tree
let mut bits_used = 0u32;
let mut hash_used = 0u32;
@ -308,6 +303,15 @@ impl PartialMerkleTree {
Ok(TxMerkleNode::from_byte_array(hash_merkle_root.to_byte_array()))
}
/// Calculates the height of the tree.
fn calc_tree_height(&self) -> u32 {
let mut height = 0;
while self.calc_tree_width(height) > 1 {
height += 1;
}
height
}
/// Helper function to efficiently calculate the number of nodes at given height
/// in the merkle tree
#[inline]
@ -767,4 +771,54 @@ mod tests {
let block_hex = include_str!("../../tests/data/block_13b8a.hex");
deserialize(&Vec::from_hex(block_hex).unwrap()).unwrap()
}
macro_rules! check_calc_tree_width {
($($test_name:ident, $num_transactions:literal, $height:literal, $expected_width:literal);* $(;)?) => {
$(
#[test]
fn $test_name() {
let pmt = PartialMerkleTree {
num_transactions: $num_transactions,
bits: vec![],
hashes: vec![],
};
let got = pmt.calc_tree_width($height);
assert_eq!(got, $expected_width)
}
)*
}
}
// tree_width_<id> <num txs> <height> <expected_width>
//
// height 0 is the bottom of the tree, where the leaves are.
check_calc_tree_width! {
tree_width_01, 1, 0, 1;
//
tree_width_02, 2, 0, 2;
tree_width_03, 2, 1, 1;
//
tree_width_04, 3, 0, 3;
tree_width_05, 3, 1, 2;
tree_width_06, 3, 2, 1;
//
tree_width_07, 4, 0, 4;
tree_width_08, 4, 1, 2;
tree_width_09, 4, 2, 1;
//
tree_width_10, 5, 0, 5;
tree_width_11, 5, 1, 3;
tree_width_12, 5, 2, 2;
tree_width_13, 5, 3, 1;
//
tree_width_14, 6, 0, 6;
tree_width_15, 6, 1, 3;
tree_width_16, 6, 2, 2;
tree_width_17, 6, 3, 1;
//
tree_width_18, 7, 0, 7;
tree_width_19, 7, 1, 4;
tree_width_20, 7, 2, 2;
tree_width_21, 7, 3, 1;
}
}