1- use std:: collections:: HashSet ;
1+ use std:: collections:: HashMap ;
22use bitcoin:: { Network , ScriptBuf } ;
33use bitcoin:: taproot:: { LeafVersion , TapTree , ScriptLeaves , TapLeafHash , TaprootMerkleBranch , TapNodeHash } ;
44use bitcoin:: p2mr:: { P2mrBuilder , P2mrControlBlock , P2mrSpendInfo } ;
@@ -136,21 +136,20 @@ fn process_test_vector_p2mr(test_vector: &TestVector) -> anyhow::Result<()> {
136136 // Use of TaprootBuilder avoids user error in constructing branches manually and ensures Merkle tree correctness and determinism
137137 let mut p2mr_builder: P2mrBuilder = P2mrBuilder :: new ( ) ;
138138
139- let mut control_block_data : Vec < ( ScriptBuf , LeafVersion ) > = Vec :: new ( ) ;
139+ let mut script_to_id : HashMap < ScriptBuf , u8 > = HashMap :: new ( ) ;
140140
141141 // 1) traverse test vector script tree and add leaves to P2MR builder
142142 if let Some ( script_tree) = tv_script_tree {
143143
144144 script_tree. traverse_with_right_subtree_first ( 0 , Direction :: Root , & mut |node, depth, direction| {
145145
146146 if let TVScriptTree :: Leaf ( tv_leaf) = node {
147-
147+
148148 let tv_leaf_script_bytes = hex:: decode ( & tv_leaf. script ) . unwrap ( ) ;
149-
150- // NOTE: IOT to execute script_info.control_block(..), will add these to a vector
149+
151150 let tv_leaf_script_buf = ScriptBuf :: from_bytes ( tv_leaf_script_bytes. clone ( ) ) ;
152151 let tv_leaf_version = LeafVersion :: from_consensus ( tv_leaf. leaf_version ) . unwrap ( ) ;
153- control_block_data . push ( ( tv_leaf_script_buf. clone ( ) , tv_leaf_version ) ) ;
152+ script_to_id . insert ( tv_leaf_script_buf. clone ( ) , tv_leaf . id ) ;
154153
155154 let mut modified_depth = depth + 1 ;
156155 if direction == Direction :: Root {
@@ -194,14 +193,10 @@ fn process_test_vector_p2mr(test_vector: &TestVector) -> anyhow::Result<()> {
194193 ) ;
195194 debug ! ( "just passed merkle root validation: {}" , test_vector_merkle_root) ;
196195
197- let test_vector_leaf_hashes_vec: Vec < String > = test_vector. intermediary . leaf_hashes . clone ( ) ;
198- let test_vector_leaf_hash_set: HashSet < String > = test_vector_leaf_hashes_vec. iter ( ) . cloned ( ) . collect ( ) ;
199- let test_vector_control_blocks_vec = & test_vector. expected . script_path_control_blocks ;
200- let test_vector_control_blocks_set: HashSet < String > = test_vector_control_blocks_vec. as_ref ( ) . unwrap ( ) . iter ( ) . cloned ( ) . collect ( ) ;
196+ let expected_control_blocks = test_vector. expected . script_path_control_blocks . as_ref ( ) . unwrap ( ) ;
201197 let tap_tree: TapTree = p2mr_builder. clone ( ) . into_inner ( ) . try_into_taptree ( ) . unwrap ( ) ;
202198 let script_leaves: ScriptLeaves = tap_tree. script_leaves ( ) ;
203199
204- // TO-DO: Investigate why the ordering of script leaves seems to be reverse of test vectors.
205200 // 3) Iterate through leaves of derived script tree and verify both script leaf hashes and control blocks
206201 for derived_leaf in script_leaves {
207202
@@ -211,34 +206,21 @@ fn process_test_vector_p2mr(test_vector: &TestVector) -> anyhow::Result<()> {
211206
212207 let derived_leaf_hash: TapLeafHash = TapLeafHash :: from_script ( script, version) ;
213208 let leaf_hash = hex:: encode ( derived_leaf_hash. as_raw_hash ( ) . to_byte_array ( ) ) ;
214- assert ! (
215- test_vector_leaf_hash_set. contains( & leaf_hash) ,
216- "Leaf hash not found in expected set for {}" , leaf_hash
217- ) ;
218- debug ! ( "just passed leaf_hash validation: {}" , leaf_hash) ;
219-
220- // Each leaf in the script tree has a corresponding control block.
221- // Specific to P2TR, the 3 sections of the control block (control byte, public key & merkle path) are highlighted here:
222- // https://learnmeabitcoin.com/technical/upgrades/taproot/#script-path-spend-control-block
223- // The control block, which includes the Merkle path, must be 33 + 32 * n bytes, where n is the number of Merkle path hashes (n ≥ 0).
224- // There is no consensus limit on n, but large Merkle trees increase the witness size, impacting block weight.
225- // NOTE: Control blocks could have also been obtained from spend_info.control_block(..) using the data in control_block_data
226- debug ! ( "merkle_branch nodes: {:?}" , merkle_branch) ;
209+
210+ let leaf_id = script_to_id. get ( script)
211+ . unwrap_or_else ( || panic ! ( "leaf script not found in script_to_id map: {}" , hex:: encode( script. as_bytes( ) ) ) ) ;
212+
227213 let derived_control_block: P2mrControlBlock = P2mrControlBlock {
228214 merkle_branch : merkle_branch. clone ( ) ,
229215 } ;
230- let serialized_control_block = derived_control_block. serialize ( ) ;
231- debug ! ( "derived_control_block: {:?}, merkle_branch size: {}, control_block size: {}, serialized size: {}" ,
232- derived_control_block,
233- merkle_branch. len( ) ,
234- derived_control_block. size( ) ,
235- serialized_control_block. len( ) ) ;
236- let derived_serialized_control_block = hex:: encode ( serialized_control_block) ;
237- assert ! (
238- test_vector_control_blocks_set. contains( & derived_serialized_control_block) ,
239- "Control block mismatch: {}, expected: {:?}" , derived_serialized_control_block, test_vector_control_blocks_set
216+ let derived_serialized_control_block = hex:: encode ( derived_control_block. serialize ( ) ) ;
217+
218+ let expected_cb = & expected_control_blocks[ * leaf_id as usize ] ;
219+ assert_eq ! (
220+ derived_serialized_control_block, * expected_cb,
221+ "Control block mismatch for leaf id {}: derived {}, expected {}" , leaf_id, derived_serialized_control_block, expected_cb
240222 ) ;
241- debug ! ( "leaf_hash: {}, derived_serialized_control_block: {}" , leaf_hash, derived_serialized_control_block) ;
223+ debug ! ( "leaf_id: {}, leaf_hash: {}, derived_serialized_control_block: {}" , leaf_id , leaf_hash, derived_serialized_control_block) ;
242224
243225 }
244226
0 commit comments