1#[cfg(feature = "utxo-commitments")]
7use crate::utxo_commitments::data_structures::{
8 UtxoCommitment, UtxoCommitmentError, UtxoCommitmentResult,
9};
10#[cfg(feature = "utxo-commitments")]
11use blvm_consensus::types::{Hash, Natural, OutPoint, UTXO};
12#[cfg(feature = "utxo-commitments")]
13use blvm_spec_lock::spec_locked;
14#[cfg(feature = "utxo-commitments")]
15use sha2::{Digest, Sha256};
16#[cfg(feature = "utxo-commitments")]
17use sparse_merkle_tree::default_store::DefaultStore;
18#[cfg(feature = "utxo-commitments")]
19use sparse_merkle_tree::traits::{Hasher, Value};
20#[cfg(feature = "utxo-commitments")]
21use sparse_merkle_tree::{SparseMerkleTree, H256};
22#[cfg(feature = "utxo-commitments")]
23use std::collections::HashMap;
24
25#[cfg(feature = "utxo-commitments")]
27#[derive(Default, Clone, Debug)]
28pub struct UtxoHasher {
29 hasher: Sha256,
30}
31
32#[cfg(feature = "utxo-commitments")]
33impl Hasher for UtxoHasher {
34 fn write_h256(&mut self, h: &H256) {
35 self.hasher.update(h.as_slice());
37 }
38
39 fn write_byte(&mut self, b: u8) {
40 self.hasher.update(&[b]);
41 }
42
43 fn finish(self) -> H256 {
44 let hash = self.hasher.finalize();
45 let mut bytes = [0u8; 32];
46 bytes.copy_from_slice(&hash);
47 H256::from(bytes)
48 }
49}
50
51#[cfg(feature = "utxo-commitments")]
53#[derive(Clone, Debug, PartialEq, Eq, Default)]
54pub struct UtxoValue {
55 pub data: Vec<u8>,
56}
57
58#[cfg(feature = "utxo-commitments")]
59impl Value for UtxoValue {
60 fn to_h256(&self) -> H256 {
61 let mut hasher = Sha256::new();
62 hasher.update(&self.data);
63 let hash = hasher.finalize();
64 let mut bytes = [0u8; 32];
65 bytes.copy_from_slice(&hash);
66 H256::from(bytes)
67 }
68
69 fn zero() -> Self {
70 Self { data: Vec::new() }
71 }
72}
73
74#[cfg(feature = "utxo-commitments")]
79pub struct UtxoMerkleTree {
80 tree: SparseMerkleTree<UtxoHasher, UtxoValue, DefaultStore<UtxoValue>>,
81 #[allow(dead_code)] utxo_index: HashMap<OutPoint, usize>,
83 total_supply: u64,
84 utxo_count: u64,
85}
86
87#[cfg(feature = "utxo-commitments")]
88impl UtxoMerkleTree {
89 pub fn new() -> UtxoCommitmentResult<Self> {
91 let store = DefaultStore::default();
92 let tree = SparseMerkleTree::new_with_store(store).map_err(|e| {
93 UtxoCommitmentError::MerkleTreeError(format!("Failed to create tree: {:?}", e))
94 })?;
95
96 Ok(Self {
97 tree,
98 utxo_index: HashMap::new(),
99 total_supply: 0,
100 utxo_count: 0,
101 })
102 }
103
104 pub fn root(&self) -> Hash {
106 let root_h256 = self.tree.root();
107 let mut hash = [0u8; 32];
108 hash.copy_from_slice(root_h256.as_slice());
109 hash
110 }
111
112 pub fn insert(&mut self, outpoint: OutPoint, utxo: UTXO) -> UtxoCommitmentResult<Hash> {
114 let key = self.hash_outpoint(&outpoint);
116
117 let value = self.serialize_utxo(&utxo)?;
119 let utxo_value = UtxoValue { data: value };
120
121 let root_h256 = self
123 .tree
124 .update(key, utxo_value)
125 .map_err(|e| UtxoCommitmentError::MerkleTreeError(format!("Update failed: {:?}", e)))?;
126
127 let old_supply = self.total_supply;
129 self.total_supply = self
130 .total_supply
131 .checked_add(utxo.value as u64)
132 .ok_or_else(|| {
133 UtxoCommitmentError::MerkleTreeError("Total supply overflow".to_string())
134 })?;
135 self.utxo_count = self.utxo_count.checked_add(1).ok_or_else(|| {
136 UtxoCommitmentError::MerkleTreeError("UTXO count overflow".to_string())
137 })?;
138
139 debug_assert!(
141 self.total_supply >= old_supply,
142 "Total supply ({}) must be >= previous supply ({})",
143 self.total_supply,
144 old_supply
145 );
146
147 let mut hash = [0u8; 32];
149 hash.copy_from_slice(root_h256.as_slice());
150 Ok(hash)
151 }
152
153 pub fn remove(&mut self, outpoint: &OutPoint, utxo: &UTXO) -> UtxoCommitmentResult<Hash> {
155 let key = self.hash_outpoint(outpoint);
157
158 let zero_value = UtxoValue::zero();
160
161 let root_h256 = self
163 .tree
164 .update(key, zero_value)
165 .map_err(|e| UtxoCommitmentError::MerkleTreeError(format!("Remove failed: {:?}", e)))?;
166
167 let old_supply = self.total_supply;
169 let old_count = self.utxo_count;
170
171 self.total_supply = self.total_supply.saturating_sub(utxo.value as u64);
172 self.utxo_count = self.utxo_count.saturating_sub(1);
173
174 debug_assert!(
176 self.total_supply <= old_supply,
177 "Total supply ({}) must be <= previous supply ({})",
178 self.total_supply,
179 old_supply
180 );
181
182 debug_assert!(
184 self.utxo_count <= old_count,
185 "UTXO count ({}) must be <= previous count ({})",
186 self.utxo_count,
187 old_count
188 );
189
190 debug_assert!(
192 true,
194 "Total supply ({}) must be within u64 bounds",
195 self.total_supply
196 );
197 debug_assert!(
198 true,
200 "UTXO count ({}) must be within u64 bounds",
201 self.utxo_count
202 );
203
204 let mut hash = [0u8; 32];
206 hash.copy_from_slice(root_h256.as_slice());
207 Ok(hash)
208 }
209
210 pub fn get(&self, outpoint: &OutPoint) -> UtxoCommitmentResult<Option<UTXO>> {
212 let key = self.hash_outpoint(outpoint);
213
214 match self.tree.get(&key) {
215 Ok(value) => {
216 if value.to_h256() == H256::zero() || value.to_h256() == UtxoValue::zero().to_h256()
218 {
219 Ok(None)
220 } else {
221 let serialized_data = &value.data;
223
224 match self.deserialize_utxo(serialized_data) {
226 Ok(utxo) => Ok(Some(utxo)),
227 Err(e) => {
228 Err(UtxoCommitmentError::InvalidUtxo(format!(
230 "Failed to deserialize UTXO: {}",
231 e
232 )))
233 }
234 }
235 }
236 }
237 Err(_) => Ok(None),
238 }
239 }
240
241 #[spec_locked("11.4")]
243 pub fn generate_commitment(&self, block_hash: Hash, block_height: Natural) -> UtxoCommitment {
244 let merkle_root = self.root();
245 UtxoCommitment::new(
246 merkle_root,
247 self.total_supply,
248 self.utxo_count,
249 block_height,
250 block_hash,
251 )
252 }
253
254 pub fn total_supply(&self) -> u64 {
256 self.total_supply
257 }
258
259 pub fn utxo_count(&self) -> u64 {
261 self.utxo_count
262 }
263
264 pub fn generate_proof(
268 &self,
269 outpoint: &OutPoint,
270 ) -> UtxoCommitmentResult<sparse_merkle_tree::MerkleProof> {
271 let key = self.hash_outpoint(outpoint);
272 let keys = vec![key];
273
274 self.tree.merkle_proof(keys).map_err(|e| {
275 UtxoCommitmentError::MerkleTreeError(format!("Failed to generate proof: {:?}", e))
276 })
277 }
278
279 pub fn serialize_proof_for_wire(
284 proof: sparse_merkle_tree::MerkleProof,
285 ) -> UtxoCommitmentResult<Vec<u8>> {
286 let (leaves_bitmap, merkle_path) = proof.take();
287 let mut buf = Vec::new();
289 buf.extend_from_slice(&(leaves_bitmap.len() as u32).to_le_bytes());
290 for h in &leaves_bitmap {
291 buf.extend_from_slice(h.as_slice());
292 }
293 buf.extend_from_slice(&(merkle_path.len() as u32).to_le_bytes());
294 for mv in &merkle_path {
295 match mv {
296 sparse_merkle_tree::merge::MergeValue::Value(v) => {
297 buf.push(0);
298 buf.extend_from_slice(v.as_slice());
299 }
300 sparse_merkle_tree::merge::MergeValue::MergeWithZero {
301 base_node,
302 zero_bits,
303 zero_count,
304 } => {
305 buf.push(1);
306 buf.extend_from_slice(base_node.as_slice());
307 buf.extend_from_slice(zero_bits.as_slice());
308 buf.push(*zero_count);
309 }
310 }
311 }
312 Ok(buf)
313 }
314
315 pub fn deserialize_proof_from_wire(
317 bytes: &[u8],
318 ) -> UtxoCommitmentResult<sparse_merkle_tree::MerkleProof> {
319 use sparse_merkle_tree::merge::MergeValue;
320 if bytes.len() < 8 {
321 return Err(UtxoCommitmentError::MerkleTreeError(
322 "proof too short".to_string(),
323 ));
324 }
325 let mut pos = 0;
326 let leaves_len = u32::from_le_bytes(bytes[pos..pos + 4].try_into().unwrap()) as usize;
327 pos += 4;
328 let min_after_header = match leaves_len.checked_mul(32).and_then(|b| b.checked_add(4)) {
330 Some(n) => n,
331 None => {
332 return Err(UtxoCommitmentError::MerkleTreeError(
333 "invalid proof: leaves count overflow".to_string(),
334 ));
335 }
336 };
337 let end = match pos.checked_add(min_after_header) {
338 Some(e) => e,
339 None => {
340 return Err(UtxoCommitmentError::MerkleTreeError(
341 "invalid proof: size overflow".to_string(),
342 ));
343 }
344 };
345 if end > bytes.len() {
346 return Err(UtxoCommitmentError::MerkleTreeError(
347 "proof truncated at leaves_bitmap".to_string(),
348 ));
349 }
350 let mut leaves_bitmap = Vec::with_capacity(leaves_len);
351 for _ in 0..leaves_len {
352 let mut arr = [0u8; 32];
353 arr.copy_from_slice(&bytes[pos..pos + 32]);
354 leaves_bitmap.push(H256::from(arr));
355 pos += 32;
356 }
357 let path_len = u32::from_le_bytes(bytes[pos..pos + 4].try_into().unwrap()) as usize;
358 pos += 4;
359 let max_path = (bytes.len().saturating_sub(pos)) / 33;
361 if path_len > max_path {
362 return Err(UtxoCommitmentError::MerkleTreeError(
363 "proof truncated or path count impossible for input size".to_string(),
364 ));
365 }
366 let mut merkle_path = Vec::with_capacity(path_len);
367 for _ in 0..path_len {
368 if pos >= bytes.len() {
369 return Err(UtxoCommitmentError::MerkleTreeError(
370 "proof truncated at merkle_path".to_string(),
371 ));
372 }
373 let tag = bytes[pos];
374 pos += 1;
375 match tag {
376 0 => {
377 if pos + 32 > bytes.len() {
378 return Err(UtxoCommitmentError::MerkleTreeError(
379 "proof truncated at MergeValue::Value".to_string(),
380 ));
381 }
382 let mut arr = [0u8; 32];
383 arr.copy_from_slice(&bytes[pos..pos + 32]);
384 merkle_path.push(MergeValue::Value(H256::from(arr)));
385 pos += 32;
386 }
387 1 => {
388 if pos + 65 > bytes.len() {
389 return Err(UtxoCommitmentError::MerkleTreeError(
390 "proof truncated at MergeValue::MergeWithZero".to_string(),
391 ));
392 }
393 let mut base = [0u8; 32];
394 base.copy_from_slice(&bytes[pos..pos + 32]);
395 let mut bits = [0u8; 32];
396 bits.copy_from_slice(&bytes[pos + 32..pos + 64]);
397 let zc = bytes[pos + 64];
398 merkle_path.push(MergeValue::MergeWithZero {
399 base_node: H256::from(base),
400 zero_bits: H256::from(bits),
401 zero_count: zc,
402 });
403 pos += 65;
404 }
405 _ => {
406 return Err(UtxoCommitmentError::MerkleTreeError(format!(
407 "invalid proof tag: {}",
408 tag
409 )));
410 }
411 }
412 }
413 Ok(sparse_merkle_tree::MerkleProof::new(
414 leaves_bitmap,
415 merkle_path,
416 ))
417 }
418
419 pub fn verify_commitment_supply(
424 &self,
425 commitment: &UtxoCommitment,
426 ) -> UtxoCommitmentResult<bool> {
427 use blvm_consensus::economic::total_supply;
428
429 let expected_supply = total_supply(commitment.block_height) as u64;
430 let matches = commitment.total_supply == expected_supply;
431
432 if !matches {
433 return Err(UtxoCommitmentError::VerificationFailed(format!(
434 "Supply mismatch: commitment has {}, expected {}",
435 commitment.total_supply, expected_supply
436 )));
437 }
438
439 Ok(true)
440 }
441
442 pub fn from_utxo_set(utxo_set: &crate::types::UtxoSet) -> UtxoCommitmentResult<Self> {
447 let mut tree = Self::new()?;
448 for (outpoint, utxo) in utxo_set {
449 tree.insert(outpoint.clone(), utxo.as_ref().clone())?;
450 }
451 Ok(tree)
452 }
453
454 pub fn update_from_utxo_set(
468 &mut self,
469 new_utxo_set: &crate::types::UtxoSet,
470 old_utxo_set: &crate::types::UtxoSet,
471 ) -> UtxoCommitmentResult<Hash> {
472 for (outpoint, old_utxo) in old_utxo_set {
474 if !new_utxo_set.contains_key(outpoint) {
475 self.remove(outpoint, old_utxo.as_ref())?;
476 }
477 }
478
479 for (outpoint, new_utxo) in new_utxo_set {
481 match old_utxo_set.get(outpoint) {
482 Some(old_utxo) if old_utxo == new_utxo => {}
483 _ => {
484 if let Some(old_utxo) = old_utxo_set.get(outpoint) {
485 self.remove(outpoint, old_utxo.as_ref())?;
486 }
487 self.insert(outpoint.clone(), new_utxo.as_ref().clone())?;
488 }
489 }
490 }
491
492 Ok(self.root())
493 }
494
495 pub fn to_utxo_set(&self) -> UtxoCommitmentResult<crate::types::UtxoSet> {
501 Err(UtxoCommitmentError::MerkleTreeError(
505 "UtxoMerkleTree iteration not efficiently supported. Use update_from_utxo_set() instead.".to_string()
506 ))
507 }
508
509 pub fn verify_commitment_root(&self, commitment: &UtxoCommitment) -> bool {
511 let tree_root = self.root();
512 commitment.merkle_root == tree_root
513 }
514
515 pub fn verify_utxo_proof(
532 commitment: &UtxoCommitment,
533 outpoint: &OutPoint,
534 utxo: &UTXO,
535 proof: sparse_merkle_tree::MerkleProof,
536 ) -> UtxoCommitmentResult<bool> {
537 let key = Self::hash_outpoint_static(outpoint);
539
540 let utxo_bytes = Self::serialize_utxo_static(utxo)?;
542
543 let utxo_value = UtxoValue { data: utxo_bytes };
546 let value_h256 = utxo_value.to_h256();
547
548 let root_h256 = H256::from(commitment.merkle_root);
550
551 let leaves = vec![(key, value_h256)];
553
554 let is_valid = proof
556 .verify::<UtxoHasher>(&root_h256, leaves)
557 .map_err(|e| {
558 UtxoCommitmentError::VerificationFailed(format!(
559 "Proof verification failed: {:?}",
560 e
561 ))
562 })?;
563
564 Ok(is_valid)
565 }
566
567 fn hash_outpoint(&self, outpoint: &OutPoint) -> H256 {
571 Self::hash_outpoint_static(outpoint)
572 }
573
574 fn hash_outpoint_static(outpoint: &OutPoint) -> H256 {
578 let mut hasher = Sha256::new();
579 hasher.update(&outpoint.hash);
580 hasher.update(&outpoint.index.to_be_bytes());
581 let hash = hasher.finalize();
582 let mut bytes = [0u8; 32];
583 bytes.copy_from_slice(&hash);
584 H256::from(bytes)
585 }
586
587 fn serialize_utxo(&self, utxo: &UTXO) -> UtxoCommitmentResult<Vec<u8>> {
589 Self::serialize_utxo_static(utxo)
590 }
591
592 fn serialize_utxo_static(utxo: &UTXO) -> UtxoCommitmentResult<Vec<u8>> {
598 let mut bytes = Vec::with_capacity(17 + utxo.script_pubkey.len());
599 bytes.extend_from_slice(&utxo.value.to_be_bytes());
600 bytes.extend_from_slice(&utxo.height.to_be_bytes());
601 bytes.push(if utxo.is_coinbase { 1 } else { 0 });
602 bytes.push(utxo.script_pubkey.len() as u8);
603 bytes.extend_from_slice(utxo.script_pubkey.as_ref());
604 Ok(bytes)
605 }
606
607 fn deserialize_utxo(&self, data: &[u8]) -> UtxoCommitmentResult<UTXO> {
609 if data.len() < 18 {
610 return Err(UtxoCommitmentError::InvalidUtxo(
611 "Data too short".to_string(),
612 ));
613 }
614
615 let mut offset = 0;
616 let value = i64::from_be_bytes(
617 data[offset..offset + 8]
618 .try_into()
619 .map_err(|_| UtxoCommitmentError::InvalidUtxo("Invalid value".to_string()))?,
620 );
621 offset += 8;
622
623 let height = u64::from_be_bytes(
624 data[offset..offset + 8]
625 .try_into()
626 .map_err(|_| UtxoCommitmentError::InvalidUtxo("Invalid height".to_string()))?,
627 );
628 offset += 8;
629
630 let is_coinbase = data[offset] != 0;
631 offset += 1;
632
633 let script_len = data[offset] as usize;
634 offset += 1;
635
636 if data.len() < offset + script_len {
637 return Err(UtxoCommitmentError::InvalidUtxo(
638 "Script length mismatch".to_string(),
639 ));
640 }
641
642 let script_pubkey =
643 crate::types::SharedByteString::from(&data[offset..offset + script_len]);
644
645 Ok(UTXO {
646 value,
647 script_pubkey,
648 height,
649 is_coinbase,
650 })
651 }
652}
653
654#[cfg(feature = "utxo-commitments")]
655impl Default for UtxoMerkleTree {
656 fn default() -> Self {
662 Self::new().unwrap_or_else(|e| {
663 panic!(
664 "Failed to create default UtxoMerkleTree: {:?}. This indicates a critical system error.",
665 e
666 )
667 })
668 }
669}
670
671#[cfg(not(feature = "utxo-commitments"))]
673pub struct UtxoMerkleTree;
674
675#[cfg(not(feature = "utxo-commitments"))]
676impl UtxoMerkleTree {
677 pub fn new() -> Result<Self, String> {
678 Err("UTXO commitments feature not enabled".to_string())
679 }
680}
681
682#[cfg(all(test, feature = "utxo-commitments"))]
683mod deserialize_proof_from_wire_tests {
684 use super::UtxoMerkleTree;
685
686 #[test]
688 fn wire_proof_rejects_without_space_for_path_length() {
689 const DATA: [u8; 39] = [
690 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0a, 0x00, 0x00, 0xff, 0xff, 0xff,
691 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x00, 0x00, 0x00, 0x00, 0x00,
692 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
693 ];
694 assert!(UtxoMerkleTree::deserialize_proof_from_wire(&DATA).is_err());
695 }
696}
697
698#[doc(hidden)]
714const _UTXO_MERKLE_TREE_SPEC: () = ();
715
716