fix: review fixes for P2MR implementation

lcovar · lcovar · commit 95389bd64103 · 2026-04-07T15:09:52.000-07:00
- Remove find_leaf_version() redundant DFS, store leaf version during compute_node traversal instead - Enable P2MR address encoding on all Bitcoin networks (mainnet + testnets), not just BitGo Signet - Add P2MR fixtures to bitcoin.json, testnet.json, bitcoinPublicSignet.json - Validate mainnet bc1z addresses from BIP-360 fixtures directly - Add missing Rust tests for lightning contract and three-leaf alternative vectors (now all 6 construction vectors covered in Rust) BTC-3241
diff --git a/packages/wasm-utxo/src/address/networks.rs b/packages/wasm-utxo/src/address/networks.rs
@@ -185,8 +185,9 @@ impl Network {
         let taproot = segwit && matches!(self.mainnet(), Network::Bitcoin);
 
         // P2MR (BIP-360) support:
-        // Currently only enabled on the BitGo Signet for testing.
-        let p2mr = matches!(self, Network::BitcoinBitGoSignet);
+        // Enabled on all Bitcoin networks (mainnet + testnets) for address encoding.
+        // Backend activation is controlled separately.
+        let p2mr = matches!(self.mainnet(), Network::Bitcoin);
 
         OutputScriptSupport {
             segwit,
diff --git a/packages/wasm-utxo/src/p2mr/mod.rs b/packages/wasm-utxo/src/p2mr/mod.rs
@@ -108,23 +108,20 @@ pub struct P2mrTreeInfo {
     pub leaves: Vec<P2mrLeafInfo>,
 }
 
+/// Intermediate leaf data collected during tree traversal: (leaf_hash, leaf_version, merkle_path).
+type LeafCollector = Vec<([u8; 32], u8, Vec<[u8; 32]>)>;
+
 /// Build a P2MR Merkle tree from a script tree definition.
 ///
 /// Returns the Merkle root and per-leaf spending info (leaf hash + control block).
 /// Leaves are returned in left-to-right DFS order matching the input tree structure.
 pub fn build_p2mr_tree(tree: &ScriptTreeNode) -> P2mrTreeInfo {
-    // Compute the merkle root and collect leaf hashes with their merkle proof paths.
-    // Leaves are output in sorted DFS order (matching tap_branch_hash lex sorting).
-    let mut leaves: Vec<([u8; 32], Vec<[u8; 32]>)> = Vec::new();
+    let mut leaves: LeafCollector = Vec::new();
     let merkle_root = compute_node(tree, &mut leaves);
 
     let leaf_infos = leaves
         .into_iter()
-        .map(|(leaf_hash, path)| {
-            // Determine leaf version from the leaf hash by looking up the tree.
-            // We need the leaf version for the control byte. Extract it from the tree.
-            let leaf_version =
-                find_leaf_version(tree, &leaf_hash).unwrap_or(TAPSCRIPT_LEAF_VERSION);
+        .map(|(leaf_hash, leaf_version, path)| {
             let mut control_block = vec![p2mr_control_byte(leaf_version)];
             for sibling in &path {
                 control_block.extend_from_slice(sibling);
@@ -144,16 +141,16 @@ pub fn build_p2mr_tree(tree: &ScriptTreeNode) -> P2mrTreeInfo {
 
 /// Recursively compute the hash of a tree node, collecting leaf info along the way.
 ///
-/// Leaves are output in input DFS order (left subtree before right subtree,
-/// preserving the tree structure as given).
-fn compute_node(node: &ScriptTreeNode, leaves: &mut Vec<([u8; 32], Vec<[u8; 32]>)>) -> [u8; 32] {
+/// Each leaf entry stores (leaf_hash, leaf_version, merkle_path_siblings).
+/// Leaves are output in input DFS order (left subtree before right subtree).
+fn compute_node(node: &ScriptTreeNode, leaves: &mut LeafCollector) -> [u8; 32] {
     match node {
         ScriptTreeNode::Leaf {
             script,
             leaf_version,
         } => {
             let hash = tap_leaf_hash(script, *leaf_version);
-            leaves.push((hash, Vec::new()));
+            leaves.push((hash, *leaf_version, Vec::new()));
             hash
         }
         ScriptTreeNode::Branch(left, right) => {
@@ -165,36 +162,17 @@ fn compute_node(node: &ScriptTreeNode, leaves: &mut Vec<([u8; 32], Vec<[u8; 32]>
             // Add sibling hashes to the merkle proof paths.
             // Left subtree leaves need right_hash as sibling, and vice versa.
             for leaf in leaves[left_start..right_start].iter_mut() {
-                leaf.1.push(right_hash);
+                leaf.2.push(right_hash);
             }
             for leaf in leaves[right_start..].iter_mut() {
-                leaf.1.push(left_hash);
+                leaf.2.push(left_hash);
             }
 
             tap_branch_hash(&left_hash, &right_hash)
         }
     }
 }
 
-/// Find the leaf version for a leaf with the given hash (DFS search).
-fn find_leaf_version(node: &ScriptTreeNode, target_hash: &[u8; 32]) -> Option<u8> {
-    match node {
-        ScriptTreeNode::Leaf {
-            script,
-            leaf_version,
-        } => {
-            if tap_leaf_hash(script, *leaf_version) == *target_hash {
-                Some(*leaf_version)
-            } else {
-                None
-            }
-        }
-        ScriptTreeNode::Branch(left, right) => {
-            find_leaf_version(left, target_hash).or_else(|| find_leaf_version(right, target_hash))
-        }
-    }
-}
-
 /// Verify a P2MR control block against a leaf hash and expected merkle root.
 ///
 /// Walks the merkle path in the control block, combining with `tap_branch_hash`
@@ -389,6 +367,133 @@ mod tests {
         );
     }
 
+    /// Test against BIP-360 spec vector: p2mr_simple_lightning_contract
+    #[test]
+    fn test_simple_lightning_contract() {
+        let script0 = hex::decode(
+            "029000b275209997a497d964fc1a62885b05a51166a65a90df00492c8d7cf61d6accf54803beac",
+        )
+        .unwrap();
+        let script1 = hex::decode(
+            "a8206c60f404f8167a38fc70eaf8aa17ac351023bef86bcb9d1086a19afe95bd533388204edfcf9dfe6c0b5c83d1ab3f78d1b39a46ebac6798e08e19761f5ed89ec83c10ac",
+        )
+        .unwrap();
+
+        let tree = ScriptTreeNode::Branch(
+            Box::new(ScriptTreeNode::Leaf {
+                script: script0,
+                leaf_version: 192,
+            }),
+            Box::new(ScriptTreeNode::Leaf {
+                script: script1,
+                leaf_version: 192,
+            }),
+        );
+        let info = build_p2mr_tree(&tree);
+
+        assert_eq!(
+            hex::encode(info.merkle_root),
+            "41646f8c1fe2a96ddad7f5471bc4fee7da98794ef8c45a4f4fc6a559d60c9f6b"
+        );
+
+        // Verify leaf hashes
+        let leaf_hashes: Vec<String> = info
+            .leaves
+            .iter()
+            .map(|l| hex::encode(l.leaf_hash))
+            .collect();
+        assert!(leaf_hashes.contains(
+            &"c81451874bd9ebd4b6fd4bba1f84cdfb533c532365d22a0a702205ff658b17c9".to_string()
+        ));
+        assert!(leaf_hashes.contains(
+            &"632c8632b4f29c6291416e23135cf78ecb82e525788ea5ed6483e3c6ce943b42".to_string()
+        ));
+
+        assert_control_blocks_match(
+            &info,
+            &[
+                "c1c81451874bd9ebd4b6fd4bba1f84cdfb533c532365d22a0a702205ff658b17c9",
+                "c1632c8632b4f29c6291416e23135cf78ecb82e525788ea5ed6483e3c6ce943b42",
+            ],
+        );
+
+        // Verify scriptPubKey and address
+        let spk = build_p2mr_script_pubkey(&info.merkle_root);
+        assert_eq!(
+            hex::encode(spk.as_bytes()),
+            "522041646f8c1fe2a96ddad7f5471bc4fee7da98794ef8c45a4f4fc6a559d60c9f6b"
+        );
+    }
+
+    /// Test against BIP-360 spec vector: p2mr_three_leaf_alternative
+    #[test]
+    fn test_three_leaf_alternative() {
+        let script_a =
+            hex::decode("2071981521ad9fc9036687364118fb6ccd2035b96a423c59c5430e98310a11abe2ac")
+                .unwrap();
+        let script_b =
+            hex::decode("20d5094d2dbe9b76e2c245a2b89b6006888952e2faa6a149ae318d69e520617748ac")
+                .unwrap();
+        let script_c =
+            hex::decode("20c440b462ad48c7a77f94cd4532d8f2119dcebbd7c9764557e62726419b08ad4cac")
+                .unwrap();
+
+        let tree = ScriptTreeNode::Branch(
+            Box::new(ScriptTreeNode::Leaf {
+                script: script_a,
+                leaf_version: 192,
+            }),
+            Box::new(ScriptTreeNode::Branch(
+                Box::new(ScriptTreeNode::Leaf {
+                    script: script_b,
+                    leaf_version: 192,
+                }),
+                Box::new(ScriptTreeNode::Leaf {
+                    script: script_c,
+                    leaf_version: 192,
+                }),
+            )),
+        );
+        let info = build_p2mr_tree(&tree);
+
+        assert_eq!(
+            hex::encode(info.merkle_root),
+            "2f6b2c5397b6d68ca18e09a3f05161668ffe93a988582d55c6f07bd5b3329def"
+        );
+
+        // Verify leaf hashes
+        let leaf_hashes: Vec<String> = info
+            .leaves
+            .iter()
+            .map(|l| hex::encode(l.leaf_hash))
+            .collect();
+        assert!(leaf_hashes.contains(
+            &"f154e8e8e17c31d3462d7132589ed29353c6fafdb884c5a6e04ea938834f0d9d".to_string()
+        ));
+        assert!(leaf_hashes.contains(
+            &"737ed1fe30bc42b8022d717b44f0d93516617af64a64753b7a06bf16b26cd711".to_string()
+        ));
+        assert!(leaf_hashes.contains(
+            &"d7485025fceb78b9ed667db36ed8b8dc7b1f0b307ac167fa516fe4352b9f4ef7".to_string()
+        ));
+
+        assert_control_blocks_match(
+            &info,
+            &[
+                "c13cd369a528b326bc9d2133cbd2ac21451acb31681a410434672c8e34fe757e91",
+                "c1737ed1fe30bc42b8022d717b44f0d93516617af64a64753b7a06bf16b26cd711f154e8e8e17c31d3462d7132589ed29353c6fafdb884c5a6e04ea938834f0d9d",
+                "c1d7485025fceb78b9ed667db36ed8b8dc7b1f0b307ac167fa516fe4352b9f4ef7f154e8e8e17c31d3462d7132589ed29353c6fafdb884c5a6e04ea938834f0d9d",
+            ],
+        );
+
+        // Verify scriptPubKey
+        let spk = build_p2mr_script_pubkey(&info.merkle_root);
+        assert_eq!(
+            hex::encode(spk.as_bytes()),
+            "52202f6b2c5397b6d68ca18e09a3f05161668ffe93a988582d55c6f07bd5b3329def"
+        );
+    }
+
     /// Test against BIP-360 spec vector: p2mr_three_leaf_complex
     /// Tree structure: [A, [B, C]]
     #[test]
diff --git a/packages/wasm-utxo/test/fixtures/address/bitcoin.json b/packages/wasm-utxo/test/fixtures/address/bitcoin.json
@@ -9,11 +9,7 @@
     "00141e231c7f9b3415daaa53ee5a7e12e120f00ec212",
     "bc1qrc33clumxs2a42jnaed8uyhpyrcqassje6kugr"
   ],
-  [
-    "p2sh",
-    "a91411510d2560794b3ed7bf734bc0e030e70e4db42d87",
-    "33GaWMLihvhqWJ9YNd7xo8diSh6otdfi5w"
-  ],
+  ["p2sh", "a91411510d2560794b3ed7bf734bc0e030e70e4db42d87", "33GaWMLihvhqWJ9YNd7xo8diSh6otdfi5w"],
   [
     "p2shP2wsh",
     "a9140c4e25aa3282fa35888f5e1eedb876265328312587",
@@ -44,11 +40,7 @@
     "001491b8f56f155030f74259be43dff4d94a6258d84a",
     "bc1qjxu02mc42qc0wsjehepalaxeff393kz2mzwu70"
   ],
-  [
-    "p2sh",
-    "a914d640bad0fafe2eeac9fc0e0f09fb899066263ebf87",
-    "3MDt56c1ME4Vi8aYP9DJdhWngxo3yDxGJb"
-  ],
+  ["p2sh", "a914d640bad0fafe2eeac9fc0e0f09fb899066263ebf87", "3MDt56c1ME4Vi8aYP9DJdhWngxo3yDxGJb"],
   [
     "p2shP2wsh",
     "a914696cab5f237c954fc1fade8c6b234fe93e0e80f287",
@@ -79,11 +71,7 @@
     "00140a058aec7588fca80070436b020c352c2891b680",
     "bc1qpgzc4mr43r72sqrsgd4syrp49s5frd5qq7et93"
   ],
-  [
-    "p2sh",
-    "a914056f5a1c07fe38d27e554b88bd857f64bda8eb6f87",
-    "32BkhNhA8a7byNxgz6RR3jrG8T1mCXqiC8"
-  ],
+  ["p2sh", "a914056f5a1c07fe38d27e554b88bd857f64bda8eb6f87", "32BkhNhA8a7byNxgz6RR3jrG8T1mCXqiC8"],
   [
     "p2shP2wsh",
     "a914f7db4f654f1211a63165cfdaf1170e96d433bc1387",
@@ -114,11 +102,7 @@
     "00145a8451539186feb4578b4f5613df6991e3078230",
     "bc1qt2z9z5u3smltg4utfatp8hmfj83s0q3s62ygc4"
   ],
-  [
-    "p2sh",
-    "a9149829fb41e3c2bcf6a164310a7acbf0adcc0c3ee187",
-    "3FZavdFukrNR5aJcL7dmbc18tZ9tBkVtqK"
-  ],
+  ["p2sh", "a9149829fb41e3c2bcf6a164310a7acbf0adcc0c3ee187", "3FZavdFukrNR5aJcL7dmbc18tZ9tBkVtqK"],
   [
     "p2shP2wsh",
     "a91493f1dd87104175795a1e37f5245461827237a05787",
@@ -138,5 +122,20 @@
     "p2trMusig2",
     "512085078b6ce45af8c4dea63248a5fae283d8edcca75f186395b237706c4bb42a36",
     "bc1ps5rckm8yttuvfh4xxfy2t7hzs0vwmn98tuvx89djxacxcja59gmq8lx0zu"
+  ],
+  [
+    "p2mr",
+    "5220c525714a7f49c28aedbbba78c005931a81c234b2f6c99a73e4d06082adc8bf2b",
+    "bc1zc5jhzjnlf8pg4mdmhfuvqpvnr2quyd9j7mye5uly6psg9twghu4ssr0v9k"
+  ],
+  [
+    "p2mr",
+    "5220ab179431c28d3b68fb798957faf5497d69c883c6fb1e1cd9f81483d87bac90cc",
+    "bc1z4vtegvwz35ak37me39tl4a2f045u3q7xlv0pek0czjpas7avjrxqz20g2y"
+  ],
+  [
+    "p2mr",
+    "5220ccbd66c6f7e8fdab47b3a486f59d28262be857f30d4773f2d5ea47f7761ce0e2",
+    "bc1zej7kd3hhar76k3an5jr0t8fgyc47s4lnp4rh8uk4afrlwasuur3qzgewqq"
   ]
 ]
diff --git a/packages/wasm-utxo/test/fixtures/address/bitcoinPublicSignet.json b/packages/wasm-utxo/test/fixtures/address/bitcoinPublicSignet.json
@@ -9,11 +9,7 @@
     "00141e231c7f9b3415daaa53ee5a7e12e120f00ec212",
     "tb1qrc33clumxs2a42jnaed8uyhpyrcqassjnud0ns"
   ],
-  [
-    "p2sh",
-    "a91411510d2560794b3ed7bf734bc0e030e70e4db42d87",
-    "2Mtpna6GkKPDBi5n63kjqR5cyf3JybU8XmH"
-  ],
+  ["p2sh", "a91411510d2560794b3ed7bf734bc0e030e70e4db42d87", "2Mtpna6GkKPDBi5n63kjqR5cyf3JybU8XmH"],
   [
     "p2shP2wsh",
     "a9140c4e25aa3282fa35888f5e1eedb876265328312587",
@@ -44,11 +40,7 @@
     "001491b8f56f155030f74259be43dff4d94a6258d84a",
     "tb1qjxu02mc42qc0wsjehepalaxeff393kz23y409u"
   ],
-  [
-    "p2sh",
-    "a914d640bad0fafe2eeac9fc0e0f09fb899066263ebf87",
-    "2NCn68qY2xgZquvD64GqBFeW3uK1Dqn2PA4"
-  ],
+  ["p2sh", "a914d640bad0fafe2eeac9fc0e0f09fb899066263ebf87", "2NCn68qY2xgZquvD64GqBFeW3uK1Dqn2PA4"],
   [
     "p2shP2wsh",
     "a914696cab5f237c954fc1fade8c6b234fe93e0e80f287",
@@ -79,11 +71,7 @@
     "00140a058aec7588fca80070436b020c352c2891b680",
     "tb1qpgzc4mr43r72sqrsgd4syrp49s5frd5q2czc7z"
   ],
-  [
-    "p2sh",
-    "a914056f5a1c07fe38d27e554b88bd857f64bda8eb6f87",
-    "2Msjxm7dBk2cxBAbEfE3HfgqXLoDw5wBFpQ"
-  ],
+  ["p2sh", "a914056f5a1c07fe38d27e554b88bd857f64bda8eb6f87", "2Msjxm7dBk2cxBAbEfE3HfgqXLoDw5wBFpQ"],
   [
     "p2shP2wsh",
     "a914f7db4f654f1211a63165cfdaf1170e96d433bc1387",
@@ -114,11 +102,7 @@
     "00145a8451539186feb4578b4f5613df6991e3078230",
     "tb1qt2z9z5u3smltg4utfatp8hmfj83s0q3ssvlmrx"
   ],
-  [
-    "p2sh",
-    "a9149829fb41e3c2bcf6a164310a7acbf0adcc0c3ee187",
-    "2N77nzNBwNJsmHMwA1FFeDYzQ6uN3yhET3E"
-  ],
+  ["p2sh", "a9149829fb41e3c2bcf6a164310a7acbf0adcc0c3ee187", "2N77nzNBwNJsmHMwA1FFeDYzQ6uN3yhET3E"],
   [
     "p2shP2wsh",
     "a91493f1dd87104175795a1e37f5245461827237a05787",
@@ -138,5 +122,20 @@
     "p2trMusig2",
     "512085078b6ce45af8c4dea63248a5fae283d8edcca75f186395b237706c4bb42a36",
     "tb1ps5rckm8yttuvfh4xxfy2t7hzs0vwmn98tuvx89djxacxcja59gmqshsqcn"
+  ],
+  [
+    "p2mr",
+    "5220c525714a7f49c28aedbbba78c005931a81c234b2f6c99a73e4d06082adc8bf2b",
+    "tb1zc5jhzjnlf8pg4mdmhfuvqpvnr2quyd9j7mye5uly6psg9twghu4s8terle"
+  ],
+  [
+    "p2mr",
+    "5220ab179431c28d3b68fb798957faf5497d69c883c6fb1e1cd9f81483d87bac90cc",
+    "tb1z4vtegvwz35ak37me39tl4a2f045u3q7xlv0pek0czjpas7avjrxq4ze8st"
+  ],
+  [
+    "p2mr",
+    "5220ccbd66c6f7e8fdab47b3a486f59d28262be857f30d4773f2d5ea47f7761ce0e2",
+    "tb1zej7kd3hhar76k3an5jr0t8fgyc47s4lnp4rh8uk4afrlwasuur3q4q0p60"
   ]
-]
+]
diff --git a/packages/wasm-utxo/test/fixtures/address/testnet.json b/packages/wasm-utxo/test/fixtures/address/testnet.json
diff --git a/packages/wasm-utxo/test/p2mr.ts b/packages/wasm-utxo/test/p2mr.ts
