blvm_protocol/utxo_commitments/
merkle_tree.rs1#[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::{H256, SparseMerkleTree};
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 let mut hash = [0u8; 32];
192 hash.copy_from_slice(root_h256.as_slice());
193 Ok(hash)
194 }
195
196 pub fn get(&self, outpoint: &OutPoint) -> UtxoCommitmentResult<Option<UTXO>> {
198 let key = self.hash_outpoint(outpoint);
199
200 match self.tree.get(&key) {
201 Ok(value) => {
202 if value.to_h256() == H256::zero() || value.to_h256() == UtxoValue::zero().to_h256()
204 {
205 Ok(None)
206 } else {
207 let serialized_data = &value.data;
209
210 match self.deserialize_utxo(serialized_data) {
212 Ok(utxo) => Ok(Some(utxo)),
213 Err(e) => {
214 Err(UtxoCommitmentError::InvalidUtxo(format!(
216 "Failed to deserialize UTXO: {e}"
217 )))
218 }
219 }
220 }
221 }
222 Err(_) => Ok(None),
223 }
224 }
225
226 #[spec_locked("11.4", "GenerateCommitment")]
228 #[blvm_spec_lock::ensures(result.block_height == block_height)]
229 #[blvm_spec_lock::ensures(result.block_hash == block_hash)]
230 pub fn generate_commitment(&self, block_hash: Hash, block_height: Natural) -> UtxoCommitment {
231 let merkle_root = self.root();
232 UtxoCommitment::new(
233 merkle_root,
234 self.total_supply,
235 self.utxo_count,
236 block_height,
237 block_hash,
238 )
239 }
240
241 pub fn total_supply(&self) -> u64 {
243 self.total_supply
244 }
245
246 pub fn utxo_count(&self) -> u64 {
248 self.utxo_count
249 }
250
251 pub fn generate_proof(
255 &self,
256 outpoint: &OutPoint,
257 ) -> UtxoCommitmentResult<sparse_merkle_tree::MerkleProof> {
258 let key = self.hash_outpoint(outpoint);
259 let keys = vec![key];
260
261 self.tree.merkle_proof(keys).map_err(|e| {
262 UtxoCommitmentError::MerkleTreeError(format!("Failed to generate proof: {e:?}"))
263 })
264 }
265
266 pub fn serialize_proof_for_wire(
271 proof: sparse_merkle_tree::MerkleProof,
272 ) -> UtxoCommitmentResult<Vec<u8>> {
273 let (leaves_bitmap, merkle_path) = proof.take();
274 let mut buf = Vec::new();
276 buf.extend_from_slice(&(leaves_bitmap.len() as u32).to_le_bytes());
277 for h in &leaves_bitmap {
278 buf.extend_from_slice(h.as_slice());
279 }
280 buf.extend_from_slice(&(merkle_path.len() as u32).to_le_bytes());
281 for mv in &merkle_path {
282 match mv {
283 sparse_merkle_tree::merge::MergeValue::Value(v) => {
284 buf.push(0);
285 buf.extend_from_slice(v.as_slice());
286 }
287 sparse_merkle_tree::merge::MergeValue::MergeWithZero {
288 base_node,
289 zero_bits,
290 zero_count,
291 } => {
292 buf.push(1);
293 buf.extend_from_slice(base_node.as_slice());
294 buf.extend_from_slice(zero_bits.as_slice());
295 buf.push(*zero_count);
296 }
297 }
298 }
299 Ok(buf)
300 }
301
302 pub fn deserialize_proof_from_wire(
304 bytes: &[u8],
305 ) -> UtxoCommitmentResult<sparse_merkle_tree::MerkleProof> {
306 use sparse_merkle_tree::merge::MergeValue;
307 if bytes.len() < 8 {
308 return Err(UtxoCommitmentError::MerkleTreeError(
309 "proof too short".to_string(),
310 ));
311 }
312 let mut pos = 0;
313 let leaves_len = u32::from_le_bytes(bytes[pos..pos + 4].try_into().unwrap()) as usize;
314 pos += 4;
315 let min_after_header = match leaves_len.checked_mul(32).and_then(|b| b.checked_add(4)) {
317 Some(n) => n,
318 None => {
319 return Err(UtxoCommitmentError::MerkleTreeError(
320 "invalid proof: leaves count overflow".to_string(),
321 ));
322 }
323 };
324 let end = match pos.checked_add(min_after_header) {
325 Some(e) => e,
326 None => {
327 return Err(UtxoCommitmentError::MerkleTreeError(
328 "invalid proof: size overflow".to_string(),
329 ));
330 }
331 };
332 if end > bytes.len() {
333 return Err(UtxoCommitmentError::MerkleTreeError(
334 "proof truncated at leaves_bitmap".to_string(),
335 ));
336 }
337 let mut leaves_bitmap = Vec::with_capacity(leaves_len);
338 for _ in 0..leaves_len {
339 let mut arr = [0u8; 32];
340 arr.copy_from_slice(&bytes[pos..pos + 32]);
341 leaves_bitmap.push(H256::from(arr));
342 pos += 32;
343 }
344 let path_len = u32::from_le_bytes(bytes[pos..pos + 4].try_into().unwrap()) as usize;
345 pos += 4;
346 let max_path = (bytes.len().saturating_sub(pos)) / 33;
348 if path_len > max_path {
349 return Err(UtxoCommitmentError::MerkleTreeError(
350 "proof truncated or path count impossible for input size".to_string(),
351 ));
352 }
353 let mut merkle_path = Vec::with_capacity(path_len);
354 for _ in 0..path_len {
355 if pos >= bytes.len() {
356 return Err(UtxoCommitmentError::MerkleTreeError(
357 "proof truncated at merkle_path".to_string(),
358 ));
359 }
360 let tag = bytes[pos];
361 pos += 1;
362 match tag {
363 0 => {
364 if pos + 32 > bytes.len() {
365 return Err(UtxoCommitmentError::MerkleTreeError(
366 "proof truncated at MergeValue::Value".to_string(),
367 ));
368 }
369 let mut arr = [0u8; 32];
370 arr.copy_from_slice(&bytes[pos..pos + 32]);
371 merkle_path.push(MergeValue::Value(H256::from(arr)));
372 pos += 32;
373 }
374 1 => {
375 if pos + 65 > bytes.len() {
376 return Err(UtxoCommitmentError::MerkleTreeError(
377 "proof truncated at MergeValue::MergeWithZero".to_string(),
378 ));
379 }
380 let mut base = [0u8; 32];
381 base.copy_from_slice(&bytes[pos..pos + 32]);
382 let mut bits = [0u8; 32];
383 bits.copy_from_slice(&bytes[pos + 32..pos + 64]);
384 let zc = bytes[pos + 64];
385 merkle_path.push(MergeValue::MergeWithZero {
386 base_node: H256::from(base),
387 zero_bits: H256::from(bits),
388 zero_count: zc,
389 });
390 pos += 65;
391 }
392 _ => {
393 return Err(UtxoCommitmentError::MerkleTreeError(format!(
394 "invalid proof tag: {tag}"
395 )));
396 }
397 }
398 }
399 Ok(sparse_merkle_tree::MerkleProof::new(
400 leaves_bitmap,
401 merkle_path,
402 ))
403 }
404
405 pub fn verify_commitment_supply(
410 &self,
411 commitment: &UtxoCommitment,
412 ) -> UtxoCommitmentResult<bool> {
413 use blvm_consensus::economic::total_supply;
414
415 let expected_supply = total_supply(commitment.block_height) as u64;
416 let matches = commitment.total_supply == expected_supply;
417
418 if !matches {
419 return Err(UtxoCommitmentError::VerificationFailed(format!(
420 "Supply mismatch: commitment has {}, expected {}",
421 commitment.total_supply, expected_supply
422 )));
423 }
424
425 Ok(true)
426 }
427
428 pub fn from_utxo_set(utxo_set: &crate::types::UtxoSet) -> UtxoCommitmentResult<Self> {
433 let mut tree = Self::new()?;
434 for (outpoint, utxo) in utxo_set {
435 tree.insert(*outpoint, utxo.as_ref().clone())?;
436 }
437 Ok(tree)
438 }
439
440 pub fn update_from_utxo_set(
454 &mut self,
455 new_utxo_set: &crate::types::UtxoSet,
456 old_utxo_set: &crate::types::UtxoSet,
457 ) -> UtxoCommitmentResult<Hash> {
458 for (outpoint, old_utxo) in old_utxo_set {
460 if !new_utxo_set.contains_key(outpoint) {
461 self.remove(outpoint, old_utxo.as_ref())?;
462 }
463 }
464
465 for (outpoint, new_utxo) in new_utxo_set {
467 match old_utxo_set.get(outpoint) {
468 Some(old_utxo) if old_utxo == new_utxo => {}
469 _ => {
470 if let Some(old_utxo) = old_utxo_set.get(outpoint) {
471 self.remove(outpoint, old_utxo.as_ref())?;
472 }
473 self.insert(*outpoint, new_utxo.as_ref().clone())?;
474 }
475 }
476 }
477
478 Ok(self.root())
479 }
480
481 pub fn to_utxo_set(&self) -> UtxoCommitmentResult<crate::types::UtxoSet> {
487 Err(UtxoCommitmentError::MerkleTreeError(
491 "UtxoMerkleTree iteration not efficiently supported. Use update_from_utxo_set() instead.".to_string()
492 ))
493 }
494
495 pub fn verify_commitment_root(&self, commitment: &UtxoCommitment) -> bool {
497 let tree_root = self.root();
498 commitment.merkle_root == tree_root
499 }
500
501 pub fn verify_utxo_proof(
518 commitment: &UtxoCommitment,
519 outpoint: &OutPoint,
520 utxo: &UTXO,
521 proof: sparse_merkle_tree::MerkleProof,
522 ) -> UtxoCommitmentResult<bool> {
523 let key = Self::hash_outpoint_static(outpoint);
525
526 let utxo_bytes = Self::serialize_utxo_static(utxo)?;
528
529 let utxo_value = UtxoValue { data: utxo_bytes };
532 let value_h256 = utxo_value.to_h256();
533
534 let root_h256 = H256::from(commitment.merkle_root);
536
537 let leaves = vec![(key, value_h256)];
539
540 let is_valid = proof
542 .verify::<UtxoHasher>(&root_h256, leaves)
543 .map_err(|e| {
544 UtxoCommitmentError::VerificationFailed(format!("Proof verification failed: {e:?}"))
545 })?;
546
547 Ok(is_valid)
548 }
549
550 fn hash_outpoint(&self, outpoint: &OutPoint) -> H256 {
554 Self::hash_outpoint_static(outpoint)
555 }
556
557 fn hash_outpoint_static(outpoint: &OutPoint) -> H256 {
561 let mut hasher = Sha256::new();
562 hasher.update(outpoint.hash);
563 hasher.update(outpoint.index.to_be_bytes());
564 let hash = hasher.finalize();
565 let mut bytes = [0u8; 32];
566 bytes.copy_from_slice(&hash);
567 H256::from(bytes)
568 }
569
570 fn serialize_utxo(&self, utxo: &UTXO) -> UtxoCommitmentResult<Vec<u8>> {
572 Self::serialize_utxo_static(utxo)
573 }
574
575 fn serialize_utxo_static(utxo: &UTXO) -> UtxoCommitmentResult<Vec<u8>> {
581 let mut bytes = Vec::with_capacity(17 + utxo.script_pubkey.len());
582 bytes.extend_from_slice(&utxo.value.to_be_bytes());
583 bytes.extend_from_slice(&utxo.height.to_be_bytes());
584 bytes.push(if utxo.is_coinbase { 1 } else { 0 });
585 bytes.push(utxo.script_pubkey.len() as u8);
586 bytes.extend_from_slice(utxo.script_pubkey.as_ref());
587 Ok(bytes)
588 }
589
590 fn deserialize_utxo(&self, data: &[u8]) -> UtxoCommitmentResult<UTXO> {
592 if data.len() < 18 {
593 return Err(UtxoCommitmentError::InvalidUtxo(
594 "Data too short".to_string(),
595 ));
596 }
597
598 let mut offset = 0;
599 let value = i64::from_be_bytes(
600 data[offset..offset + 8]
601 .try_into()
602 .map_err(|_| UtxoCommitmentError::InvalidUtxo("Invalid value".to_string()))?,
603 );
604 offset += 8;
605
606 let height = u64::from_be_bytes(
607 data[offset..offset + 8]
608 .try_into()
609 .map_err(|_| UtxoCommitmentError::InvalidUtxo("Invalid height".to_string()))?,
610 );
611 offset += 8;
612
613 let is_coinbase = data[offset] != 0;
614 offset += 1;
615
616 let script_len = data[offset] as usize;
617 offset += 1;
618
619 if data.len() < offset + script_len {
620 return Err(UtxoCommitmentError::InvalidUtxo(
621 "Script length mismatch".to_string(),
622 ));
623 }
624
625 let script_pubkey =
626 crate::types::SharedByteString::from(&data[offset..offset + script_len]);
627
628 Ok(UTXO {
629 value,
630 script_pubkey,
631 height,
632 is_coinbase,
633 })
634 }
635}
636
637#[cfg(feature = "utxo-commitments")]
638impl Default for UtxoMerkleTree {
639 fn default() -> Self {
645 Self::new().unwrap_or_else(|e| {
646 panic!(
647 "Failed to create default UtxoMerkleTree: {e:?}. This indicates a critical system error."
648 )
649 })
650 }
651}
652
653#[cfg(not(feature = "utxo-commitments"))]
655pub struct UtxoMerkleTree;
656
657#[cfg(not(feature = "utxo-commitments"))]
658impl UtxoMerkleTree {
659 pub fn new() -> Result<Self, String> {
660 Err("UTXO commitments feature not enabled".to_string())
661 }
662}
663
664#[cfg(all(test, feature = "utxo-commitments"))]
665mod deserialize_proof_from_wire_tests {
666 use super::UtxoMerkleTree;
667
668 #[test]
670 fn wire_proof_rejects_without_space_for_path_length() {
671 const DATA: [u8; 39] = [
672 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0a, 0x00, 0x00, 0xff, 0xff, 0xff,
673 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x00, 0x00, 0x00, 0x00, 0x00,
674 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
675 ];
676 assert!(UtxoMerkleTree::deserialize_proof_from_wire(&DATA).is_err());
677 }
678}
679
680#[doc(hidden)]
696const _UTXO_MERKLE_TREE_SPEC: () = ();
697
698