1use crate::block::{connect_block, BlockValidationContext};
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 = BlockValidationContext::from_connect_block_ibd_args(
272 recent_headers,
273 network_time,
274 network,
275 None,
276 None,
277 );
278 let (validation_result, new_utxo_set, undo_log) =
279 connect_block(block, &witnesses, utxo_set, new_height, &context)?;
280
281 if !matches!(validation_result, ValidationResult::Valid) {
282 return Err(crate::error::ConsensusError::ConsensusRuleViolation(
283 format!("Invalid block at height {new_height} during reorganization").into(),
284 ));
285 }
286
287 let block_hash = calculate_block_hash(&block.header);
289
290 if let Some(ref store_undo_log) = store_undo_log_for_block {
292 if let Err(e) = store_undo_log(&block_hash, &undo_log) {
293 #[cfg(any(debug_assertions, feature = "profile"))]
296 eprintln!("Warning: Failed to store undo log for block {block_hash:?}: {e}");
297 #[cfg(not(any(debug_assertions, feature = "profile")))]
298 let _ = e;
299 }
300 }
301
302 connected_undo_logs.insert(block_hash, undo_log);
304
305 utxo_set = new_utxo_set;
306 connected_blocks.push(block.clone());
307 }
308
309 Ok(ReorganizationResult {
311 new_utxo_set: utxo_set,
312 new_height,
313 common_ancestor: common_ancestor_header,
314 disconnected_blocks: current_chain[disconnect_start..].to_vec(),
315 connected_blocks,
316 reorganization_depth: current_chain.len() - disconnect_start,
317 connected_block_undo_logs: connected_undo_logs,
318 })
319}
320
321#[spec_locked("11.3")]
385pub fn update_mempool_after_reorg<F>(
386 mempool: &mut crate::mempool::Mempool,
387 reorg_result: &ReorganizationResult,
388 utxo_set: &UtxoSet,
389 get_tx_by_id: Option<F>,
390) -> Result<Vec<Hash>>
391where
392 F: Fn(&Hash) -> Option<Transaction>,
393{
394 use crate::mempool::update_mempool_after_block;
395
396 let mut all_removed = Vec::new();
397
398 for block in &reorg_result.connected_blocks {
400 let removed = update_mempool_after_block(mempool, block, utxo_set)?;
401 all_removed.extend(removed);
402 }
403
404 let mut spent_outpoints = std::collections::HashSet::new();
407 for block in &reorg_result.connected_blocks {
408 for tx in &block.transactions {
409 if !crate::transaction::is_coinbase(tx) {
410 for input in &tx.inputs {
411 spent_outpoints.insert(input.prevout);
412 }
413 }
414 }
415 }
416
417 if let Some(lookup) = get_tx_by_id {
419 let mut invalid_tx_ids = Vec::new();
420 for &tx_id in mempool.iter() {
421 if let Some(tx) = lookup(&tx_id) {
422 for input in &tx.inputs {
424 if spent_outpoints.contains(&input.prevout) {
425 invalid_tx_ids.push(tx_id);
426 break;
427 }
428 }
429 }
430 }
431
432 for tx_id in invalid_tx_ids {
434 if mempool.remove(&tx_id) {
435 all_removed.push(tx_id);
436 }
437 }
438 }
439
440 Ok(all_removed)
448}
449
450#[spec_locked("11.3")]
452pub fn update_mempool_after_reorg_simple(
453 mempool: &mut crate::mempool::Mempool,
454 reorg_result: &ReorganizationResult,
455 utxo_set: &UtxoSet,
456) -> Result<Vec<Hash>> {
457 update_mempool_after_reorg(
458 mempool,
459 reorg_result,
460 utxo_set,
461 None::<fn(&Hash) -> Option<Transaction>>,
462 )
463}
464
465struct CommonAncestorResult {
467 header: BlockHeader,
468 new_chain_index: usize,
469 current_chain_index: usize,
470}
471
472#[spec_locked("11.3")]
479fn find_common_ancestor(
480 new_chain: &[Block],
481 current_chain: &[Block],
482) -> Result<CommonAncestorResult> {
483 if new_chain.is_empty() || current_chain.is_empty() {
484 return Err(crate::error::ConsensusError::ConsensusRuleViolation(
485 "Cannot find common ancestor: empty chain".into(),
486 ));
487 }
488
489 let min_len = new_chain.len().min(current_chain.len());
491
492 for distance_from_tip in 0..min_len {
495 let new_idx = new_chain.len() - 1 - distance_from_tip;
496 let current_idx = current_chain.len() - 1 - distance_from_tip;
497
498 let new_hash = calculate_block_hash(&new_chain[new_idx].header);
499 let current_hash = calculate_block_hash(¤t_chain[current_idx].header);
500
501 if new_hash == current_hash {
503 return Ok(CommonAncestorResult {
504 header: new_chain[new_idx].header.clone(),
505 new_chain_index: new_idx,
506 current_chain_index: current_idx,
507 });
508 }
509 }
510
511 if !new_chain.is_empty() && !current_chain.is_empty() {
514 let new_genesis_hash = calculate_block_hash(&new_chain[0].header);
515 let current_genesis_hash = calculate_block_hash(¤t_chain[0].header);
516 if new_genesis_hash == current_genesis_hash {
517 return Ok(CommonAncestorResult {
518 header: new_chain[0].header.clone(),
519 new_chain_index: 0,
520 current_chain_index: 0,
521 });
522 }
523 }
524
525 Err(crate::error::ConsensusError::ConsensusRuleViolation(
527 "Chains do not share a common ancestor".into(),
528 ))
529}
530
531#[spec_locked("11.3.1")]
544fn disconnect_block(
545 _block: &Block,
546 undo_log: &BlockUndoLog,
547 mut utxo_set: UtxoSet,
548 _height: Natural,
549) -> Result<UtxoSet> {
550 assert!(
552 !_block.transactions.is_empty(),
553 "Block must have at least one transaction"
554 );
555 assert!(
556 _height <= i64::MAX as u64,
557 "Block height {_height} must fit in i64"
558 );
559 assert!(
560 utxo_set.len() <= u32::MAX as usize,
561 "UTXO set size {} must not exceed maximum",
562 utxo_set.len()
563 );
564 assert!(
566 undo_log.entries.len() <= 10_000,
567 "Undo log entry count {} must be reasonable",
568 undo_log.entries.len()
569 );
570
571 for (i, entry) in undo_log.entries.iter().enumerate() {
574 assert!(i < undo_log.entries.len(), "Entry index {i} out of bounds");
576 if entry.new_utxo.is_some() {
578 utxo_set.remove(&entry.outpoint);
579 }
580
581 if let Some(previous_utxo) = &entry.previous_utxo {
583 utxo_set.insert(entry.outpoint, std::sync::Arc::clone(previous_utxo));
584 }
585 }
586
587 Ok(utxo_set)
588}
589
590#[track_caller] #[allow(clippy::redundant_comparisons)] #[spec_locked("11.3")]
594pub fn should_reorganize(new_chain: &[Block], current_chain: &[Block]) -> Result<bool> {
595 assert!(
597 new_chain.len() <= 10_000,
598 "New chain length {} must be reasonable",
599 new_chain.len()
600 );
601 assert!(
602 current_chain.len() <= 10_000,
603 "Current chain length {} must be reasonable",
604 current_chain.len()
605 );
606
607 if new_chain.len() > current_chain.len() {
609 #[allow(clippy::eq_op)]
611 {
612 assert!(true == true || false == false, "Result must be boolean");
613 }
614 return Ok(true);
615 }
616
617 if new_chain.len() == current_chain.len() {
619 let new_work = calculate_chain_work(new_chain)?;
620 let current_work = calculate_chain_work(current_chain)?;
621 let result = new_work > current_work;
622 return Ok(result);
625 }
626
627 let result = false;
629 Ok(result)
631}
632
633#[spec_locked("11.3")]
641fn calculate_chain_work(chain: &[Block]) -> Result<u128> {
642 let mut total_work = 0u128;
643
644 for block in chain {
645 let target = expand_target(block.header.bits)?;
646 if target > 0 {
649 let work_contribution = if target == u128::MAX {
654 0 } else {
656 u128::MAX.checked_div(target + 1).unwrap_or(0)
658 };
659
660 let old_total = total_work;
663 total_work = total_work.saturating_add(work_contribution);
664
665 debug_assert!(
667 total_work >= old_total,
668 "Total work ({total_work}) must be >= previous total ({old_total})"
669 );
670 }
671 }
673
674 Ok(total_work)
677}
678
679fn expand_target(bits: Natural) -> Result<u128> {
681 let exponent = (bits >> 24) as u8;
682 let mantissa = bits & 0x00ffffff;
683
684 if exponent <= 3 {
685 let shift = 8 * (3 - exponent);
686 Ok((mantissa as u128) >> shift)
687 } else {
688 if exponent > 19 {
691 return Err(crate::error::ConsensusError::InvalidProofOfWork(
692 "Target too large".into(),
693 ));
694 }
695 let shift = 8 * (exponent - 3);
697 let mantissa_u128 = mantissa as u128;
699 let expanded = mantissa_u128.checked_shl(shift as u32).ok_or_else(|| {
700 crate::error::ConsensusError::InvalidProofOfWork("Target expansion overflow".into())
701 })?;
702 Ok(expanded)
703 }
704}
705
706#[allow(dead_code)] fn calculate_tx_id(tx: &Transaction) -> Hash {
709 let mut hash = [0u8; 32];
710 hash[0] = (tx.version & 0xff) as u8;
711 hash[1] = (tx.inputs.len() & 0xff) as u8;
712 hash[2] = (tx.outputs.len() & 0xff) as u8;
713 hash[3] = (tx.lock_time & 0xff) as u8;
714 hash
715}
716
717fn calculate_block_hash(header: &BlockHeader) -> Hash {
722 use sha2::{Digest, Sha256};
723
724 let mut bytes = Vec::with_capacity(80);
726 bytes.extend_from_slice(&header.version.to_le_bytes());
727 bytes.extend_from_slice(&header.prev_block_hash);
728 bytes.extend_from_slice(&header.merkle_root);
729 bytes.extend_from_slice(&header.timestamp.to_le_bytes());
730 bytes.extend_from_slice(&header.bits.to_le_bytes());
731 bytes.extend_from_slice(&header.nonce.to_le_bytes());
732
733 let first_hash = Sha256::digest(&bytes);
735 let second_hash = Sha256::digest(first_hash);
736
737 let mut hash = [0u8; 32];
738 hash.copy_from_slice(&second_hash);
739 hash
740}
741
742mod undo_entry_serde {
747 use crate::types::UTXO;
748 use serde::{Deserialize, Deserializer, Serialize, Serializer};
749 use std::sync::Arc;
750
751 pub fn serialize<S>(opt: &Option<Arc<UTXO>>, s: S) -> Result<S::Ok, S::Error>
752 where
753 S: Serializer,
754 {
755 opt.as_ref().map(|a| a.as_ref()).serialize(s)
756 }
757
758 pub fn deserialize<'de, D>(d: D) -> Result<Option<Arc<UTXO>>, D::Error>
759 where
760 D: Deserializer<'de>,
761 {
762 Option::<UTXO>::deserialize(d).map(|opt| opt.map(Arc::new))
763 }
764}
765
766#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
771pub struct UndoEntry {
772 pub outpoint: OutPoint,
774 #[serde(with = "undo_entry_serde")]
776 pub previous_utxo: Option<std::sync::Arc<UTXO>>,
777 #[serde(with = "undo_entry_serde")]
779 pub new_utxo: Option<std::sync::Arc<UTXO>>,
780}
781
782#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
790pub struct BlockUndoLog {
791 pub entries: Vec<UndoEntry>,
794}
795
796impl BlockUndoLog {
797 pub fn new() -> Self {
799 Self {
800 entries: Vec::new(),
801 }
802 }
803
804 pub fn push(&mut self, entry: UndoEntry) {
806 self.entries.push(entry);
807 }
808
809 pub fn is_empty(&self) -> bool {
811 self.entries.is_empty()
812 }
813}
814
815impl Default for BlockUndoLog {
816 fn default() -> Self {
817 Self::new()
818 }
819}
820
821#[derive(Debug, Clone)]
823pub struct ReorganizationResult {
824 pub new_utxo_set: UtxoSet,
825 pub new_height: Natural,
826 pub common_ancestor: BlockHeader,
827 pub disconnected_blocks: Vec<Block>,
828 pub connected_blocks: Vec<Block>,
829 pub reorganization_depth: usize,
830 pub connected_block_undo_logs: HashMap<Hash, BlockUndoLog>,
833}
834
835#[cfg(test)]
849mod property_tests {
850 use super::*;
851 use proptest::prelude::*;
852
853 fn chain_len_range() -> std::ops::Range<usize> {
855 if std::env::var("CARGO_TARPAULIN").is_ok() || std::env::var("TARPAULIN").is_ok() {
856 1..3 } else {
858 1..5
859 }
860 }
861
862 fn chain_len_range_det() -> std::ops::Range<usize> {
864 if std::env::var("CARGO_TARPAULIN").is_ok() || std::env::var("TARPAULIN").is_ok() {
865 0..3 } else {
867 0..10
868 }
869 }
870
871 fn arb_block() -> impl Strategy<Value = Block> {
873 (0x00000000u64..0x1d00ffffu64).prop_map(|bits| Block {
874 header: BlockHeader {
875 version: 1,
876 prev_block_hash: [0u8; 32],
877 merkle_root: [0u8; 32],
878 timestamp: 0,
879 bits,
880 nonce: 0,
881 },
882 transactions: Box::new([]),
883 })
884 }
885
886 proptest! {
888 #[test]
889 fn prop_should_reorganize_max_work(
890 new_chain in proptest::collection::vec(arb_block(), chain_len_range()),
891 current_chain in proptest::collection::vec(arb_block(), chain_len_range())
892 ) {
893 let new_work = calculate_chain_work(&new_chain);
895 let current_work = calculate_chain_work(¤t_chain);
896
897 if let (Ok(new_w), Ok(current_w)) = (new_work, current_work) {
899 let should_reorg = should_reorganize(&new_chain, ¤t_chain).unwrap_or(false);
901
902 if new_chain.len() > current_chain.len() {
908 prop_assert!(should_reorg, "Must reorganize when new chain is longer");
910 } else if new_chain.len() == current_chain.len() {
911 if new_w > current_w {
913 prop_assert!(should_reorg, "Must reorganize when new chain has more work (equal length)");
914 } else {
915 prop_assert!(!should_reorg, "Must not reorganize when new chain has less or equal work (equal length)");
916 }
917 } else {
918 prop_assert!(!should_reorg, "Must not reorganize when new chain is shorter");
920 }
921 }
922 }
924 }
925
926 proptest! {
928 #[test]
929 fn prop_calculate_chain_work_deterministic(
930 chain in proptest::collection::vec(arb_block(), chain_len_range_det())
931 ) {
932 let work1 = calculate_chain_work(&chain);
934 let work2 = calculate_chain_work(&chain);
935
936 match (work1, work2) {
938 (Ok(w1), Ok(w2)) => {
939 prop_assert_eq!(w1, w2, "Chain work calculation must be deterministic");
940 },
941 (Err(_), Err(_)) => {
942 },
944 _ => {
945 prop_assert!(false, "Chain work calculation must be deterministic (both succeed or both fail)");
946 }
947 }
948 }
949 }
950
951 proptest! {
953 #[test]
954 fn prop_expand_target_valid_range(
955 bits in 0x00000000u64..0x1d00ffffu64
956 ) {
957 let result = expand_target(bits);
958
959 match result {
960 Ok(target) => {
961 let _ = target;
965 },
966 Err(_) => {
967 }
969 }
970 }
971 }
972
973 proptest! {
975 #[test]
976 fn prop_should_reorganize_equal_length(
977 chain1 in proptest::collection::vec(arb_block(), 1..3),
978 chain2 in proptest::collection::vec(arb_block(), 1..3)
979 ) {
980 let len = chain1.len().min(chain2.len());
982 let chain1 = &chain1[..len];
983 let chain2 = &chain2[..len];
984
985 let work1 = calculate_chain_work(chain1);
986 let work2 = calculate_chain_work(chain2);
987
988 if let (Ok(w1), Ok(w2)) = (work1, work2) {
990 let should_reorg = should_reorganize(chain1, chain2).unwrap_or(false);
991
992 if w1 > w2 {
994 prop_assert!(should_reorg, "Must reorganize when first chain has more work");
995 } else {
996 prop_assert!(!should_reorg, "Must not reorganize when first chain has less or equal work");
997 }
998 }
999 }
1001 }
1002}
1003
1004#[cfg(test)]
1005mod tests {
1006 use super::*;
1007
1008 #[test]
1009 fn test_should_reorganize_longer_chain() {
1010 let new_chain = vec![create_test_block(), create_test_block()];
1011 let current_chain = vec![create_test_block()];
1012
1013 assert!(should_reorganize(&new_chain, ¤t_chain).unwrap());
1014 }
1015
1016 #[test]
1017 fn test_should_reorganize_same_length_more_work() {
1018 let mut new_chain = vec![create_test_block()];
1019 let mut current_chain = vec![create_test_block()];
1020
1021 new_chain[0].header.bits = 0x0200ffff; current_chain[0].header.bits = 0x0300ffff; assert!(should_reorganize(&new_chain, ¤t_chain).unwrap());
1026 }
1027
1028 #[test]
1029 fn test_should_not_reorganize_shorter_chain() {
1030 let new_chain = vec![create_test_block()];
1031 let current_chain = vec![create_test_block(), create_test_block()];
1032
1033 assert!(!should_reorganize(&new_chain, ¤t_chain).unwrap());
1034 }
1035
1036 #[test]
1037 fn test_find_common_ancestor() {
1038 let new_chain = vec![create_test_block()];
1039 let current_chain = vec![create_test_block()];
1040
1041 let ancestor = find_common_ancestor(&new_chain, ¤t_chain).unwrap();
1042 assert_eq!(ancestor.header.version, 4);
1043 assert_eq!(ancestor.new_chain_index, 0);
1044 assert_eq!(ancestor.current_chain_index, 0);
1045 }
1046
1047 #[test]
1048 fn test_find_common_ancestor_empty_chain() {
1049 let new_chain = vec![];
1050 let current_chain = vec![create_test_block()];
1051
1052 let result = find_common_ancestor(&new_chain, ¤t_chain);
1053 assert!(result.is_err());
1054 }
1055
1056 #[test]
1057 fn test_calculate_chain_work() {
1058 let mut block = create_test_block();
1059 block.header.bits = 0x0300ffff;
1061 let chain = vec![block];
1062 let work = calculate_chain_work(&chain).unwrap();
1063 assert!(work > 0);
1064 }
1065
1066 #[test]
1067 fn test_reorganize_chain() {
1068 let ancestor = create_test_block_at_height(0);
1072 let mut new_block = create_test_block_at_height(1);
1073 new_block.header.nonce = 42; let new_chain = vec![ancestor.clone(), new_block];
1077 let current_chain = vec![ancestor];
1078 let utxo_set = UtxoSet::default();
1079
1080 let result = reorganize_chain(
1082 &new_chain,
1083 ¤t_chain,
1084 utxo_set,
1085 1,
1086 crate::types::Network::Regtest,
1087 );
1088 match result {
1089 Ok(reorg_result) => {
1090 assert_eq!(reorg_result.new_height, 2);
1093 assert_eq!(reorg_result.connected_blocks.len(), 1);
1094 assert_eq!(reorg_result.connected_block_undo_logs.len(), 1);
1095 }
1096 Err(_) => {
1097 }
1099 }
1100 }
1101
1102 #[test]
1103 fn test_reorganize_chain_deep_reorg() {
1104 let mut block1 = create_test_block();
1106 block1.header.nonce = 1;
1107 let mut block2 = create_test_block();
1108 block2.header.nonce = 2;
1109 let mut block3 = create_test_block();
1110 block3.header.nonce = 3;
1111 let new_chain = vec![block1, block2, block3];
1112
1113 let mut current_block1 = create_test_block();
1114 current_block1.header.nonce = 10;
1115 let mut current_block2 = create_test_block();
1116 current_block2.header.nonce = 11;
1117 let current_chain = vec![current_block1, current_block2];
1118 let utxo_set = UtxoSet::default();
1119
1120 let result = reorganize_chain(
1121 &new_chain,
1122 ¤t_chain,
1123 utxo_set,
1124 2,
1125 crate::types::Network::Regtest,
1126 );
1127 match result {
1128 Ok(reorg_result) => {
1129 assert_eq!(reorg_result.connected_blocks.len(), 3);
1130 assert_eq!(reorg_result.reorganization_depth, 2);
1131 assert_eq!(reorg_result.connected_block_undo_logs.len(), 3);
1133 }
1134 Err(_) => {
1135 }
1137 }
1138 }
1139
1140 #[test]
1141 fn test_undo_log_storage_and_retrieval() {
1142 use crate::block::connect_block;
1143 use crate::segwit::Witness;
1144
1145 let block = create_test_block_at_height(1);
1146 let mut utxo_set = UtxoSet::default();
1147
1148 let tx_id = calculate_tx_id(&block.transactions[0]);
1150 let outpoint = OutPoint {
1151 hash: tx_id,
1152 index: 0,
1153 };
1154 let utxo = UTXO {
1155 value: 5_000_000_000,
1156 script_pubkey: vec![0x51].into(),
1157 height: 1,
1158 is_coinbase: false,
1159 };
1160 utxo_set.insert(outpoint.clone(), std::sync::Arc::new(utxo.clone()));
1161
1162 let witnesses: Vec<Vec<Witness>> = block
1164 .transactions
1165 .iter()
1166 .map(|tx| tx.inputs.iter().map(|_| Vec::new()).collect())
1167 .collect();
1168 let ctx = BlockValidationContext::for_network(crate::types::Network::Regtest);
1169 let (result, new_utxo_set, undo_log) =
1170 connect_block(&block, &witnesses, utxo_set.clone(), 1, &ctx).unwrap();
1171
1172 assert!(matches!(result, crate::types::ValidationResult::Valid));
1173
1174 assert!(
1176 !undo_log.entries.is_empty(),
1177 "Undo log should contain entries"
1178 );
1179
1180 let block_hash = calculate_block_hash(&block.header);
1182
1183 let mut undo_log_storage: HashMap<Hash, BlockUndoLog> = HashMap::new();
1185 undo_log_storage.insert(block_hash, undo_log.clone());
1186
1187 let retrieved_undo_log = undo_log_storage.get(&block_hash);
1189 assert!(
1190 retrieved_undo_log.is_some(),
1191 "Should be able to retrieve undo log"
1192 );
1193 assert_eq!(
1194 retrieved_undo_log.unwrap().entries.len(),
1195 undo_log.entries.len()
1196 );
1197
1198 let disconnected_utxo_set = disconnect_block(&block, &undo_log, new_utxo_set, 1).unwrap();
1200
1201 assert!(
1203 disconnected_utxo_set.contains_key(&outpoint),
1204 "Disconnected UTXO set should contain restored UTXO"
1205 );
1206 }
1207
1208 #[test]
1209 fn test_reorganize_with_undo_log_callback() {
1210 use crate::block::connect_block;
1211 use crate::segwit::Witness;
1212
1213 let block = create_test_block_at_height(1);
1215 let utxo_set = UtxoSet::default();
1216 let witnesses: Vec<Vec<Witness>> = block
1217 .transactions
1218 .iter()
1219 .map(|tx| tx.inputs.iter().map(|_| Vec::new()).collect())
1220 .collect();
1221
1222 let ctx = BlockValidationContext::for_network(crate::types::Network::Regtest);
1223 let (result, connected_utxo_set, undo_log) =
1224 connect_block(&block, &witnesses, utxo_set.clone(), 1, &ctx).unwrap();
1225
1226 if !matches!(result, crate::types::ValidationResult::Valid) {
1227 eprintln!("Block validation failed: {:?}", result);
1228 }
1229 assert!(matches!(result, crate::types::ValidationResult::Valid));
1230
1231 let block_hash = calculate_block_hash(&block.header);
1233 let mut undo_log_storage: HashMap<Hash, BlockUndoLog> = HashMap::new();
1234 undo_log_storage.insert(block_hash, undo_log);
1235
1236 let get_undo_log =
1238 |hash: &Hash| -> Option<BlockUndoLog> { undo_log_storage.get(hash).cloned() };
1239
1240 let mut new_block = create_test_block_at_height(2);
1243 new_block.header.nonce = 42; let new_chain = vec![block.clone(), new_block];
1245 let current_chain = vec![block];
1246 let empty_witnesses: Vec<Vec<Vec<Witness>>> = new_chain
1247 .iter()
1248 .map(|b| {
1249 b.transactions
1250 .iter()
1251 .map(|tx| tx.inputs.iter().map(|_| Vec::new()).collect())
1252 .collect()
1253 })
1254 .collect();
1255
1256 let reorg_result = reorganize_chain_with_witnesses(
1257 &new_chain,
1258 &empty_witnesses,
1259 None,
1260 ¤t_chain,
1261 connected_utxo_set,
1262 1,
1263 None::<fn(&Block) -> Option<Vec<Witness>>>,
1264 None::<fn(Natural) -> Option<Vec<BlockHeader>>>,
1265 Some(get_undo_log),
1266 None::<fn(&Hash, &BlockUndoLog) -> Result<()>>, crate::types::Network::Regtest,
1268 );
1269
1270 match reorg_result {
1272 Ok(result) => {
1273 assert!(!result.connected_block_undo_logs.is_empty());
1275 }
1276 Err(_) => {
1277 }
1279 }
1280 }
1281
1282 #[test]
1283 fn test_reorganize_chain_empty_new_chain() {
1284 let new_chain = vec![];
1285 let current_chain = vec![create_test_block()];
1286 let utxo_set = UtxoSet::default();
1287
1288 let result = reorganize_chain(
1289 &new_chain,
1290 ¤t_chain,
1291 utxo_set,
1292 1,
1293 crate::types::Network::Regtest,
1294 );
1295 assert!(result.is_err());
1296 }
1297
1298 #[test]
1299 fn test_reorganize_chain_empty_current_chain() {
1300 let new_chain = vec![create_test_block()];
1301 let current_chain = vec![];
1302 let utxo_set = UtxoSet::default();
1303
1304 let result = reorganize_chain(
1305 &new_chain,
1306 ¤t_chain,
1307 utxo_set,
1308 0,
1309 crate::types::Network::Regtest,
1310 );
1311 assert!(result.is_err());
1312 }
1313
1314 #[test]
1315 fn test_disconnect_block() {
1316 let block = create_test_block();
1317 let mut utxo_set = UtxoSet::default();
1318
1319 let tx_id = calculate_tx_id(&block.transactions[0]);
1321 let outpoint = OutPoint {
1322 hash: tx_id,
1323 index: 0,
1324 };
1325 let utxo = UTXO {
1326 value: 50_000_000_000,
1327 script_pubkey: vec![0x51].into(),
1328 height: 1,
1329 is_coinbase: false,
1330 };
1331 utxo_set.insert(outpoint, std::sync::Arc::new(utxo));
1332
1333 let empty_undo_log = BlockUndoLog::new();
1335 let result = disconnect_block(&block, &empty_undo_log, utxo_set, 1);
1336 assert!(result.is_ok());
1337 }
1338
1339 #[test]
1340 fn test_calculate_chain_work_empty_chain() {
1341 let chain = vec![];
1342 let work = calculate_chain_work(&chain).unwrap();
1343 assert_eq!(work, 0);
1344 }
1345
1346 #[test]
1347 fn test_calculate_chain_work_multiple_blocks() {
1348 let mut chain = vec![create_test_block(), create_test_block()];
1349 chain[0].header.bits = 0x0300ffff;
1351 chain[1].header.bits = 0x0200ffff;
1352
1353 let work = calculate_chain_work(&chain).unwrap();
1354 assert!(work > 0);
1355 }
1356
1357 #[test]
1358 fn test_expand_target_edge_cases() {
1359 let result = expand_target(0x00000000);
1361 assert!(result.is_ok());
1362
1363 let result = expand_target(0x03ffffff);
1365 assert!(result.is_ok());
1366
1367 let result = expand_target(0x14000000); assert!(result.is_err());
1370 }
1371
1372 #[test]
1373 fn test_calculate_tx_id_different_transactions() {
1374 let tx1 = Transaction {
1375 version: 1,
1376 inputs: vec![].into(),
1377 outputs: vec![].into(),
1378 lock_time: 0,
1379 };
1380
1381 let tx2 = Transaction {
1382 version: 2,
1383 inputs: vec![].into(),
1384 outputs: vec![].into(),
1385 lock_time: 0,
1386 };
1387
1388 let id1 = calculate_tx_id(&tx1);
1389 let id2 = calculate_tx_id(&tx2);
1390
1391 assert_ne!(id1, id2);
1392 }
1393
1394 fn encode_bip34_height(height: u64) -> Vec<u8> {
1399 if height == 0 {
1400 return vec![0x00, 0xff]; }
1403 let mut height_bytes = Vec::new();
1404 let mut n = height;
1405 while n > 0 {
1406 height_bytes.push((n & 0xff) as u8);
1407 n >>= 8;
1408 }
1409 if height_bytes.last().map_or(false, |&b| b & 0x80 != 0) {
1411 height_bytes.push(0x00);
1412 }
1413 let mut script_sig = Vec::with_capacity(1 + height_bytes.len() + 1);
1414 script_sig.push(height_bytes.len() as u8); script_sig.extend_from_slice(&height_bytes);
1416 if script_sig.len() < 2 {
1418 script_sig.push(0xff);
1419 }
1420 script_sig
1421 }
1422
1423 fn create_test_block_at_height(height: u64) -> Block {
1425 use crate::mining::calculate_merkle_root;
1426
1427 let script_sig = encode_bip34_height(height);
1428 let coinbase_tx = Transaction {
1429 version: 1,
1430 inputs: vec![TransactionInput {
1431 prevout: OutPoint {
1432 hash: [0; 32].into(),
1433 index: 0xffffffff,
1434 },
1435 script_sig,
1436 sequence: 0xffffffff,
1437 }]
1438 .into(),
1439 outputs: vec![TransactionOutput {
1440 value: 5_000_000_000,
1441 script_pubkey: vec![0x51].into(),
1442 }]
1443 .into(),
1444 lock_time: 0,
1445 };
1446
1447 let merkle_root =
1448 calculate_merkle_root(&[coinbase_tx.clone()]).expect("Failed to calculate merkle root");
1449
1450 Block {
1451 header: BlockHeader {
1452 version: 4,
1453 prev_block_hash: [0; 32],
1454 merkle_root,
1455 timestamp: 1231006505,
1456 bits: 0x207fffff, nonce: 0,
1458 },
1459 transactions: vec![coinbase_tx].into_boxed_slice(),
1460 }
1461 }
1462
1463 fn create_test_block() -> Block {
1465 create_test_block_at_height(0)
1466 }
1467}