1use crate::block::connect_block;
4use crate::error::Result;
5use crate::segwit::Witness;
6use crate::types::*;
7use blvm_spec_lock::spec_locked;
8use std::collections::HashMap;
9
10#[spec_locked("11.3")]
15pub fn reorganize_chain(
16 new_chain: &[Block],
17 current_chain: &[Block],
18 current_utxo_set: UtxoSet,
19 current_height: Natural,
20 network: crate::types::Network,
21) -> Result<ReorganizationResult> {
22 assert!(
24 current_height <= i64::MAX as u64,
25 "Current height {current_height} must fit in i64"
26 );
27 assert!(
28 current_utxo_set.len() <= u32::MAX as usize,
29 "Current UTXO set size {} exceeds maximum",
30 current_utxo_set.len()
31 );
32 assert!(
33 new_chain.len() <= 10_000,
34 "New chain length {} must be reasonable",
35 new_chain.len()
36 );
37 assert!(
38 current_chain.len() <= 10_000,
39 "Current chain length {} must be reasonable",
40 current_chain.len()
41 );
42
43 let empty_witnesses: Vec<Vec<Vec<Witness>>> = new_chain
46 .iter()
47 .map(|block| {
48 block
49 .transactions
50 .iter()
51 .map(|tx| tx.inputs.iter().map(|_| Vec::new()).collect())
52 .collect()
53 })
54 .collect();
55 assert!(
57 empty_witnesses.len() == new_chain.len(),
58 "Witness count {} must match new chain block count {}",
59 empty_witnesses.len(),
60 new_chain.len()
61 );
62
63 reorganize_chain_with_witnesses(
64 new_chain,
65 &empty_witnesses,
66 None, current_chain,
68 current_utxo_set,
69 current_height,
70 None::<fn(&Block) -> Option<Vec<Witness>>>, None::<fn(Natural) -> Option<Vec<BlockHeader>>>, None::<fn(&Hash) -> Option<BlockUndoLog>>, None::<fn(&Hash, &BlockUndoLog) -> Result<()>>, network,
75 )
76}
77
78#[allow(clippy::too_many_arguments)]
99#[spec_locked("11.3")]
100pub fn reorganize_chain_with_witnesses(
101 new_chain: &[Block],
102 new_chain_witnesses: &[Vec<Vec<Witness>>], new_chain_headers: Option<&[BlockHeader]>,
105 current_chain: &[Block],
106 current_utxo_set: UtxoSet,
107 current_height: Natural,
108 _get_witnesses_for_block: Option<impl Fn(&Block) -> Option<Vec<Witness>>>,
109 _get_headers_for_height: Option<impl Fn(Natural) -> Option<Vec<BlockHeader>>>,
110 get_undo_log_for_block: Option<impl Fn(&Hash) -> Option<BlockUndoLog>>,
111 store_undo_log_for_block: Option<impl Fn(&Hash, &BlockUndoLog) -> Result<()>>,
112 network: crate::types::Network,
113) -> Result<ReorganizationResult> {
114 assert!(
116 current_height <= i64::MAX as u64,
117 "Current height {current_height} must fit in i64"
118 );
119 assert!(
120 current_utxo_set.len() <= u32::MAX as usize,
121 "Current UTXO set size {} exceeds maximum",
122 current_utxo_set.len()
123 );
124 assert!(
125 new_chain.len() <= 10_000,
126 "New chain length {} must be reasonable",
127 new_chain.len()
128 );
129 assert!(
130 current_chain.len() <= 10_000,
131 "Current chain length {} must be reasonable",
132 current_chain.len()
133 );
134 assert!(
135 new_chain_witnesses.len() == new_chain.len(),
136 "New chain witness count {} must match block count {}",
137 new_chain_witnesses.len(),
138 new_chain.len()
139 );
140
141 let common_ancestor = find_common_ancestor(new_chain, current_chain)?;
143 let common_ancestor_header = common_ancestor.header;
144 let common_ancestor_index = common_ancestor.new_chain_index;
145 let current_ancestor_index = common_ancestor.current_chain_index;
146
147 assert!(
149 common_ancestor_index < new_chain.len(),
150 "Common ancestor index {} must be < new chain length {}",
151 common_ancestor_index,
152 new_chain.len()
153 );
154 assert!(
155 current_ancestor_index < current_chain.len(),
156 "Common ancestor index {} must be < current chain length {}",
157 current_ancestor_index,
158 current_chain.len()
159 );
160
161 let mut utxo_set = current_utxo_set;
167 assert!(
169 utxo_set.len() <= u32::MAX as usize,
170 "UTXO set size {} must not exceed maximum",
171 utxo_set.len()
172 );
173
174 let disconnect_start = current_ancestor_index + 1;
176 assert!(
178 disconnect_start <= current_chain.len(),
179 "Disconnect start {} must be <= current chain length {}",
180 disconnect_start,
181 current_chain.len()
182 );
183
184 let mut disconnected_undo_logs: HashMap<Hash, BlockUndoLog> = HashMap::new();
185 assert!(
187 disconnected_undo_logs.is_empty(),
188 "Disconnected undo logs must start empty"
189 );
190
191 for i in (disconnect_start..current_chain.len()).rev() {
192 assert!(i < current_chain.len(), "Block index {i} out of bounds");
194 if let Some(block) = current_chain.get(i) {
195 assert!(
197 !block.transactions.is_empty(),
198 "Block at index {i} must have at least one transaction"
199 );
200
201 let block_hash = calculate_block_hash(&block.header);
202 assert!(block_hash != [0u8; 32], "Block hash must be non-zero");
204
205 let undo_log = if let Some(ref get_undo_log) = get_undo_log_for_block {
208 get_undo_log(&block_hash).unwrap_or_else(|| {
209 BlockUndoLog::new()
213 })
214 } else {
215 BlockUndoLog::new()
218 };
219
220 utxo_set = disconnect_block(block, &undo_log, utxo_set, (i as Natural) + 1)?;
221 disconnected_undo_logs.insert(block_hash, undo_log);
222 }
223 }
224
225 let blocks_after_ancestor = (current_chain.len() - 1 - current_ancestor_index) as Natural;
232 let common_ancestor_height = current_height.saturating_sub(blocks_after_ancestor);
233 let mut new_height = common_ancestor_height;
234 let mut connected_blocks = Vec::new();
235 let mut connected_undo_logs: HashMap<Hash, BlockUndoLog> = HashMap::new();
236
237 if new_chain_witnesses.len() != new_chain.len() {
239 return Err(crate::error::ConsensusError::ConsensusRuleViolation(
240 format!(
241 "Witness count {} does not match block count {}",
242 new_chain_witnesses.len(),
243 new_chain.len()
244 )
245 .into(),
246 ));
247 }
248
249 for (i, block) in new_chain.iter().enumerate().skip(common_ancestor_index + 1) {
251 new_height += 1;
252 let witnesses = new_chain_witnesses.get(i).cloned().unwrap_or_else(|| {
255 block
256 .transactions
257 .iter()
258 .map(|tx| tx.inputs.iter().map(|_| Vec::new()).collect())
259 .collect()
260 });
261
262 let recent_headers = new_chain_headers;
267
268 let network_time = block.header.timestamp;
271 let context = crate::block::block_validation_context_for_connect_ibd(
272 recent_headers,
273 network_time,
274 network,
275 );
276 let (validation_result, new_utxo_set, undo_log) =
277 connect_block(block, &witnesses, utxo_set, new_height, &context)?;
278
279 if !matches!(validation_result, ValidationResult::Valid) {
280 return Err(crate::error::ConsensusError::ConsensusRuleViolation(
281 format!("Invalid block at height {new_height} during reorganization").into(),
282 ));
283 }
284
285 let block_hash = calculate_block_hash(&block.header);
287
288 if let Some(ref store_undo_log) = store_undo_log_for_block {
290 if let Err(e) = store_undo_log(&block_hash, &undo_log) {
291 #[cfg(any(debug_assertions, feature = "profile"))]
294 eprintln!("Warning: Failed to store undo log for block {block_hash:?}: {e}");
295 #[cfg(not(any(debug_assertions, feature = "profile")))]
296 let _ = e;
297 }
298 }
299
300 connected_undo_logs.insert(block_hash, undo_log);
302
303 utxo_set = new_utxo_set;
304 connected_blocks.push(block.clone());
305 }
306
307 Ok(ReorganizationResult {
309 new_utxo_set: utxo_set,
310 new_height,
311 common_ancestor: common_ancestor_header,
312 disconnected_blocks: current_chain[disconnect_start..].to_vec(),
313 connected_blocks,
314 reorganization_depth: current_chain.len() - disconnect_start,
315 connected_block_undo_logs: connected_undo_logs,
316 })
317}
318
319#[spec_locked("11.3")]
383pub fn update_mempool_after_reorg<F>(
384 mempool: &mut crate::mempool::Mempool,
385 reorg_result: &ReorganizationResult,
386 utxo_set: &UtxoSet,
387 get_tx_by_id: Option<F>,
388) -> Result<Vec<Hash>>
389where
390 F: Fn(&Hash) -> Option<Transaction>,
391{
392 use crate::mempool::update_mempool_after_block;
393
394 let mut all_removed = Vec::new();
395
396 for block in &reorg_result.connected_blocks {
398 let removed = update_mempool_after_block(mempool, block, utxo_set)?;
399 all_removed.extend(removed);
400 }
401
402 let mut spent_outpoints = std::collections::HashSet::new();
405 for block in &reorg_result.connected_blocks {
406 for tx in &block.transactions {
407 if !crate::transaction::is_coinbase(tx) {
408 for input in &tx.inputs {
409 spent_outpoints.insert(input.prevout);
410 }
411 }
412 }
413 }
414
415 if let Some(lookup) = get_tx_by_id {
417 let mut invalid_tx_ids = Vec::new();
418 for &tx_id in mempool.iter() {
419 if let Some(tx) = lookup(&tx_id) {
420 for input in &tx.inputs {
422 if spent_outpoints.contains(&input.prevout) {
423 invalid_tx_ids.push(tx_id);
424 break;
425 }
426 }
427 }
428 }
429
430 for tx_id in invalid_tx_ids {
432 if mempool.remove(&tx_id) {
433 all_removed.push(tx_id);
434 }
435 }
436 }
437
438 Ok(all_removed)
446}
447
448#[spec_locked("11.3")]
450pub fn update_mempool_after_reorg_simple(
451 mempool: &mut crate::mempool::Mempool,
452 reorg_result: &ReorganizationResult,
453 utxo_set: &UtxoSet,
454) -> Result<Vec<Hash>> {
455 update_mempool_after_reorg(
456 mempool,
457 reorg_result,
458 utxo_set,
459 None::<fn(&Hash) -> Option<Transaction>>,
460 )
461}
462
463struct CommonAncestorResult {
465 header: BlockHeader,
466 new_chain_index: usize,
467 current_chain_index: usize,
468}
469
470#[spec_locked("11.3")]
477fn find_common_ancestor(
478 new_chain: &[Block],
479 current_chain: &[Block],
480) -> Result<CommonAncestorResult> {
481 if new_chain.is_empty() || current_chain.is_empty() {
482 return Err(crate::error::ConsensusError::ConsensusRuleViolation(
483 "Cannot find common ancestor: empty chain".into(),
484 ));
485 }
486
487 let min_len = new_chain.len().min(current_chain.len());
489
490 for distance_from_tip in 0..min_len {
493 let new_idx = new_chain.len() - 1 - distance_from_tip;
494 let current_idx = current_chain.len() - 1 - distance_from_tip;
495
496 let new_hash = calculate_block_hash(&new_chain[new_idx].header);
497 let current_hash = calculate_block_hash(¤t_chain[current_idx].header);
498
499 if new_hash == current_hash {
501 return Ok(CommonAncestorResult {
502 header: new_chain[new_idx].header.clone(),
503 new_chain_index: new_idx,
504 current_chain_index: current_idx,
505 });
506 }
507 }
508
509 if !new_chain.is_empty() && !current_chain.is_empty() {
512 let new_genesis_hash = calculate_block_hash(&new_chain[0].header);
513 let current_genesis_hash = calculate_block_hash(¤t_chain[0].header);
514 if new_genesis_hash == current_genesis_hash {
515 return Ok(CommonAncestorResult {
516 header: new_chain[0].header.clone(),
517 new_chain_index: 0,
518 current_chain_index: 0,
519 });
520 }
521 }
522
523 Err(crate::error::ConsensusError::ConsensusRuleViolation(
525 "Chains do not share a common ancestor".into(),
526 ))
527}
528
529#[spec_locked("11.3.1")]
542fn disconnect_block(
543 _block: &Block,
544 undo_log: &BlockUndoLog,
545 mut utxo_set: UtxoSet,
546 _height: Natural,
547) -> Result<UtxoSet> {
548 assert!(
550 !_block.transactions.is_empty(),
551 "Block must have at least one transaction"
552 );
553 assert!(
554 _height <= i64::MAX as u64,
555 "Block height {_height} must fit in i64"
556 );
557 assert!(
558 utxo_set.len() <= u32::MAX as usize,
559 "UTXO set size {} must not exceed maximum",
560 utxo_set.len()
561 );
562 assert!(
564 undo_log.entries.len() <= 10_000,
565 "Undo log entry count {} must be reasonable",
566 undo_log.entries.len()
567 );
568
569 for (i, entry) in undo_log.entries.iter().enumerate() {
572 assert!(i < undo_log.entries.len(), "Entry index {i} out of bounds");
574 if entry.new_utxo.is_some() {
576 utxo_set.remove(&entry.outpoint);
577 }
578
579 if let Some(previous_utxo) = &entry.previous_utxo {
581 utxo_set.insert(entry.outpoint, std::sync::Arc::clone(previous_utxo));
582 }
583 }
584
585 Ok(utxo_set)
586}
587
588#[track_caller] #[allow(clippy::redundant_comparisons)] #[spec_locked("11.3")]
592pub fn should_reorganize(new_chain: &[Block], current_chain: &[Block]) -> Result<bool> {
593 assert!(
595 new_chain.len() <= 10_000,
596 "New chain length {} must be reasonable",
597 new_chain.len()
598 );
599 assert!(
600 current_chain.len() <= 10_000,
601 "Current chain length {} must be reasonable",
602 current_chain.len()
603 );
604
605 if new_chain.len() > current_chain.len() {
607 #[allow(clippy::eq_op)]
609 {
610 assert!(true == true || false == false, "Result must be boolean");
611 }
612 return Ok(true);
613 }
614
615 if new_chain.len() == current_chain.len() {
617 let new_work = calculate_chain_work(new_chain)?;
618 let current_work = calculate_chain_work(current_chain)?;
619 let result = new_work > current_work;
620 return Ok(result);
623 }
624
625 let result = false;
627 Ok(result)
629}
630
631#[spec_locked("11.3")]
639fn calculate_chain_work(chain: &[Block]) -> Result<u128> {
640 let mut total_work = 0u128;
641
642 for block in chain {
643 let target = expand_target(block.header.bits)?;
644 if target > 0 {
647 let work_contribution = if target == u128::MAX {
652 0 } else {
654 u128::MAX.checked_div(target + 1).unwrap_or(0)
656 };
657
658 let old_total = total_work;
661 total_work = total_work.saturating_add(work_contribution);
662
663 debug_assert!(
665 total_work >= old_total,
666 "Total work ({total_work}) must be >= previous total ({old_total})"
667 );
668 }
669 }
671
672 Ok(total_work)
675}
676
677fn expand_target(bits: Natural) -> Result<u128> {
679 let exponent = (bits >> 24) as u8;
680 let mantissa = bits & 0x00ffffff;
681
682 if exponent <= 3 {
683 let shift = 8 * (3 - exponent);
684 Ok((mantissa as u128) >> shift)
685 } else {
686 if exponent > 19 {
689 return Err(crate::error::ConsensusError::InvalidProofOfWork(
690 "Target too large".into(),
691 ));
692 }
693 let shift = 8 * (exponent - 3);
695 let mantissa_u128 = mantissa as u128;
697 let expanded = mantissa_u128.checked_shl(shift as u32).ok_or_else(|| {
698 crate::error::ConsensusError::InvalidProofOfWork("Target expansion overflow".into())
699 })?;
700 Ok(expanded)
701 }
702}
703
704#[allow(dead_code)] fn calculate_tx_id(tx: &Transaction) -> Hash {
707 let mut hash = [0u8; 32];
708 hash[0] = (tx.version & 0xff) as u8;
709 hash[1] = (tx.inputs.len() & 0xff) as u8;
710 hash[2] = (tx.outputs.len() & 0xff) as u8;
711 hash[3] = (tx.lock_time & 0xff) as u8;
712 hash
713}
714
715fn calculate_block_hash(header: &BlockHeader) -> Hash {
720 use sha2::{Digest, Sha256};
721
722 let mut bytes = Vec::with_capacity(80);
724 bytes.extend_from_slice(&header.version.to_le_bytes());
725 bytes.extend_from_slice(&header.prev_block_hash);
726 bytes.extend_from_slice(&header.merkle_root);
727 bytes.extend_from_slice(&header.timestamp.to_le_bytes());
728 bytes.extend_from_slice(&header.bits.to_le_bytes());
729 bytes.extend_from_slice(&header.nonce.to_le_bytes());
730
731 let first_hash = Sha256::digest(&bytes);
733 let second_hash = Sha256::digest(first_hash);
734
735 let mut hash = [0u8; 32];
736 hash.copy_from_slice(&second_hash);
737 hash
738}
739
740mod undo_entry_serde {
745 use crate::types::UTXO;
746 use serde::{Deserialize, Deserializer, Serialize, Serializer};
747 use std::sync::Arc;
748
749 pub fn serialize<S>(opt: &Option<Arc<UTXO>>, s: S) -> Result<S::Ok, S::Error>
750 where
751 S: Serializer,
752 {
753 opt.as_ref().map(|a| a.as_ref()).serialize(s)
754 }
755
756 pub fn deserialize<'de, D>(d: D) -> Result<Option<Arc<UTXO>>, D::Error>
757 where
758 D: Deserializer<'de>,
759 {
760 Option::<UTXO>::deserialize(d).map(|opt| opt.map(Arc::new))
761 }
762}
763
764#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
769pub struct UndoEntry {
770 pub outpoint: OutPoint,
772 #[serde(with = "undo_entry_serde")]
774 pub previous_utxo: Option<std::sync::Arc<UTXO>>,
775 #[serde(with = "undo_entry_serde")]
777 pub new_utxo: Option<std::sync::Arc<UTXO>>,
778}
779
780#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
788pub struct BlockUndoLog {
789 pub entries: Vec<UndoEntry>,
792}
793
794impl BlockUndoLog {
795 pub fn new() -> Self {
797 Self {
798 entries: Vec::new(),
799 }
800 }
801
802 pub fn push(&mut self, entry: UndoEntry) {
804 self.entries.push(entry);
805 }
806
807 pub fn is_empty(&self) -> bool {
809 self.entries.is_empty()
810 }
811}
812
813impl Default for BlockUndoLog {
814 fn default() -> Self {
815 Self::new()
816 }
817}
818
819#[derive(Debug, Clone)]
821pub struct ReorganizationResult {
822 pub new_utxo_set: UtxoSet,
823 pub new_height: Natural,
824 pub common_ancestor: BlockHeader,
825 pub disconnected_blocks: Vec<Block>,
826 pub connected_blocks: Vec<Block>,
827 pub reorganization_depth: usize,
828 pub connected_block_undo_logs: HashMap<Hash, BlockUndoLog>,
831}
832
833#[cfg(test)]
847mod property_tests {
848 use super::*;
849 use proptest::prelude::*;
850
851 fn chain_len_range() -> std::ops::Range<usize> {
853 if std::env::var("CARGO_TARPAULIN").is_ok() || std::env::var("TARPAULIN").is_ok() {
854 1..3 } else {
856 1..5
857 }
858 }
859
860 fn chain_len_range_det() -> std::ops::Range<usize> {
862 if std::env::var("CARGO_TARPAULIN").is_ok() || std::env::var("TARPAULIN").is_ok() {
863 0..3 } else {
865 0..10
866 }
867 }
868
869 fn arb_block() -> impl Strategy<Value = Block> {
871 (0x00000000u64..0x1d00ffffu64).prop_map(|bits| Block {
872 header: BlockHeader {
873 version: 1,
874 prev_block_hash: [0u8; 32],
875 merkle_root: [0u8; 32],
876 timestamp: 0,
877 bits,
878 nonce: 0,
879 },
880 transactions: Box::new([]),
881 })
882 }
883
884 proptest! {
886 #[test]
887 fn prop_should_reorganize_max_work(
888 new_chain in proptest::collection::vec(arb_block(), chain_len_range()),
889 current_chain in proptest::collection::vec(arb_block(), chain_len_range())
890 ) {
891 let new_work = calculate_chain_work(&new_chain);
893 let current_work = calculate_chain_work(¤t_chain);
894
895 if let (Ok(new_w), Ok(current_w)) = (new_work, current_work) {
897 let should_reorg = should_reorganize(&new_chain, ¤t_chain).unwrap_or(false);
899
900 if new_chain.len() > current_chain.len() {
906 prop_assert!(should_reorg, "Must reorganize when new chain is longer");
908 } else if new_chain.len() == current_chain.len() {
909 if new_w > current_w {
911 prop_assert!(should_reorg, "Must reorganize when new chain has more work (equal length)");
912 } else {
913 prop_assert!(!should_reorg, "Must not reorganize when new chain has less or equal work (equal length)");
914 }
915 } else {
916 prop_assert!(!should_reorg, "Must not reorganize when new chain is shorter");
918 }
919 }
920 }
922 }
923
924 proptest! {
926 #[test]
927 fn prop_calculate_chain_work_deterministic(
928 chain in proptest::collection::vec(arb_block(), chain_len_range_det())
929 ) {
930 let work1 = calculate_chain_work(&chain);
932 let work2 = calculate_chain_work(&chain);
933
934 match (work1, work2) {
936 (Ok(w1), Ok(w2)) => {
937 prop_assert_eq!(w1, w2, "Chain work calculation must be deterministic");
938 },
939 (Err(_), Err(_)) => {
940 },
942 _ => {
943 prop_assert!(false, "Chain work calculation must be deterministic (both succeed or both fail)");
944 }
945 }
946 }
947 }
948
949 proptest! {
951 #[test]
952 fn prop_expand_target_valid_range(
953 bits in 0x00000000u64..0x1d00ffffu64
954 ) {
955 let result = expand_target(bits);
956
957 match result {
958 Ok(target) => {
959 let _ = target;
963 },
964 Err(_) => {
965 }
967 }
968 }
969 }
970
971 proptest! {
973 #[test]
974 fn prop_should_reorganize_equal_length(
975 chain1 in proptest::collection::vec(arb_block(), 1..3),
976 chain2 in proptest::collection::vec(arb_block(), 1..3)
977 ) {
978 let len = chain1.len().min(chain2.len());
980 let chain1 = &chain1[..len];
981 let chain2 = &chain2[..len];
982
983 let work1 = calculate_chain_work(chain1);
984 let work2 = calculate_chain_work(chain2);
985
986 if let (Ok(w1), Ok(w2)) = (work1, work2) {
988 let should_reorg = should_reorganize(chain1, chain2).unwrap_or(false);
989
990 if w1 > w2 {
992 prop_assert!(should_reorg, "Must reorganize when first chain has more work");
993 } else {
994 prop_assert!(!should_reorg, "Must not reorganize when first chain has less or equal work");
995 }
996 }
997 }
999 }
1000}
1001
1002#[cfg(test)]
1003mod tests {
1004 use super::*;
1005
1006 #[test]
1007 fn test_should_reorganize_longer_chain() {
1008 let new_chain = vec![create_test_block(), create_test_block()];
1009 let current_chain = vec![create_test_block()];
1010
1011 assert!(should_reorganize(&new_chain, ¤t_chain).unwrap());
1012 }
1013
1014 #[test]
1015 fn test_should_reorganize_same_length_more_work() {
1016 let mut new_chain = vec![create_test_block()];
1017 let mut current_chain = vec![create_test_block()];
1018
1019 new_chain[0].header.bits = 0x0200ffff; current_chain[0].header.bits = 0x0300ffff; assert!(should_reorganize(&new_chain, ¤t_chain).unwrap());
1024 }
1025
1026 #[test]
1027 fn test_should_not_reorganize_shorter_chain() {
1028 let new_chain = vec![create_test_block()];
1029 let current_chain = vec![create_test_block(), create_test_block()];
1030
1031 assert!(!should_reorganize(&new_chain, ¤t_chain).unwrap());
1032 }
1033
1034 #[test]
1035 fn test_find_common_ancestor() {
1036 let new_chain = vec![create_test_block()];
1037 let current_chain = vec![create_test_block()];
1038
1039 let ancestor = find_common_ancestor(&new_chain, ¤t_chain).unwrap();
1040 assert_eq!(ancestor.header.version, 4);
1041 assert_eq!(ancestor.new_chain_index, 0);
1042 assert_eq!(ancestor.current_chain_index, 0);
1043 }
1044
1045 #[test]
1046 fn test_find_common_ancestor_empty_chain() {
1047 let new_chain = vec![];
1048 let current_chain = vec![create_test_block()];
1049
1050 let result = find_common_ancestor(&new_chain, ¤t_chain);
1051 assert!(result.is_err());
1052 }
1053
1054 #[test]
1055 fn test_calculate_chain_work() {
1056 let mut block = create_test_block();
1057 block.header.bits = 0x0300ffff;
1059 let chain = vec![block];
1060 let work = calculate_chain_work(&chain).unwrap();
1061 assert!(work > 0);
1062 }
1063
1064 #[test]
1065 fn test_reorganize_chain() {
1066 let ancestor = create_test_block_at_height(0);
1070 let mut new_block = create_test_block_at_height(1);
1071 new_block.header.nonce = 42; let new_chain = vec![ancestor.clone(), new_block];
1075 let current_chain = vec![ancestor];
1076 let utxo_set = UtxoSet::default();
1077
1078 let result = reorganize_chain(
1080 &new_chain,
1081 ¤t_chain,
1082 utxo_set,
1083 1,
1084 crate::types::Network::Regtest,
1085 );
1086 match result {
1087 Ok(reorg_result) => {
1088 assert_eq!(reorg_result.new_height, 2);
1091 assert_eq!(reorg_result.connected_blocks.len(), 1);
1092 assert_eq!(reorg_result.connected_block_undo_logs.len(), 1);
1093 }
1094 Err(_) => {
1095 }
1097 }
1098 }
1099
1100 #[test]
1101 fn test_reorganize_chain_deep_reorg() {
1102 let mut block1 = create_test_block();
1104 block1.header.nonce = 1;
1105 let mut block2 = create_test_block();
1106 block2.header.nonce = 2;
1107 let mut block3 = create_test_block();
1108 block3.header.nonce = 3;
1109 let new_chain = vec![block1, block2, block3];
1110
1111 let mut current_block1 = create_test_block();
1112 current_block1.header.nonce = 10;
1113 let mut current_block2 = create_test_block();
1114 current_block2.header.nonce = 11;
1115 let current_chain = vec![current_block1, current_block2];
1116 let utxo_set = UtxoSet::default();
1117
1118 let result = reorganize_chain(
1119 &new_chain,
1120 ¤t_chain,
1121 utxo_set,
1122 2,
1123 crate::types::Network::Regtest,
1124 );
1125 match result {
1126 Ok(reorg_result) => {
1127 assert_eq!(reorg_result.connected_blocks.len(), 3);
1128 assert_eq!(reorg_result.reorganization_depth, 2);
1129 assert_eq!(reorg_result.connected_block_undo_logs.len(), 3);
1131 }
1132 Err(_) => {
1133 }
1135 }
1136 }
1137
1138 #[test]
1139 fn test_undo_log_storage_and_retrieval() {
1140 use crate::block::connect_block;
1141 use crate::segwit::Witness;
1142
1143 let block = create_test_block_at_height(1);
1144 let mut utxo_set = UtxoSet::default();
1145
1146 let tx_id = calculate_tx_id(&block.transactions[0]);
1148 let outpoint = OutPoint {
1149 hash: tx_id,
1150 index: 0,
1151 };
1152 let utxo = UTXO {
1153 value: 5_000_000_000,
1154 script_pubkey: vec![0x51].into(),
1155 height: 1,
1156 is_coinbase: false,
1157 };
1158 utxo_set.insert(outpoint.clone(), std::sync::Arc::new(utxo.clone()));
1159
1160 let witnesses: Vec<Vec<Witness>> = block
1162 .transactions
1163 .iter()
1164 .map(|tx| tx.inputs.iter().map(|_| Vec::new()).collect())
1165 .collect();
1166 let ctx = crate::block::BlockValidationContext::for_network(crate::types::Network::Regtest);
1167 let (result, new_utxo_set, undo_log) =
1168 connect_block(&block, &witnesses, utxo_set.clone(), 1, &ctx).unwrap();
1169
1170 assert!(matches!(result, crate::types::ValidationResult::Valid));
1171
1172 assert!(
1174 !undo_log.entries.is_empty(),
1175 "Undo log should contain entries"
1176 );
1177
1178 let block_hash = calculate_block_hash(&block.header);
1180
1181 let mut undo_log_storage: HashMap<Hash, BlockUndoLog> = HashMap::new();
1183 undo_log_storage.insert(block_hash, undo_log.clone());
1184
1185 let retrieved_undo_log = undo_log_storage.get(&block_hash);
1187 assert!(
1188 retrieved_undo_log.is_some(),
1189 "Should be able to retrieve undo log"
1190 );
1191 assert_eq!(
1192 retrieved_undo_log.unwrap().entries.len(),
1193 undo_log.entries.len()
1194 );
1195
1196 let disconnected_utxo_set = disconnect_block(&block, &undo_log, new_utxo_set, 1).unwrap();
1198
1199 assert!(
1201 disconnected_utxo_set.contains_key(&outpoint),
1202 "Disconnected UTXO set should contain restored UTXO"
1203 );
1204 }
1205
1206 #[test]
1207 fn test_reorganize_with_undo_log_callback() {
1208 use crate::block::connect_block;
1209 use crate::segwit::Witness;
1210
1211 let block = create_test_block_at_height(1);
1213 let utxo_set = UtxoSet::default();
1214 let witnesses: Vec<Vec<Witness>> = block
1215 .transactions
1216 .iter()
1217 .map(|tx| tx.inputs.iter().map(|_| Vec::new()).collect())
1218 .collect();
1219
1220 let ctx = crate::block::BlockValidationContext::for_network(crate::types::Network::Regtest);
1221 let (result, connected_utxo_set, undo_log) =
1222 connect_block(&block, &witnesses, utxo_set.clone(), 1, &ctx).unwrap();
1223
1224 if !matches!(result, crate::types::ValidationResult::Valid) {
1225 eprintln!("Block validation failed: {:?}", result);
1226 }
1227 assert!(matches!(result, crate::types::ValidationResult::Valid));
1228
1229 let block_hash = calculate_block_hash(&block.header);
1231 let mut undo_log_storage: HashMap<Hash, BlockUndoLog> = HashMap::new();
1232 undo_log_storage.insert(block_hash, undo_log);
1233
1234 let get_undo_log =
1236 |hash: &Hash| -> Option<BlockUndoLog> { undo_log_storage.get(hash).cloned() };
1237
1238 let mut new_block = create_test_block_at_height(2);
1241 new_block.header.nonce = 42; let new_chain = vec![block.clone(), new_block];
1243 let current_chain = vec![block];
1244 let empty_witnesses: Vec<Vec<Vec<Witness>>> = new_chain
1245 .iter()
1246 .map(|b| {
1247 b.transactions
1248 .iter()
1249 .map(|tx| tx.inputs.iter().map(|_| Vec::new()).collect())
1250 .collect()
1251 })
1252 .collect();
1253
1254 let reorg_result = reorganize_chain_with_witnesses(
1255 &new_chain,
1256 &empty_witnesses,
1257 None,
1258 ¤t_chain,
1259 connected_utxo_set,
1260 1,
1261 None::<fn(&Block) -> Option<Vec<Witness>>>,
1262 None::<fn(Natural) -> Option<Vec<BlockHeader>>>,
1263 Some(get_undo_log),
1264 None::<fn(&Hash, &BlockUndoLog) -> Result<()>>, crate::types::Network::Regtest,
1266 );
1267
1268 match reorg_result {
1270 Ok(result) => {
1271 assert!(!result.connected_block_undo_logs.is_empty());
1273 }
1274 Err(_) => {
1275 }
1277 }
1278 }
1279
1280 #[test]
1281 fn test_reorganize_chain_empty_new_chain() {
1282 let new_chain = vec![];
1283 let current_chain = vec![create_test_block()];
1284 let utxo_set = UtxoSet::default();
1285
1286 let result = reorganize_chain(
1287 &new_chain,
1288 ¤t_chain,
1289 utxo_set,
1290 1,
1291 crate::types::Network::Regtest,
1292 );
1293 assert!(result.is_err());
1294 }
1295
1296 #[test]
1297 fn test_reorganize_chain_empty_current_chain() {
1298 let new_chain = vec![create_test_block()];
1299 let current_chain = vec![];
1300 let utxo_set = UtxoSet::default();
1301
1302 let result = reorganize_chain(
1303 &new_chain,
1304 ¤t_chain,
1305 utxo_set,
1306 0,
1307 crate::types::Network::Regtest,
1308 );
1309 assert!(result.is_err());
1310 }
1311
1312 #[test]
1313 fn test_disconnect_block() {
1314 let block = create_test_block();
1315 let mut utxo_set = UtxoSet::default();
1316
1317 let tx_id = calculate_tx_id(&block.transactions[0]);
1319 let outpoint = OutPoint {
1320 hash: tx_id,
1321 index: 0,
1322 };
1323 let utxo = UTXO {
1324 value: 50_000_000_000,
1325 script_pubkey: vec![0x51].into(),
1326 height: 1,
1327 is_coinbase: false,
1328 };
1329 utxo_set.insert(outpoint, std::sync::Arc::new(utxo));
1330
1331 let empty_undo_log = BlockUndoLog::new();
1333 let result = disconnect_block(&block, &empty_undo_log, utxo_set, 1);
1334 assert!(result.is_ok());
1335 }
1336
1337 #[test]
1338 fn test_calculate_chain_work_empty_chain() {
1339 let chain = vec![];
1340 let work = calculate_chain_work(&chain).unwrap();
1341 assert_eq!(work, 0);
1342 }
1343
1344 #[test]
1345 fn test_calculate_chain_work_multiple_blocks() {
1346 let mut chain = vec![create_test_block(), create_test_block()];
1347 chain[0].header.bits = 0x0300ffff;
1349 chain[1].header.bits = 0x0200ffff;
1350
1351 let work = calculate_chain_work(&chain).unwrap();
1352 assert!(work > 0);
1353 }
1354
1355 #[test]
1356 fn test_expand_target_edge_cases() {
1357 let result = expand_target(0x00000000);
1359 assert!(result.is_ok());
1360
1361 let result = expand_target(0x03ffffff);
1363 assert!(result.is_ok());
1364
1365 let result = expand_target(0x14000000); assert!(result.is_err());
1368 }
1369
1370 #[test]
1371 fn test_calculate_tx_id_different_transactions() {
1372 let tx1 = Transaction {
1373 version: 1,
1374 inputs: vec![].into(),
1375 outputs: vec![].into(),
1376 lock_time: 0,
1377 };
1378
1379 let tx2 = Transaction {
1380 version: 2,
1381 inputs: vec![].into(),
1382 outputs: vec![].into(),
1383 lock_time: 0,
1384 };
1385
1386 let id1 = calculate_tx_id(&tx1);
1387 let id2 = calculate_tx_id(&tx2);
1388
1389 assert_ne!(id1, id2);
1390 }
1391
1392 fn encode_bip34_height(height: u64) -> Vec<u8> {
1397 if height == 0 {
1398 return vec![0x00, 0xff]; }
1401 let mut height_bytes = Vec::new();
1402 let mut n = height;
1403 while n > 0 {
1404 height_bytes.push((n & 0xff) as u8);
1405 n >>= 8;
1406 }
1407 if height_bytes.last().map_or(false, |&b| b & 0x80 != 0) {
1409 height_bytes.push(0x00);
1410 }
1411 let mut script_sig = Vec::with_capacity(1 + height_bytes.len() + 1);
1412 script_sig.push(height_bytes.len() as u8); script_sig.extend_from_slice(&height_bytes);
1414 if script_sig.len() < 2 {
1416 script_sig.push(0xff);
1417 }
1418 script_sig
1419 }
1420
1421 fn create_test_block_at_height(height: u64) -> Block {
1423 use crate::mining::calculate_merkle_root;
1424
1425 let script_sig = encode_bip34_height(height);
1426 let coinbase_tx = Transaction {
1427 version: 1,
1428 inputs: vec![TransactionInput {
1429 prevout: OutPoint {
1430 hash: [0; 32].into(),
1431 index: 0xffffffff,
1432 },
1433 script_sig,
1434 sequence: 0xffffffff,
1435 }]
1436 .into(),
1437 outputs: vec![TransactionOutput {
1438 value: 5_000_000_000,
1439 script_pubkey: vec![0x51].into(),
1440 }]
1441 .into(),
1442 lock_time: 0,
1443 };
1444
1445 let merkle_root =
1446 calculate_merkle_root(&[coinbase_tx.clone()]).expect("Failed to calculate merkle root");
1447
1448 Block {
1449 header: BlockHeader {
1450 version: 4,
1451 prev_block_hash: [0; 32],
1452 merkle_root,
1453 timestamp: 1231006505,
1454 bits: 0x207fffff, nonce: 0,
1456 },
1457 transactions: vec![coinbase_tx].into_boxed_slice(),
1458 }
1459 }
1460
1461 fn create_test_block() -> Block {
1463 create_test_block_at_height(0)
1464 }
1465}