use alloc::collections::BTreeMap;
use alloc::string::String;
use alloc::vec::Vec;
use crate::hash::Hash;
use crate::store::{SealStore, StoreError};
#[derive(Clone, Debug)]
pub struct ReorgEvent {
pub chain: String,
pub fork_height: u64,
pub depth: u64,
pub old_tip: [u8; 32],
pub new_tip: [u8; 32],
}
pub struct ReorgMonitor {
tips: BTreeMap<String, (u64, [u8; 32])>,
reorgs: Vec<ReorgEvent>,
}
impl ReorgMonitor {
pub fn new() -> Self {
Self {
tips: BTreeMap::new(),
reorgs: Vec::new(),
}
}
pub fn update_tip(
&mut self,
chain: &str,
height: u64,
tip_hash: [u8; 32],
) -> Option<ReorgEvent> {
if let Some(&(prev_height, prev_hash)) = self.tips.get(chain) {
if height <= prev_height {
let depth = prev_height.saturating_sub(height) + 1;
let event = ReorgEvent {
chain: chain.to_string(),
fork_height: height,
depth,
old_tip: prev_hash,
new_tip: tip_hash,
};
self.tips.insert(chain.to_string(), (height, tip_hash));
self.reorgs.push(event.clone());
return Some(event);
}
if height > prev_height + 1 {
}
}
self.tips.insert(chain.to_string(), (height, tip_hash));
None
}
pub fn handle_reorg(
&self,
event: &ReorgEvent,
store: &mut dyn SealStore,
) -> Result<usize, StoreError> {
let mut removed = 0;
removed += store.remove_anchors_after(&event.chain, event.fork_height)?;
removed += store.remove_seals_after(&event.chain, event.fork_height)?;
Ok(removed)
}
pub fn recent_reorgs(&self, chain: &str) -> Vec<&ReorgEvent> {
self.reorgs.iter().filter(|r| r.chain == chain).collect()
}
pub fn tip(&self, chain: &str) -> Option<(u64, [u8; 32])> {
self.tips.get(chain).copied()
}
}
impl Default for ReorgMonitor {
fn default() -> Self {
Self::new()
}
}
pub struct PublicationTracker {
pending: BTreeMap<String, Vec<PendingPublication>>,
pub timeout_seconds: u64,
}
#[derive(Clone, Debug)]
pub struct PendingPublication {
pub tx_hash: Vec<u8>,
pub commitment_hash: Hash,
pub submitted_at: u64,
}
impl PublicationTracker {
pub fn new(timeout_seconds: u64) -> Self {
Self {
pending: BTreeMap::new(),
timeout_seconds,
}
}
pub fn track_publication(
&mut self,
chain: &str,
tx_hash: Vec<u8>,
commitment_hash: Hash,
submitted_at: u64,
) {
self.pending
.entry(chain.to_string())
.or_default()
.push(PendingPublication {
tx_hash,
commitment_hash,
submitted_at,
});
}
pub fn confirm_publication(&mut self, chain: &str, tx_hash: &[u8]) -> bool {
if let Some(pending) = self.pending.get_mut(chain) {
let before = pending.len();
pending.retain(|p| p.tx_hash != tx_hash);
return pending.len() < before;
}
false
}
pub fn timed_out(&self, chain: &str, current_time: u64) -> Vec<&PendingPublication> {
self.pending
.get(chain)
.map(|p| {
p.iter()
.filter(|pp| {
current_time.saturating_sub(pp.submitted_at) > self.timeout_seconds
})
.collect()
})
.unwrap_or_default()
}
pub fn pending_count(&self, chain: &str) -> usize {
self.pending.get(chain).map(|p| p.len()).unwrap_or(0)
}
pub fn clear_chain(&mut self, chain: &str) {
self.pending.remove(chain);
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::store::InMemorySealStore;
#[test]
fn test_reorg_monitor_normal_update() {
let mut monitor = ReorgMonitor::new();
assert!(monitor.update_tip("bitcoin", 100, [1u8; 32]).is_none());
assert!(monitor.update_tip("bitcoin", 101, [2u8; 32]).is_none());
}
#[test]
fn test_reorg_monitor_detect_reorg() {
let mut monitor = ReorgMonitor::new();
monitor.update_tip("bitcoin", 100, [1u8; 32]);
monitor.update_tip("bitcoin", 101, [2u8; 32]);
let event = monitor.update_tip("bitcoin", 99, [3u8; 32]);
assert!(event.is_some());
let event = event.unwrap();
assert_eq!(event.chain, "bitcoin");
assert_eq!(event.fork_height, 99);
assert_eq!(event.old_tip, [2u8; 32]);
assert_eq!(event.new_tip, [3u8; 32]);
}
#[test]
fn test_reorg_monitor_handle_reorg() {
let mut monitor = ReorgMonitor::new();
monitor.update_tip("bitcoin", 100, [1u8; 32]);
let mut store = InMemorySealStore::new();
store
.save_seal(&crate::store::SealRecord {
chain: "bitcoin".to_string(),
seal_id: vec![1],
consumed_at_height: 101,
commitment_hash: Hash::new([0xAA; 32]),
recorded_at: 1700000000,
})
.unwrap();
let event = ReorgEvent {
chain: "bitcoin".to_string(),
fork_height: 100,
depth: 1,
old_tip: [1u8; 32],
new_tip: [2u8; 32],
};
let removed = monitor.handle_reorg(&event, &mut store).unwrap();
assert_eq!(removed, 1);
}
#[test]
fn test_publication_tracker_lifecycle() {
let mut tracker = PublicationTracker::new(3600);
tracker.track_publication("bitcoin", vec![1, 2, 3], Hash::new([0xAA; 32]), 1700000000);
assert_eq!(tracker.pending_count("bitcoin"), 1);
assert!(tracker.confirm_publication("bitcoin", &[1, 2, 3]));
assert_eq!(tracker.pending_count("bitcoin"), 0);
}
#[test]
fn test_publication_tracker_timeout() {
let mut tracker = PublicationTracker::new(3600);
tracker.track_publication("bitcoin", vec![1, 2, 3], Hash::new([0xAA; 32]), 1700000000);
assert!(tracker.timed_out("bitcoin", 1700003000).is_empty());
let timed_out = tracker.timed_out("bitcoin", 1700004000);
assert_eq!(timed_out.len(), 1);
}
#[test]
fn test_publication_tracker_multiple() {
let mut tracker = PublicationTracker::new(3600);
tracker.track_publication("bitcoin", vec![1], Hash::new([1u8; 32]), 1700000000);
tracker.track_publication("bitcoin", vec![2], Hash::new([2u8; 32]), 1700000000);
tracker.track_publication("ethereum", vec![3], Hash::new([3u8; 32]), 1700000000);
assert_eq!(tracker.pending_count("bitcoin"), 2);
assert_eq!(tracker.pending_count("ethereum"), 1);
tracker.confirm_publication("bitcoin", &[1]);
assert_eq!(tracker.pending_count("bitcoin"), 1);
tracker.clear_chain("bitcoin");
assert_eq!(tracker.pending_count("bitcoin"), 0);
}
}