Merge pull request #1 from conduition/360/p2mr-compute-control-block

bip360: simplify computing control blocks using explicit traversal paths

bip-0360/ref-impl/common/tests/data/p2mr_construction.json

@@ -253,6 +253,51 @@
                     "c1737ed1fe30bc42b8022d717b44f0d93516617af64a64753b7a06bf16b26cd711f154e8e8e17c31d3462d7132589ed29353c6fafdb884c5a6e04ea938834f0d9d"
                 ]
             }
+        },
+        {
+            "id": "p2mr_duplicate_leaves",
+            "objective": "Ensure P2MR control blocks can be constructed correctly even when duplicate script leaves coexist in the same tree.",
+            "given": {
+                "scriptTree": [
+                    {
+                        "id": 0,
+                        "script": "2071981521ad9fc9036687364118fb6ccd2035b96a423c59c5430e98310a11abe2ac",
+                        "asm": "71981521ad9fc9036687364118fb6ccd2035b96a423c59c5430e98310a11abe2 OP_CHECKSIG",
+                        "leafVersion": 192
+                    },
+                    [
+                        {
+                            "id": 1,
+                            "script": "20d5094d2dbe9b76e2c245a2b89b6006888952e2faa6a149ae318d69e520617748ac",
+                            "asm": "d5094d2dbe9b76e2c245a2b89b6006888952e2faa6a149ae318d69e520617748 OP_CHECKSIG",
+                            "leafVersion": 192
+                        },
+                        {
+                            "id": 2,
+                            "script": "2071981521ad9fc9036687364118fb6ccd2035b96a423c59c5430e98310a11abe2ac",
+                            "asm": "71981521ad9fc9036687364118fb6ccd2035b96a423c59c5430e98310a11abe2 OP_CHECKSIG",
+                            "leafVersion": 192
+                        }
+                    ]
+                ]
+            },
+            "intermediary": {
+                "leafHashes": [
+                    "f154e8e8e17c31d3462d7132589ed29353c6fafdb884c5a6e04ea938834f0d9d",
+                    "737ed1fe30bc42b8022d717b44f0d93516617af64a64753b7a06bf16b26cd711",
+                    "f154e8e8e17c31d3462d7132589ed29353c6fafdb884c5a6e04ea938834f0d9d"
+                ],
+                "merkleRoot": "3ab4f80153012398a7df2df273b0bd3cdb839883320ff06d62e8cc1ecb351ba7"
+            },
+            "expected": {
+                "scriptPubKey": "52203ab4f80153012398a7df2df273b0bd3cdb839883320ff06d62e8cc1ecb351ba7",
+                "bip350Address": "bc1z8260sq2nqy3e3f7l9he88v9a8ndc8xyrxg8lqmtzarxpaje4rwnsxxww7h",
+                "scriptPathControlBlocks": [
+                    "c13a9adbb72f00022f76909e0f4bbab37dacafadb59618a4a13a123e13cb6f3019",
+                    "c1f154e8e8e17c31d3462d7132589ed29353c6fafdb884c5a6e04ea938834f0d9df154e8e8e17c31d3462d7132589ed29353c6fafdb884c5a6e04ea938834f0d9d",
+                    "c1737ed1fe30bc42b8022d717b44f0d93516617af64a64753b7a06bf16b26cd711f154e8e8e17c31d3462d7132589ed29353c6fafdb884c5a6e04ea938834f0d9d"
+                ]
+            }
         }
     ]
 }

bip-0360/ref-impl/python/p2mr.py

@@ -130,52 +130,30 @@ def compute_merkle_root(tree: ScriptTree) -> str:
         raise ValueError("Invalid tree node")
 
 
-def compute_control_block(
+def compute_control_block(path: int, tree: ScriptTree) -> bytes:
-    leaf: Dict[str, Any], tree: ScriptTree, path: Optional[List[str]] = None
+    """
-) -> Optional[str]:
+    Compute the control block for a script leaf at a given position in the script tree.
-    """Compute the control block for a given leaf in a given tree"""
+    The `path` argument encodes the position as follows.
-    if path is None:
-        path = []
 
+    Starting at depth zero, follow the branches of the script tree until reaching a leaf.
+    When we encounter a branch at any depth `d` (steps from the root), we look at the bit
+    `(path >> d) & 1` to decide whether to take the left or right branch.
+    """
     if isinstance(tree, dict):
-        if tree == leaf:
+        return bytes([tree["leafVersion"] | 1])
-            version_byte = (leaf["leafVersion"] | 1) & 0xFF
+    assert isinstance(tree, list) and len(tree) == 2
-            return f"{version_byte:02x}" + "".join(path)
 
-        return None
+    control_block = b""
 
-    if isinstance(tree, list):
+    while isinstance(tree, list):
-        for i, child in enumerate(tree):
+        assert len(tree) == 2
-            # build a list of sibling roots at this level
+        sibling = tree[(path & 1) ^ 1]
-            siblings: List[str] = []
+        tree = tree[(path & 1)]
-            for j, sib in enumerate(tree):
+        control_block = bytes.fromhex(compute_merkle_root(sibling)) + control_block
-                if j != i:
+        path >>= 1
-                    siblings.append(compute_merkle_root(sib))
-            # try this child; if it (or a descendant) matches, we get a result
-            result = compute_control_block(leaf, child, siblings + path)
-            if result:
-                return result
 
-    return None
+    assert isinstance(tree, dict)
-
+    return bytes([tree["leafVersion"] | 1]) + control_block
-
-def collect_control_blocks(script_tree: ScriptTree) -> List[str]:
-    """Return control blocks for all leaves in tree declaration order.
-    Note: This ordering is for testing purposes. In practice, you would
-    compute the control block for a specific leaf at spend-time using
-    `compute_control_block(leaf, tree)`."""
-    leaf_nodes: List[Dict[str, Any]] = []
-    stack = [script_tree]
-    while stack:
-        node = stack.pop()
-        if isinstance(node, dict):
-            leaf_nodes.append(node)
-        elif isinstance(node, list):
-            stack.extend(reversed(node))
-
-    return [
-        cb for leaf in leaf_nodes if (cb := compute_control_block(leaf, script_tree))
-    ]
 
 
 #
@@ -306,6 +284,26 @@ def encode(hrp, witver, witprog):
 #
 # BIP-360 Test Code
 #
+def walk_script_tree_paths(script_tree: ScriptTree, path: int = 0, depth: int = 0) -> List[int]:
+    """Walk through a script tree and produce a list of the bit-encoded traversal paths for each leaf.
+    Used for testing compute_control_block."""
+    if isinstance(script_tree, dict):
+        return [path]
+    assert isinstance(script_tree, list) and len(script_tree) == 2
+    lchild_paths = walk_script_tree_paths(script_tree[0], path, depth + 1)
+    rchild_paths = walk_script_tree_paths(script_tree[1], path | (1 << depth), depth + 1)
+    return lchild_paths + rchild_paths
+
+
+def collect_control_blocks(script_tree: ScriptTree) -> List[str]:
+    """Return control blocks for all leaves in tree declaration order.
+    Note: This ordering is for testing purposes. In practice, you would
+    compute the control block for a specific leaf at spend-time using
+    `compute_control_block(path, tree)`."""
+    leaf_node_paths: List[int] = walk_script_tree_paths(script_tree)
+    return [compute_control_block(path, script_tree).hex() for path in leaf_node_paths]
+
+
 def extract_test_data(v: Dict[str, Any]) -> Dict[str, Any]:
     """Extract test data from a test vector, returning None for missing keys"""
     given = v.get("given", {})