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
319pub fn update_mempool_after_reorg<F>(
383 mempool: &mut crate::mempool::Mempool,
384 reorg_result: &ReorganizationResult,
385 utxo_set: &UtxoSet,
386 get_tx_by_id: Option<F>,
387) -> Result<Vec<Hash>>
388where
389 F: Fn(&Hash) -> Option<Transaction>,
390{
391 use crate::mempool::update_mempool_after_block;
392
393 let mut all_removed = Vec::new();
394
395 for block in &reorg_result.connected_blocks {
397 let removed = update_mempool_after_block(mempool, block, utxo_set)?;
398 all_removed.extend(removed);
399 }
400
401 let mut spent_outpoints = std::collections::HashSet::new();
404 for block in &reorg_result.connected_blocks {
405 for tx in &block.transactions {
406 if !crate::transaction::is_coinbase(tx) {
407 for input in &tx.inputs {
408 spent_outpoints.insert(input.prevout);
409 }
410 }
411 }
412 }
413
414 if let Some(lookup) = get_tx_by_id {
416 let mut invalid_tx_ids = Vec::new();
417 for &tx_id in mempool.iter() {
418 if let Some(tx) = lookup(&tx_id) {
419 for input in &tx.inputs {
421 if spent_outpoints.contains(&input.prevout) {
422 invalid_tx_ids.push(tx_id);
423 break;
424 }
425 }
426 }
427 }
428
429 for tx_id in invalid_tx_ids {
431 if mempool.remove(&tx_id) {
432 all_removed.push(tx_id);
433 }
434 }
435 }
436
437 Ok(all_removed)
445}
446
447pub fn update_mempool_after_reorg_simple(
449 mempool: &mut crate::mempool::Mempool,
450 reorg_result: &ReorganizationResult,
451 utxo_set: &UtxoSet,
452) -> Result<Vec<Hash>> {
453 update_mempool_after_reorg(
454 mempool,
455 reorg_result,
456 utxo_set,
457 None::<fn(&Hash) -> Option<Transaction>>,
458 )
459}
460
461struct CommonAncestorResult {
463 header: BlockHeader,
464 new_chain_index: usize,
465 current_chain_index: usize,
466}
467
468fn find_common_ancestor(
475 new_chain: &[Block],
476 current_chain: &[Block],
477) -> Result<CommonAncestorResult> {
478 if new_chain.is_empty() || current_chain.is_empty() {
479 return Err(crate::error::ConsensusError::ConsensusRuleViolation(
480 "Cannot find common ancestor: empty chain".into(),
481 ));
482 }
483
484 let min_len = new_chain.len().min(current_chain.len());
486
487 for distance_from_tip in 0..min_len {
490 let new_idx = new_chain.len() - 1 - distance_from_tip;
491 let current_idx = current_chain.len() - 1 - distance_from_tip;
492
493 let new_hash = calculate_block_hash(&new_chain[new_idx].header);
494 let current_hash = calculate_block_hash(¤t_chain[current_idx].header);
495
496 if new_hash == current_hash {
498 return Ok(CommonAncestorResult {
499 header: new_chain[new_idx].header.clone(),
500 new_chain_index: new_idx,
501 current_chain_index: current_idx,
502 });
503 }
504 }
505
506 if !new_chain.is_empty() && !current_chain.is_empty() {
509 let new_genesis_hash = calculate_block_hash(&new_chain[0].header);
510 let current_genesis_hash = calculate_block_hash(¤t_chain[0].header);
511 if new_genesis_hash == current_genesis_hash {
512 return Ok(CommonAncestorResult {
513 header: new_chain[0].header.clone(),
514 new_chain_index: 0,
515 current_chain_index: 0,
516 });
517 }
518 }
519
520 Err(crate::error::ConsensusError::ConsensusRuleViolation(
522 "Chains do not share a common ancestor".into(),
523 ))
524}
525
526#[spec_locked("11.3.1", "DisconnectBlock")]
539fn disconnect_block(
540 _block: &Block,
541 undo_log: &BlockUndoLog,
542 mut utxo_set: UtxoSet,
543 _height: Natural,
544) -> Result<UtxoSet> {
545 assert!(
547 !_block.transactions.is_empty(),
548 "Block must have at least one transaction"
549 );
550 assert!(
551 _height <= i64::MAX as u64,
552 "Block height {_height} must fit in i64"
553 );
554 assert!(
555 utxo_set.len() <= u32::MAX as usize,
556 "UTXO set size {} must not exceed maximum",
557 utxo_set.len()
558 );
559 assert!(
561 undo_log.entries.len() <= 10_000,
562 "Undo log entry count {} must be reasonable",
563 undo_log.entries.len()
564 );
565
566 for (i, entry) in undo_log.entries.iter().enumerate() {
569 assert!(i < undo_log.entries.len(), "Entry index {i} out of bounds");
571 if entry.new_utxo.is_some() {
573 utxo_set.remove(&entry.outpoint);
574 }
575
576 if let Some(previous_utxo) = &entry.previous_utxo {
578 utxo_set.insert(entry.outpoint, std::sync::Arc::clone(previous_utxo));
579 }
580 }
581
582 Ok(utxo_set)
583}
584
585#[track_caller] #[allow(clippy::redundant_comparisons)] #[spec_locked("11.3", "ShouldReorganize")]
589pub fn should_reorganize(new_chain: &[Block], current_chain: &[Block]) -> Result<bool> {
590 assert!(
592 new_chain.len() <= 10_000,
593 "New chain length {} must be reasonable",
594 new_chain.len()
595 );
596 assert!(
597 current_chain.len() <= 10_000,
598 "Current chain length {} must be reasonable",
599 current_chain.len()
600 );
601
602 if new_chain.len() > current_chain.len() {
604 #[allow(clippy::eq_op)]
606 {
607 assert!(true == true || false == false, "Result must be boolean");
608 }
609 return Ok(true);
610 }
611
612 if new_chain.len() == current_chain.len() {
614 let new_work = calculate_chain_work(new_chain)?;
615 let current_work = calculate_chain_work(current_chain)?;
616 let result = new_work > current_work;
617 return Ok(result);
620 }
621
622 let result = false;
624 Ok(result)
626}
627
628#[spec_locked("11.3", "CalculateChainWork")]
636fn calculate_chain_work(chain: &[Block]) -> Result<u128> {
637 let mut total_work = 0u128;
638
639 for block in chain {
640 let target = expand_target(block.header.bits)?;
641 if target > 0 {
644 let work_contribution = if target == u128::MAX {
649 0 } else {
651 u128::MAX.checked_div(target + 1).unwrap_or(0)
653 };
654
655 let old_total = total_work;
658 total_work = total_work.saturating_add(work_contribution);
659
660 debug_assert!(
662 total_work >= old_total,
663 "Total work ({total_work}) must be >= previous total ({old_total})"
664 );
665 }
666 }
668
669 Ok(total_work)
672}
673
674fn expand_target(bits: Natural) -> Result<u128> {
676 let exponent = (bits >> 24) as u8;
677 let mantissa = bits & 0x00ffffff;
678
679 if exponent <= 3 {
680 let shift = 8 * (3 - exponent);
681 Ok((mantissa as u128) >> shift)
682 } else {
683 if exponent > 19 {
686 return Err(crate::error::ConsensusError::InvalidProofOfWork(
687 "Target too large".into(),
688 ));
689 }
690 let shift = 8 * (exponent - 3);
692 let mantissa_u128 = mantissa as u128;
694 let expanded = mantissa_u128.checked_shl(shift as u32).ok_or_else(|| {
695 crate::error::ConsensusError::InvalidProofOfWork("Target expansion overflow".into())
696 })?;
697 Ok(expanded)
698 }
699}
700
701#[allow(dead_code)] fn calculate_tx_id(tx: &Transaction) -> Hash {
704 let mut hash = [0u8; 32];
705 hash[0] = (tx.version & 0xff) as u8;
706 hash[1] = (tx.inputs.len() & 0xff) as u8;
707 hash[2] = (tx.outputs.len() & 0xff) as u8;
708 hash[3] = (tx.lock_time & 0xff) as u8;
709 hash
710}
711
712fn calculate_block_hash(header: &BlockHeader) -> Hash {
717 use sha2::{Digest, Sha256};
718
719 let mut bytes = Vec::with_capacity(80);
721 bytes.extend_from_slice(&header.version.to_le_bytes());
722 bytes.extend_from_slice(&header.prev_block_hash);
723 bytes.extend_from_slice(&header.merkle_root);
724 bytes.extend_from_slice(&header.timestamp.to_le_bytes());
725 bytes.extend_from_slice(&header.bits.to_le_bytes());
726 bytes.extend_from_slice(&header.nonce.to_le_bytes());
727
728 let first_hash = Sha256::digest(&bytes);
730 let second_hash = Sha256::digest(first_hash);
731
732 let mut hash = [0u8; 32];
733 hash.copy_from_slice(&second_hash);
734 hash
735}
736
737mod undo_entry_serde {
742 use crate::types::UTXO;
743 use serde::{Deserialize, Deserializer, Serialize, Serializer};
744 use std::sync::Arc;
745
746 pub fn serialize<S>(opt: &Option<Arc<UTXO>>, s: S) -> Result<S::Ok, S::Error>
747 where
748 S: Serializer,
749 {
750 opt.as_ref().map(|a| a.as_ref()).serialize(s)
751 }
752
753 pub fn deserialize<'de, D>(d: D) -> Result<Option<Arc<UTXO>>, D::Error>
754 where
755 D: Deserializer<'de>,
756 {
757 Option::<UTXO>::deserialize(d).map(|opt| opt.map(Arc::new))
758 }
759}
760
761#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
766pub struct UndoEntry {
767 pub outpoint: OutPoint,
769 #[serde(with = "undo_entry_serde")]
771 pub previous_utxo: Option<std::sync::Arc<UTXO>>,
772 #[serde(with = "undo_entry_serde")]
774 pub new_utxo: Option<std::sync::Arc<UTXO>>,
775}
776
777#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
785pub struct BlockUndoLog {
786 pub entries: Vec<UndoEntry>,
789}
790
791impl BlockUndoLog {
792 pub fn new() -> Self {
794 Self {
795 entries: Vec::new(),
796 }
797 }
798
799 pub fn push(&mut self, entry: UndoEntry) {
801 self.entries.push(entry);
802 }
803
804 pub fn is_empty(&self) -> bool {
806 self.entries.is_empty()
807 }
808}
809
810impl Default for BlockUndoLog {
811 fn default() -> Self {
812 Self::new()
813 }
814}
815
816#[derive(Debug, Clone)]
818pub struct ReorganizationResult {
819 pub new_utxo_set: UtxoSet,
820 pub new_height: Natural,
821 pub common_ancestor: BlockHeader,
822 pub disconnected_blocks: Vec<Block>,
823 pub connected_blocks: Vec<Block>,
824 pub reorganization_depth: usize,
825 pub connected_block_undo_logs: HashMap<Hash, BlockUndoLog>,
828}
829
830#[cfg(test)]
844mod property_tests {
845 use super::*;
846 use proptest::prelude::*;
847
848 fn chain_len_range() -> std::ops::Range<usize> {
850 if std::env::var("CARGO_TARPAULIN").is_ok() || std::env::var("TARPAULIN").is_ok() {
851 1..3 } else {
853 1..5
854 }
855 }
856
857 fn chain_len_range_det() -> std::ops::Range<usize> {
859 if std::env::var("CARGO_TARPAULIN").is_ok() || std::env::var("TARPAULIN").is_ok() {
860 0..3 } else {
862 0..10
863 }
864 }
865
866 fn arb_block() -> impl Strategy<Value = Block> {
868 (0x00000000u64..0x1d00ffffu64).prop_map(|bits| Block {
869 header: BlockHeader {
870 version: 1,
871 prev_block_hash: [0u8; 32],
872 merkle_root: [0u8; 32],
873 timestamp: 0,
874 bits,
875 nonce: 0,
876 },
877 transactions: Box::new([]),
878 })
879 }
880
881 proptest! {
883 #[test]
884 fn prop_should_reorganize_max_work(
885 new_chain in proptest::collection::vec(arb_block(), chain_len_range()),
886 current_chain in proptest::collection::vec(arb_block(), chain_len_range())
887 ) {
888 let new_work = calculate_chain_work(&new_chain);
890 let current_work = calculate_chain_work(¤t_chain);
891
892 if let (Ok(new_w), Ok(current_w)) = (new_work, current_work) {
894 let should_reorg = should_reorganize(&new_chain, ¤t_chain).unwrap_or(false);
896
897 if new_chain.len() > current_chain.len() {
903 prop_assert!(should_reorg, "Must reorganize when new chain is longer");
905 } else if new_chain.len() == current_chain.len() {
906 if new_w > current_w {
908 prop_assert!(should_reorg, "Must reorganize when new chain has more work (equal length)");
909 } else {
910 prop_assert!(!should_reorg, "Must not reorganize when new chain has less or equal work (equal length)");
911 }
912 } else {
913 prop_assert!(!should_reorg, "Must not reorganize when new chain is shorter");
915 }
916 }
917 }
919 }
920
921 proptest! {
923 #[test]
924 fn prop_calculate_chain_work_deterministic(
925 chain in proptest::collection::vec(arb_block(), chain_len_range_det())
926 ) {
927 let work1 = calculate_chain_work(&chain);
929 let work2 = calculate_chain_work(&chain);
930
931 match (work1, work2) {
933 (Ok(w1), Ok(w2)) => {
934 prop_assert_eq!(w1, w2, "Chain work calculation must be deterministic");
935 },
936 (Err(_), Err(_)) => {
937 },
939 _ => {
940 prop_assert!(false, "Chain work calculation must be deterministic (both succeed or both fail)");
941 }
942 }
943 }
944 }
945
946 proptest! {
948 #[test]
949 fn prop_expand_target_valid_range(
950 bits in 0x00000000u64..0x1d00ffffu64
951 ) {
952 let result = expand_target(bits);
953
954 match result {
955 Ok(target) => {
956 let _ = target;
960 },
961 Err(_) => {
962 }
964 }
965 }
966 }
967
968 proptest! {
970 #[test]
971 fn prop_should_reorganize_equal_length(
972 chain1 in proptest::collection::vec(arb_block(), 1..3),
973 chain2 in proptest::collection::vec(arb_block(), 1..3)
974 ) {
975 let len = chain1.len().min(chain2.len());
977 let chain1 = &chain1[..len];
978 let chain2 = &chain2[..len];
979
980 let work1 = calculate_chain_work(chain1);
981 let work2 = calculate_chain_work(chain2);
982
983 if let (Ok(w1), Ok(w2)) = (work1, work2) {
985 let should_reorg = should_reorganize(chain1, chain2).unwrap_or(false);
986
987 if w1 > w2 {
989 prop_assert!(should_reorg, "Must reorganize when first chain has more work");
990 } else {
991 prop_assert!(!should_reorg, "Must not reorganize when first chain has less or equal work");
992 }
993 }
994 }
996 }
997}
998
999#[cfg(test)]
1000mod tests {
1001 use super::*;
1002
1003 #[test]
1004 fn test_should_reorganize_longer_chain() {
1005 let new_chain = vec![create_test_block(), create_test_block()];
1006 let current_chain = vec![create_test_block()];
1007
1008 assert!(should_reorganize(&new_chain, ¤t_chain).unwrap());
1009 }
1010
1011 #[test]
1012 fn test_should_reorganize_same_length_more_work() {
1013 let mut new_chain = vec![create_test_block()];
1014 let mut current_chain = vec![create_test_block()];
1015
1016 new_chain[0].header.bits = 0x0200ffff; current_chain[0].header.bits = 0x0300ffff; assert!(should_reorganize(&new_chain, ¤t_chain).unwrap());
1021 }
1022
1023 #[test]
1024 fn test_should_not_reorganize_shorter_chain() {
1025 let new_chain = vec![create_test_block()];
1026 let current_chain = vec![create_test_block(), create_test_block()];
1027
1028 assert!(!should_reorganize(&new_chain, ¤t_chain).unwrap());
1029 }
1030
1031 #[test]
1032 fn test_find_common_ancestor() {
1033 let new_chain = vec![create_test_block()];
1034 let current_chain = vec![create_test_block()];
1035
1036 let ancestor = find_common_ancestor(&new_chain, ¤t_chain).unwrap();
1037 assert_eq!(ancestor.header.version, 4);
1038 assert_eq!(ancestor.new_chain_index, 0);
1039 assert_eq!(ancestor.current_chain_index, 0);
1040 }
1041
1042 #[test]
1043 fn test_find_common_ancestor_empty_chain() {
1044 let new_chain = vec![];
1045 let current_chain = vec![create_test_block()];
1046
1047 let result = find_common_ancestor(&new_chain, ¤t_chain);
1048 assert!(result.is_err());
1049 }
1050
1051 #[test]
1052 fn test_calculate_chain_work() {
1053 let mut block = create_test_block();
1054 block.header.bits = 0x0300ffff;
1056 let chain = vec![block];
1057 let work = calculate_chain_work(&chain).unwrap();
1058 assert!(work > 0);
1059 }
1060
1061 #[test]
1062 fn test_reorganize_chain() {
1063 let ancestor = create_test_block_at_height(0);
1067 let mut new_block = create_test_block_at_height(1);
1068 new_block.header.nonce = 42; let new_chain = vec![ancestor.clone(), new_block];
1072 let current_chain = vec![ancestor];
1073 let utxo_set = UtxoSet::default();
1074
1075 let result = reorganize_chain(
1077 &new_chain,
1078 ¤t_chain,
1079 utxo_set,
1080 1,
1081 crate::types::Network::Regtest,
1082 );
1083 match result {
1084 Ok(reorg_result) => {
1085 assert_eq!(reorg_result.new_height, 2);
1088 assert_eq!(reorg_result.connected_blocks.len(), 1);
1089 assert_eq!(reorg_result.connected_block_undo_logs.len(), 1);
1090 }
1091 Err(_) => {
1092 }
1094 }
1095 }
1096
1097 #[test]
1098 fn test_reorganize_chain_deep_reorg() {
1099 let mut block1 = create_test_block();
1101 block1.header.nonce = 1;
1102 let mut block2 = create_test_block();
1103 block2.header.nonce = 2;
1104 let mut block3 = create_test_block();
1105 block3.header.nonce = 3;
1106 let new_chain = vec![block1, block2, block3];
1107
1108 let mut current_block1 = create_test_block();
1109 current_block1.header.nonce = 10;
1110 let mut current_block2 = create_test_block();
1111 current_block2.header.nonce = 11;
1112 let current_chain = vec![current_block1, current_block2];
1113 let utxo_set = UtxoSet::default();
1114
1115 let result = reorganize_chain(
1116 &new_chain,
1117 ¤t_chain,
1118 utxo_set,
1119 2,
1120 crate::types::Network::Regtest,
1121 );
1122 match result {
1123 Ok(reorg_result) => {
1124 assert_eq!(reorg_result.connected_blocks.len(), 3);
1125 assert_eq!(reorg_result.reorganization_depth, 2);
1126 assert_eq!(reorg_result.connected_block_undo_logs.len(), 3);
1128 }
1129 Err(_) => {
1130 }
1132 }
1133 }
1134
1135 #[test]
1136 fn test_undo_log_storage_and_retrieval() {
1137 use crate::block::connect_block;
1138 use crate::segwit::Witness;
1139
1140 let block = create_test_block_at_height(1);
1141 let mut utxo_set = UtxoSet::default();
1142
1143 let tx_id = calculate_tx_id(&block.transactions[0]);
1145 let outpoint = OutPoint {
1146 hash: tx_id,
1147 index: 0,
1148 };
1149 let utxo = UTXO {
1150 value: 5_000_000_000,
1151 script_pubkey: vec![0x51].into(),
1152 height: 1,
1153 is_coinbase: false,
1154 };
1155 utxo_set.insert(outpoint.clone(), std::sync::Arc::new(utxo.clone()));
1156
1157 let witnesses: Vec<Vec<Witness>> = block
1159 .transactions
1160 .iter()
1161 .map(|tx| tx.inputs.iter().map(|_| Vec::new()).collect())
1162 .collect();
1163 let ctx = crate::block::BlockValidationContext::for_network(crate::types::Network::Regtest);
1164 let (result, new_utxo_set, undo_log) =
1165 connect_block(&block, &witnesses, utxo_set.clone(), 1, &ctx).unwrap();
1166
1167 assert!(matches!(result, crate::types::ValidationResult::Valid));
1168
1169 assert!(
1171 !undo_log.entries.is_empty(),
1172 "Undo log should contain entries"
1173 );
1174
1175 let block_hash = calculate_block_hash(&block.header);
1177
1178 let mut undo_log_storage: HashMap<Hash, BlockUndoLog> = HashMap::new();
1180 undo_log_storage.insert(block_hash, undo_log.clone());
1181
1182 let retrieved_undo_log = undo_log_storage.get(&block_hash);
1184 assert!(
1185 retrieved_undo_log.is_some(),
1186 "Should be able to retrieve undo log"
1187 );
1188 assert_eq!(
1189 retrieved_undo_log.unwrap().entries.len(),
1190 undo_log.entries.len()
1191 );
1192
1193 let disconnected_utxo_set = disconnect_block(&block, &undo_log, new_utxo_set, 1).unwrap();
1195
1196 assert!(
1198 disconnected_utxo_set.contains_key(&outpoint),
1199 "Disconnected UTXO set should contain restored UTXO"
1200 );
1201 }
1202
1203 #[test]
1204 fn test_reorganize_with_undo_log_callback() {
1205 use crate::block::connect_block;
1206 use crate::segwit::Witness;
1207
1208 let block = create_test_block_at_height(1);
1210 let utxo_set = UtxoSet::default();
1211 let witnesses: Vec<Vec<Witness>> = block
1212 .transactions
1213 .iter()
1214 .map(|tx| tx.inputs.iter().map(|_| Vec::new()).collect())
1215 .collect();
1216
1217 let ctx = crate::block::BlockValidationContext::for_network(crate::types::Network::Regtest);
1218 let (result, connected_utxo_set, undo_log) =
1219 connect_block(&block, &witnesses, utxo_set.clone(), 1, &ctx).unwrap();
1220
1221 if !matches!(result, crate::types::ValidationResult::Valid) {
1222 eprintln!("Block validation failed: {:?}", result);
1223 }
1224 assert!(matches!(result, crate::types::ValidationResult::Valid));
1225
1226 let block_hash = calculate_block_hash(&block.header);
1228 let mut undo_log_storage: HashMap<Hash, BlockUndoLog> = HashMap::new();
1229 undo_log_storage.insert(block_hash, undo_log);
1230
1231 let get_undo_log =
1233 |hash: &Hash| -> Option<BlockUndoLog> { undo_log_storage.get(hash).cloned() };
1234
1235 let mut new_block = create_test_block_at_height(2);
1238 new_block.header.nonce = 42; let new_chain = vec![block.clone(), new_block];
1240 let current_chain = vec![block];
1241 let empty_witnesses: Vec<Vec<Vec<Witness>>> = new_chain
1242 .iter()
1243 .map(|b| {
1244 b.transactions
1245 .iter()
1246 .map(|tx| tx.inputs.iter().map(|_| Vec::new()).collect())
1247 .collect()
1248 })
1249 .collect();
1250
1251 let reorg_result = reorganize_chain_with_witnesses(
1252 &new_chain,
1253 &empty_witnesses,
1254 None,
1255 ¤t_chain,
1256 connected_utxo_set,
1257 1,
1258 None::<fn(&Block) -> Option<Vec<Witness>>>,
1259 None::<fn(Natural) -> Option<Vec<BlockHeader>>>,
1260 Some(get_undo_log),
1261 None::<fn(&Hash, &BlockUndoLog) -> Result<()>>, crate::types::Network::Regtest,
1263 );
1264
1265 match reorg_result {
1267 Ok(result) => {
1268 assert!(!result.connected_block_undo_logs.is_empty());
1270 }
1271 Err(_) => {
1272 }
1274 }
1275 }
1276
1277 #[test]
1278 fn test_reorganize_chain_empty_new_chain() {
1279 let new_chain = vec![];
1280 let current_chain = vec![create_test_block()];
1281 let utxo_set = UtxoSet::default();
1282
1283 let result = reorganize_chain(
1284 &new_chain,
1285 ¤t_chain,
1286 utxo_set,
1287 1,
1288 crate::types::Network::Regtest,
1289 );
1290 assert!(result.is_err());
1291 }
1292
1293 #[test]
1294 fn test_reorganize_chain_empty_current_chain() {
1295 let new_chain = vec![create_test_block()];
1296 let current_chain = vec![];
1297 let utxo_set = UtxoSet::default();
1298
1299 let result = reorganize_chain(
1300 &new_chain,
1301 ¤t_chain,
1302 utxo_set,
1303 0,
1304 crate::types::Network::Regtest,
1305 );
1306 assert!(result.is_err());
1307 }
1308
1309 #[test]
1310 fn test_disconnect_block() {
1311 let block = create_test_block();
1312 let mut utxo_set = UtxoSet::default();
1313
1314 let tx_id = calculate_tx_id(&block.transactions[0]);
1316 let outpoint = OutPoint {
1317 hash: tx_id,
1318 index: 0,
1319 };
1320 let utxo = UTXO {
1321 value: 50_000_000_000,
1322 script_pubkey: vec![0x51].into(),
1323 height: 1,
1324 is_coinbase: false,
1325 };
1326 utxo_set.insert(outpoint, std::sync::Arc::new(utxo));
1327
1328 let empty_undo_log = BlockUndoLog::new();
1330 let result = disconnect_block(&block, &empty_undo_log, utxo_set, 1);
1331 assert!(result.is_ok());
1332 }
1333
1334 #[test]
1335 fn test_calculate_chain_work_empty_chain() {
1336 let chain = vec![];
1337 let work = calculate_chain_work(&chain).unwrap();
1338 assert_eq!(work, 0);
1339 }
1340
1341 #[test]
1342 fn test_calculate_chain_work_multiple_blocks() {
1343 let mut chain = vec![create_test_block(), create_test_block()];
1344 chain[0].header.bits = 0x0300ffff;
1346 chain[1].header.bits = 0x0200ffff;
1347
1348 let work = calculate_chain_work(&chain).unwrap();
1349 assert!(work > 0);
1350 }
1351
1352 #[test]
1353 fn test_expand_target_edge_cases() {
1354 let result = expand_target(0x00000000);
1356 assert!(result.is_ok());
1357
1358 let result = expand_target(0x03ffffff);
1360 assert!(result.is_ok());
1361
1362 let result = expand_target(0x14000000); assert!(result.is_err());
1365 }
1366
1367 #[test]
1368 fn test_calculate_tx_id_different_transactions() {
1369 let tx1 = Transaction {
1370 version: 1,
1371 inputs: vec![].into(),
1372 outputs: vec![].into(),
1373 lock_time: 0,
1374 };
1375
1376 let tx2 = Transaction {
1377 version: 2,
1378 inputs: vec![].into(),
1379 outputs: vec![].into(),
1380 lock_time: 0,
1381 };
1382
1383 let id1 = calculate_tx_id(&tx1);
1384 let id2 = calculate_tx_id(&tx2);
1385
1386 assert_ne!(id1, id2);
1387 }
1388
1389 fn encode_bip34_height(height: u64) -> Vec<u8> {
1394 if height == 0 {
1395 return vec![0x00, 0xff]; }
1398 let mut height_bytes = Vec::new();
1399 let mut n = height;
1400 while n > 0 {
1401 height_bytes.push((n & 0xff) as u8);
1402 n >>= 8;
1403 }
1404 if height_bytes.last().map_or(false, |&b| b & 0x80 != 0) {
1406 height_bytes.push(0x00);
1407 }
1408 let mut script_sig = Vec::with_capacity(1 + height_bytes.len() + 1);
1409 script_sig.push(height_bytes.len() as u8); script_sig.extend_from_slice(&height_bytes);
1411 if script_sig.len() < 2 {
1413 script_sig.push(0xff);
1414 }
1415 script_sig
1416 }
1417
1418 fn create_test_block_at_height(height: u64) -> Block {
1420 use crate::mining::calculate_merkle_root;
1421
1422 let script_sig = encode_bip34_height(height);
1423 let coinbase_tx = Transaction {
1424 version: 1,
1425 inputs: vec![TransactionInput {
1426 prevout: OutPoint {
1427 hash: [0; 32].into(),
1428 index: 0xffffffff,
1429 },
1430 script_sig,
1431 sequence: 0xffffffff,
1432 }]
1433 .into(),
1434 outputs: vec![TransactionOutput {
1435 value: 5_000_000_000,
1436 script_pubkey: vec![0x51].into(),
1437 }]
1438 .into(),
1439 lock_time: 0,
1440 };
1441
1442 let merkle_root =
1443 calculate_merkle_root(&[coinbase_tx.clone()]).expect("Failed to calculate merkle root");
1444
1445 Block {
1446 header: BlockHeader {
1447 version: 4,
1448 prev_block_hash: [0; 32],
1449 merkle_root,
1450 timestamp: 1231006505,
1451 bits: 0x207fffff, nonce: 0,
1453 },
1454 transactions: vec![coinbase_tx].into_boxed_slice(),
1455 }
1456 }
1457
1458 fn create_test_block() -> Block {
1460 create_test_block_at_height(0)
1461 }
1462}