Skip to main content

taproot_assets_types/
mssmt.rs

1use crate::error::Error;
2use alloc::{string::ToString, vec::Vec};
3use bitcoin::hashes::{Hash, sha256::Hash as Sha256Hash};
4use bitcoin::io::Read;
5use serde::{Deserialize, Serialize};
6
7/// Represents a node in an MS-SMT (Merkle Sum Sparse Merkle Tree).
8#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
9pub struct MssmtNode {
10    /// The hash of the node.
11    pub hash: Sha256Hash,
12    /// The sum of the node.
13    pub sum: u64,
14}
15
16/// Represents a merkle proof for a MS-SMT.
17#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
18pub struct MssmtProof {
19    // Corresponds to mssmt.Proof
20    /// Siblings that should be hashed with the leaf and its parents to arrive at the root.
21    pub nodes: Vec<MssmtNode>,
22}
23
24impl MssmtProof {
25    pub fn decode_tlv<R: Read>(mut r: R) -> Result<Self, Error> {
26        // Decode the compressed MSSMT proof format:
27        // - 2 bytes: number of nodes (uint16, big endian)
28        // - For each node: 32 bytes hash + 8 bytes sum (uint64, big endian)
29        // - Packed bits (32 bytes for 255 bits max representing empty tree bits)
30
31        // Read number of nodes (uint16, big endian)
32        let mut num_nodes_bytes = [0u8; 2];
33        r.read_exact(&mut num_nodes_bytes).map_err(Error::Io)?;
34        let num_nodes = u16::from_be_bytes(num_nodes_bytes) as usize;
35
36        // Read the non-empty nodes
37        let mut explicit_nodes = Vec::with_capacity(num_nodes);
38        for _ in 0..num_nodes {
39            // Read 32-byte hash
40            let mut hash_bytes = [0u8; 32];
41            r.read_exact(&mut hash_bytes).map_err(Error::Io)?;
42            let hash = Sha256Hash::from_byte_array(hash_bytes);
43
44            // Read 8-byte sum (uint64, big endian)
45            let mut sum_bytes = [0u8; 8];
46            r.read_exact(&mut sum_bytes).map_err(Error::Io)?;
47            let sum = u64::from_be_bytes(sum_bytes);
48
49            explicit_nodes.push(MssmtNode { hash, sum });
50        }
51
52        // Read the packed bits (32 bytes for exactly 256 bits)
53        // MaxTreeLevels = 256, so MaxTreeLevels / 8 = 32 bytes
54        let mut packed_bits = [0u8; 32];
55        r.read_exact(&mut packed_bits).map_err(Error::Io)?;
56
57        // Unpack bits - matches Go's UnpackBits function
58        // Go uses little-endian bit order: LSB first within each byte
59        // Creates len(bytes)*8 = 32*8 = 256 bits
60        let mut bits = Vec::with_capacity(256);
61        for i in 0..256 {
62            let byte_idx = i / 8;
63            let bit_idx = i % 8;
64            let byte_val = packed_bits[byte_idx];
65            let bit_set = (byte_val >> bit_idx) & 1 == 1;
66            bits.push(bit_set);
67        }
68
69        // Reconstruct the full proof by combining explicit nodes with empty tree nodes
70        let mut nodes = Vec::new();
71        let mut explicit_node_idx = 0;
72
73        // Count how many bits are set to false for validation
74        let false_bits = bits.iter().filter(|&&b| !b).count();
75
76        // According to Go logic:
77        // - true bit means use empty tree node
78        // - false bit means use explicit node
79        // So false_bits should equal explicit_nodes.len()
80        if false_bits != explicit_nodes.len() {
81            return Err(Error::InvalidTlvValue(
82                0,
83                "Bit/node count mismatch: false bits != explicit nodes".to_string(),
84            ));
85        }
86
87        for (_level, bit_set) in bits.iter().enumerate() {
88            if *bit_set {
89                // Bit is set: use empty tree node (matches Go logic)
90                nodes.push(MssmtNode {
91                    hash: Sha256Hash::all_zeros(),
92                    sum: 0,
93                });
94            } else {
95                // Bit is not set: use explicit node (matches Go logic)
96                if explicit_node_idx >= explicit_nodes.len() {
97                    return Err(Error::InvalidTlvValue(
98                        0,
99                        "Insufficient explicit nodes for compressed proof".to_string(),
100                    ));
101                }
102                nodes.push(explicit_nodes[explicit_node_idx].clone());
103                explicit_node_idx += 1;
104            }
105        }
106
107        // Verify all explicit nodes were consumed
108        if explicit_node_idx != explicit_nodes.len() {
109            return Err(Error::InvalidTlvValue(
110                0,
111                "Too many explicit nodes for compressed proof".to_string(),
112            ));
113        }
114
115        Ok(MssmtProof { nodes })
116    }
117}