BIP360: Pay to Merkle Root (P2MR) (#1670)

Review comments and assistance by:
  Armin Sabouri <armins88@gmail.com>
  D++ <82842780+dplusplus1024@users.noreply.github.com>
  Jameson Lopp <jameson.lopp@gmail.com>
  jbride <jbride2001@yahoo.com>
  Joey Yandle <xoloki@gmail.com>
  Jon Atack <jon@atack.com>
  Jonas Nick <jonasd.nick@gmail.com>
  Kyle Crews <kylecrews@Kyles-Mac-Studio.local>
  Mark "Murch" Erhardt <murch@murch.one>
  notmike-5 <notmike-5@users.noreply.github.com>
  Vojtěch Strnad <43024885+vostrnad@users.noreply.github.com>

Co-authored-by: Ethan Heilman <ethan.r.heilman@gmail.com>
Co-authored-by: Isabel Foxen Duke <110147802+Isabelfoxenduke@users.noreply.github.com>

bip-0360/ref-impl/rust/src/data_structures.rs

@@ -0,0 +1,595 @@
+use std::collections::HashMap;
+use serde::{Deserialize, Serialize};
+use log::debug;
+
+// Add imports for the unified keypair
+use bitcoin::secp256k1::{SecretKey, XOnlyPublicKey};
+use bitcoinpqc::{KeyPair, Algorithm};
+
+/// Enum representing the type of leaf script to create
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum LeafScriptType {
+    /// Script requires only SLH-DSA signature
+    SlhDsaOnly,
+    /// Script requires only Schnorr signature  
+    SchnorrOnly,
+    /// Script requires both Schnorr and SLH-DSA signatures (in that order)
+    ConcatenatedSchnorrAndSlhDsaSameLeaf,
+    /// Leaves of TapTree are mixed.  Some leaves are locked using Schnorr and others are locked using SLH-DSA
+    Mixed,
+    /// Script type is not applicable
+    NotApplicable,
+}
+
+impl LeafScriptType {
+    /// Check if this script type uses SLH-DSA
+    pub fn uses_slh_dsa(&self) -> bool {
+        matches!(self, LeafScriptType::SlhDsaOnly | LeafScriptType::ConcatenatedSchnorrAndSlhDsaSameLeaf)
+    }
+
+    /// Check if this script type uses Schnorr
+    pub fn uses_schnorr(&self) -> bool {
+        matches!(self, LeafScriptType::SchnorrOnly | LeafScriptType::ConcatenatedSchnorrAndSlhDsaSameLeaf)
+    }
+
+    /// Check if this script type requires both signature types
+    pub fn requires_both(&self) -> bool {
+        matches!(self, LeafScriptType::ConcatenatedSchnorrAndSlhDsaSameLeaf)
+    }
+
+    /// Check if TapTree uses Schnorr for some leaves and SLH-DSA for others
+    pub fn uses_mixed(&self) -> bool {
+        matches!(self, LeafScriptType::Mixed)
+    }
+
+    /// Check if this script type is not applicable
+    pub fn is_not_applicable(&self) -> bool {
+        matches!(self, LeafScriptType::NotApplicable)
+    }
+
+    /// Convert to string representation for serialization
+    pub fn to_string(&self) -> String {
+        match self {
+            LeafScriptType::SlhDsaOnly => "SLH_DSA_ONLY".to_string(),
+            LeafScriptType::SchnorrOnly => "SCHNORR_ONLY".to_string(),
+            LeafScriptType::ConcatenatedSchnorrAndSlhDsaSameLeaf => "CONCATENATED_SCHNORR_AND_SLH_DSA".to_string(),
+            LeafScriptType::Mixed => "MIXED".to_string(),
+            LeafScriptType::NotApplicable => "NOT_APPLICABLE".to_string(),
+        }
+    }
+
+    /// Parse from string representation
+    pub fn from_string(s: &str) -> Self {
+        match s {
+            "SLH_DSA_ONLY" => LeafScriptType::SlhDsaOnly,
+            "SCHNORR_ONLY" => LeafScriptType::SchnorrOnly,
+            "CONCATENATED_SCHNORR_AND_SLH_DSA" => LeafScriptType::ConcatenatedSchnorrAndSlhDsaSameLeaf,
+            "MIXED" => LeafScriptType::Mixed,
+            _ => LeafScriptType::NotApplicable,
+        }
+    }
+}
+
+#[derive(Debug, Serialize)]
+pub struct TestVectors {
+    pub version: u32,
+    #[serde(rename = "test_vectors")]
+    pub test_vectors: Vec<TestVector>,
+    #[serde(skip, default = "HashMap::new")]
+    pub test_vector_map: HashMap<String, TestVector>,
+}
+
+impl<'de> Deserialize<'de> for TestVectors {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        #[derive(Deserialize)]
+        struct Helper {
+            version: u32,
+            #[serde(rename = "test_vectors")]
+            test_vectors: Vec<TestVector>,
+        }
+
+        let helper = Helper::deserialize(deserializer)?;
+        
+        let mut test_vector_map = HashMap::new();
+        for test in helper.test_vectors.iter() {
+            test_vector_map.insert(test.id.clone(), test.clone());
+        }
+
+        Ok(TestVectors {
+            version: helper.version,
+            test_vectors: helper.test_vectors,
+            test_vector_map,
+        })
+    }
+}
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct TestVector {
+    pub id: String,
+    pub objective: String,
+    pub given: TestVectorGiven,
+    pub intermediary: TestVectorIntermediary,
+    pub expected: TestVectorExpected,
+}
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct TestVectorGiven {
+
+    #[serde(rename = "internalPubkey")]
+    pub internal_pubkey: Option<String>,
+
+    #[serde(rename = "scriptTree")]
+    pub script_tree: Option<TVScriptTree>,
+
+    #[serde(rename = "scriptInputs")]
+    pub script_inputs: Option<Vec<String>>,
+    #[serde(rename = "scriptHex")]
+    pub script_hex: Option<String>,
+    #[serde(rename = "controlBlock")]
+    pub control_block: Option<String>,
+}
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct TestVectorIntermediary {
+
+    #[serde(default)]
+    #[serde(rename = "leafHashes")]
+    pub leaf_hashes: Vec<String>,
+    #[serde(rename = "merkleRoot")]
+    pub merkle_root: Option<String>
+}
+
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct TestVectorExpected {
+    #[serde(rename = "scriptPubKey")]
+    pub script_pubkey: Option<String>,
+    #[serde(rename = "bip350Address")]
+    pub bip350_address: Option<String>,
+    #[serde(default)]
+    #[serde(rename = "scriptPathControlBlocks")]
+    pub script_path_control_blocks: Option<Vec<String>>,
+    #[serde(rename = "error")]
+    pub error: Option<String>,
+    #[serde(rename = "address")]
+    pub address: Option<String>,
+    #[serde(default)]
+    pub witness: Option<String>
+}
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct TVScriptLeaf {
+    pub id: u8,
+    pub script: String,
+    #[serde(rename = "leafVersion")]
+    pub leaf_version: u8,
+}
+
+// Taproot script trees are binary trees, so each branch should have exactly two children.
+#[derive(Debug, Serialize, Clone)]
+pub enum TVScriptTree {
+    Leaf(TVScriptLeaf),
+    Branch {
+
+        // Box is used because Rust needs to know the exact size of types at compile time.
+        // Without it, we'd have an infinitely size recursive type.
+        // The enum itself is on the stack, but the Box fields within the Branch variant store pointers to heap-allocated ScriptTree values.
+        left: Box<TVScriptTree>,
+        right: Box<TVScriptTree>,
+    },
+}
+
+// Add custom deserialize implementation
+impl<'de> Deserialize<'de> for TVScriptTree {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        #[derive(Deserialize)]
+        #[serde(untagged)]
+        enum Helper {
+            Leaf(TVScriptLeaf),
+            Branch(Vec<TVScriptTree>),
+        }
+
+        match Helper::deserialize(deserializer)? {
+            Helper::Leaf(leaf) => Ok(TVScriptTree::Leaf(leaf)),
+            Helper::Branch(v) => {
+                assert!(v.len() == 2, "Branch must have exactly two children");
+                let mut iter = v.into_iter();
+                Ok(TVScriptTree::Branch {
+                    left: Box::new(iter.next().unwrap()),
+                    right: Box::new(iter.next().unwrap()),
+                })
+            }
+        }
+    }
+}
+
+// Add this enum before the TVScriptTree implementation
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
+pub enum Direction {
+    Left,
+    Right,
+    Root,
+}
+
+impl TVScriptTree {
+    /// Implements a "post-order" traversal as follows: left, right, branch
+    pub fn traverse_with_depth<F: FnMut(&TVScriptTree, u8, Direction)>(&self, depth: u8, direction: Direction, f: &mut F) {
+        match self {
+            TVScriptTree::Branch { left, right } => {
+
+                right.traverse_with_depth(depth, Direction::Right, f); // Pass Right for right subtree
+                left.traverse_with_depth(depth, Direction::Left, f);   // Pass Left for left subtree
+                f(self, depth, direction);  // Pass the current node's direction
+            }
+            TVScriptTree::Leaf { .. } => {
+                f(self, depth, direction);
+            }
+        }
+    }
+
+    /// Traverses the tree visiting right subtree leaves first, then left subtree leaves.
+    /// Depth increases by 1 at each branch level.
+    /*
+             root (depth 0)
+            /            \
+        L0 (depth 1)    (subtree) (depth 1)
+                            /          \
+                    L1 (depth 2)    L2 (depth 2)
+
+        The new traversal will visit:
+            L1 at depth 2 -> L2 at depth 2 -> L0 at depth 1
+     */
+    pub fn traverse_with_right_subtree_first<F: FnMut(&TVScriptTree, u8, Direction)>(&self, depth: u8, direction: Direction, f: &mut F) {
+        match self {
+            TVScriptTree::Branch { left, right } => {
+                let next_depth = depth + 1;
+                // Visit right subtree first
+                right.traverse_with_right_subtree_first(next_depth, Direction::Right, f);
+                // Then visit left subtree
+                left.traverse_with_right_subtree_first(next_depth, Direction::Left, f);
+            }
+            TVScriptTree::Leaf { .. } => {
+                f(self, depth, direction);
+            }
+        }
+    }
+}
+
+impl std::fmt::Display for Direction {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        match self {
+            Direction::Left => write!(f, "L"),
+            Direction::Right => write!(f, "R"),
+            Direction::Root => write!(f, "Root"),
+        }
+    }
+}
+
+pub struct ScriptTreeHashCache {
+    pub leaf_hashes: HashMap<String, String>,
+    pub branch_hashes: HashMap<u8, String>,
+}
+
+impl ScriptTreeHashCache {
+    pub fn new() -> Self {
+        Self {
+            leaf_hashes: HashMap::new(),
+            branch_hashes: HashMap::new(),
+        }
+    }
+
+    pub fn set_leaf_hash(&mut self, branch_id: u8, direction: Direction, hash: String) {
+        let key = format!("{branch_id}_{direction}");
+        debug!("set_leaf_hash: key: {}, hash: {}", key, hash);
+        self.leaf_hashes.insert(key, hash);
+    }
+
+    pub fn set_branch_hash(&mut self, branch_id: u8, hash: String) {
+        self.branch_hashes.insert(branch_id, hash);
+    }
+}
+
+fn serialize_hex<S>(bytes: &Vec<u8>, s: S) -> Result<S::Ok, S::Error>
+where
+    S: serde::Serializer,
+{
+    s.serialize_str(&hex::encode(bytes))
+}
+
+fn deserialize_hex<'de, D>(d: D) -> Result<Vec<u8>, D::Error>
+where
+    D: serde::Deserializer<'de>,
+{
+    let s = String::deserialize(d)?;
+    hex::decode(s).map_err(serde::de::Error::custom)
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct SpendDetails {
+    pub tx_hex: String,
+    #[serde(serialize_with = "serialize_hex")]
+    #[serde(deserialize_with = "deserialize_hex")]
+    pub sighash: Vec<u8>,
+    #[serde(serialize_with = "serialize_hex")]
+    #[serde(deserialize_with = "deserialize_hex")]
+    pub sig_bytes: Vec<u8>,
+    #[serde(serialize_with = "serialize_hex")]
+    #[serde(deserialize_with = "deserialize_hex")]
+    pub derived_witness_vec: Vec<u8>,
+}
+
+impl std::process::Termination for SpendDetails {
+    fn report(self) -> std::process::ExitCode {
+        if let Ok(json) = serde_json::to_string_pretty(&self) {
+            println!("{}", json);
+        } else {
+            println!("{:?}", self);
+        }
+        std::process::ExitCode::SUCCESS
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct UtxoReturn {
+
+    pub script_pubkey_hex: String,
+    pub bech32m_address: String,
+    pub bitcoin_network: bitcoin::Network,
+}
+
+impl std::process::Termination for UtxoReturn {
+    fn report(self) -> std::process::ExitCode {
+        if let Ok(json) = serde_json::to_string_pretty(&self) {
+            println!("{}", json);
+        } else {
+            println!("{:?}", self);
+        }
+        std::process::ExitCode::SUCCESS
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct TaptreeReturn {
+    pub leaf_script_priv_keys_hex: Vec<String>, // Changed to support multiple private keys
+    pub leaf_script_hex: String,
+    pub tree_root_hex: String,
+    pub control_block_hex: String,
+    /// The script type of the leaf being returned (needed for spending)
+    pub leaf_script_type: String,
+}
+
+impl std::process::Termination for TaptreeReturn {
+    fn report(self) -> std::process::ExitCode {
+        if let Ok(json) = serde_json::to_string_pretty(&self) {
+            println!("{}", json);
+        } else {
+            println!("{:?}", self);
+        }
+        std::process::ExitCode::SUCCESS
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ConstructionReturn {
+    pub taptree_return: TaptreeReturn,
+    pub utxo_return: UtxoReturn,
+}
+
+impl std::process::Termination for ConstructionReturn {
+    fn report(self) -> std::process::ExitCode {
+        if let Ok(json) = serde_json::to_string_pretty(&self) {
+            println!("{}", json);
+        } else {
+            println!("{:?}", self);
+        }
+        std::process::ExitCode::SUCCESS
+    }
+}
+
+/// A unified keypair that can contain either a Schnorr keypair or an SLH-DSA keypair
+#[derive(Debug, Clone)]
+pub enum UnifiedKeypair {
+    Schnorr(SecretKey, XOnlyPublicKey),
+    SlhDsa(KeyPair),
+}
+
+/// A container for multiple keypairs that can be used in a single leaf script
+#[derive(Debug, Clone)]
+pub struct MultiKeypair {
+    pub schnorr_keypair: Option<UnifiedKeypair>,
+    pub slh_dsa_keypair: Option<UnifiedKeypair>,
+}
+
+impl MultiKeypair {
+    /// Create a new MultiKeypair with only a Schnorr keypair
+    pub fn new_schnorr_only(schnorr_keypair: UnifiedKeypair) -> Self {
+        Self {
+            schnorr_keypair: Some(schnorr_keypair),
+            slh_dsa_keypair: None,
+        }
+    }
+
+    /// Create a new MultiKeypair with only an SLH-DSA keypair
+    pub fn new_slh_dsa_only(slh_dsa_keypair: UnifiedKeypair) -> Self {
+        Self {
+            schnorr_keypair: None,
+            slh_dsa_keypair: Some(slh_dsa_keypair),
+        }
+    }
+
+    /// Create a new MultiKeypair with both keypairs
+    pub fn new_combined(schnorr_keypair: UnifiedKeypair, slh_dsa_keypair: UnifiedKeypair) -> Self {
+        Self {
+            schnorr_keypair: Some(schnorr_keypair),
+            slh_dsa_keypair: Some(slh_dsa_keypair),
+        }
+    }
+
+    /// Get all secret key bytes for serialization (in order: schnorr, then slh_dsa if present)
+    pub fn secret_key_bytes(&self) -> Vec<Vec<u8>> {
+        let mut result = Vec::new();
+        if let Some(ref schnorr) = self.schnorr_keypair {
+            result.push(schnorr.secret_key_bytes());
+        }
+        if let Some(ref slh_dsa) = self.slh_dsa_keypair {
+            result.push(slh_dsa.secret_key_bytes());
+        }
+        result
+    }
+
+    /// Get all public key bytes for script construction (in order: schnorr, then slh_dsa if present)
+    pub fn public_key_bytes(&self) -> Vec<Vec<u8>> {
+        let mut result = Vec::new();
+        if let Some(ref schnorr) = self.schnorr_keypair {
+            result.push(schnorr.public_key_bytes());
+        }
+        if let Some(ref slh_dsa) = self.slh_dsa_keypair {
+            result.push(slh_dsa.public_key_bytes());
+        }
+        result
+    }
+
+    /// Check if this contains a Schnorr keypair
+    pub fn has_schnorr(&self) -> bool {
+        self.schnorr_keypair.is_some()
+    }
+
+    /// Check if this contains an SLH-DSA keypair
+    pub fn has_slh_dsa(&self) -> bool {
+        self.slh_dsa_keypair.is_some()
+    }
+
+    /// Get the Schnorr keypair if present
+    pub fn schnorr_keypair(&self) -> Option<&UnifiedKeypair> {
+        self.schnorr_keypair.as_ref()
+    }
+
+    /// Get the SLH-DSA keypair if present
+    pub fn slh_dsa_keypair(&self) -> Option<&UnifiedKeypair> {
+        self.slh_dsa_keypair.as_ref()
+    }
+}
+
+/// Information about a single leaf in a mixed-type tree
+/// Used when different leaves in the same tree use different algorithms
+#[derive(Debug, Clone)]
+pub struct MixedLeafInfo {
+    /// The leaf index in the tree
+    pub leaf_index: u32,
+    /// The script type for this specific leaf
+    pub leaf_script_type: LeafScriptType,
+    /// The keypairs for this leaf
+    pub keypairs: MultiKeypair,
+    /// The script for this leaf
+    pub script: Vec<u8>,
+}
+
+impl MixedLeafInfo {
+    /// Create a new MixedLeafInfo for a Schnorr-only leaf
+    pub fn new_schnorr(leaf_index: u32, keypairs: MultiKeypair, script: Vec<u8>) -> Self {
+        Self {
+            leaf_index,
+            leaf_script_type: LeafScriptType::SchnorrOnly,
+            keypairs,
+            script,
+        }
+    }
+
+    /// Create a new MixedLeafInfo for an SLH-DSA-only leaf
+    pub fn new_slh_dsa(leaf_index: u32, keypairs: MultiKeypair, script: Vec<u8>) -> Self {
+        Self {
+            leaf_index,
+            leaf_script_type: LeafScriptType::SlhDsaOnly,
+            keypairs,
+            script,
+        }
+    }
+
+    /// Create a new MixedLeafInfo for a combined Schnorr+SLH-DSA leaf
+    pub fn new_combined(leaf_index: u32, keypairs: MultiKeypair, script: Vec<u8>) -> Self {
+        Self {
+            leaf_index,
+            leaf_script_type: LeafScriptType::ConcatenatedSchnorrAndSlhDsaSameLeaf,
+            keypairs,
+            script,
+        }
+    }
+
+    /// Get the secret key bytes for this leaf
+    pub fn secret_key_bytes(&self) -> Vec<Vec<u8>> {
+        self.keypairs.secret_key_bytes()
+    }
+
+    /// Get the public key bytes for this leaf
+    pub fn public_key_bytes(&self) -> Vec<Vec<u8>> {
+        self.keypairs.public_key_bytes()
+    }
+}
+
+impl UnifiedKeypair {
+    /// Create a new Schnorr keypair
+    pub fn new_schnorr(secret_key: SecretKey, public_key: XOnlyPublicKey) -> Self {
+        UnifiedKeypair::Schnorr(secret_key, public_key)
+    }
+
+    /// Create a new SLH-DSA keypair
+    pub fn new_slh_dsa(keypair: KeyPair) -> Self {
+        UnifiedKeypair::SlhDsa(keypair)
+    }
+
+    /// Get the secret key bytes for serialization
+    pub fn secret_key_bytes(&self) -> Vec<u8> {
+        match self {
+            UnifiedKeypair::Schnorr(secret_key, _) => secret_key.secret_bytes().to_vec(),
+            UnifiedKeypair::SlhDsa(keypair) => keypair.secret_key.bytes.clone(),
+        }
+    }
+
+    /// Get the public key bytes for script construction
+    pub fn public_key_bytes(&self) -> Vec<u8> {
+        match self {
+            UnifiedKeypair::Schnorr(_, public_key) => public_key.serialize().to_vec(),
+            UnifiedKeypair::SlhDsa(keypair) => keypair.public_key.bytes.clone(),
+        }
+    }
+
+    /// Get the algorithm type
+    pub fn algorithm(&self) -> &'static str {
+        match self {
+            UnifiedKeypair::Schnorr(_, _) => "Schnorr",
+            UnifiedKeypair::SlhDsa(_) => "SLH-DSA",
+        }
+    }
+
+    /// Check if this is a Schnorr keypair
+    pub fn is_schnorr(&self) -> bool {
+        matches!(self, UnifiedKeypair::Schnorr(_, _))
+    }
+
+    /// Check if this is an SLH-DSA keypair
+    pub fn is_slh_dsa(&self) -> bool {
+        matches!(self, UnifiedKeypair::SlhDsa(_))
+    }
+
+    /// Get the underlying Schnorr keypair if this is a Schnorr keypair
+    pub fn as_schnorr(&self) -> Option<(&SecretKey, &XOnlyPublicKey)> {
+        match self {
+            UnifiedKeypair::Schnorr(secret_key, public_key) => Some((secret_key, public_key)),
+            _ => None,
+        }
+    }
+
+    /// Get the underlying SLH-DSA keypair if this is an SLH-DSA keypair
+    pub fn as_slh_dsa(&self) -> Option<&KeyPair> {
+        match self {
+            UnifiedKeypair::SlhDsa(keypair) => Some(keypair),
+            _ => None,
+        }
+    }
+}