1use crate::{
2 ChainEvent, CheckError, Hash, Headers, Height, Input, Params, Reask, State, TxId, Word,
3};
4use std::collections::{BTreeMap, BTreeSet};
5
6#[derive(Clone, Debug, PartialEq, Eq)]
8pub struct ChainUpdate {
9 pub txid: TxId,
11 pub result: Result<(), CheckError>,
13 pub reask: Option<Reask>,
15}
16
17#[derive(Clone, Debug)]
21pub struct Tracker {
22 params: Params,
23 states: BTreeMap<TxId, State>,
24 mined: BTreeMap<Height, BTreeSet<TxId>>,
25 headers: BTreeMap<Hash, BTreeSet<TxId>>,
26 suspect: BTreeSet<TxId>,
27}
28
29#[derive(Debug, PartialEq, Eq)]
30struct Index {
31 height: Option<Height>,
32 header: Option<Hash>,
33 suspect: bool,
34}
35impl Index {
36 fn of(state: &State) -> Self {
37 Self {
38 height: state.word().height(),
39 header: match state.word() {
40 Word::Mined(m) => Some(m.checked().header().hash.clone()),
41 _ => None,
42 },
43 suspect: state.suspect(),
44 }
45 }
46}
47
48impl Tracker {
49 pub fn new(params: Params) -> Self {
51 Self {
52 params,
53 states: BTreeMap::new(),
54 mined: BTreeMap::new(),
55 headers: BTreeMap::new(),
56 suspect: BTreeSet::new(),
57 }
58 }
59 pub fn track(&mut self, txid: impl Into<TxId>) -> &State {
61 let txid = txid.into().to_ascii_lowercase();
62 self.states
63 .entry(txid.clone())
64 .or_insert_with(|| State::new(txid))
65 }
66 pub fn get(&self, txid: &str) -> Option<&State> {
68 self.states.get(&txid.to_ascii_lowercase())
69 }
70 pub fn len(&self) -> usize {
72 self.states.len()
73 }
74 pub fn is_empty(&self) -> bool {
76 self.states.is_empty()
77 }
78 pub fn apply<H: Headers + ?Sized>(
81 &mut self,
82 txid: &str,
83 headers: &H,
84 input: Input,
85 ) -> Result<(), CheckError> {
86 self.apply_transition(txid, |params, state| state.step(params, headers, input))
87 }
88 fn apply_transition(
89 &mut self,
90 txid: &str,
91 transition: impl FnOnce(&Params, &mut State) -> Result<(), CheckError>,
92 ) -> Result<(), CheckError> {
93 let txid = txid.to_ascii_lowercase();
94 let state = self
95 .states
96 .entry(txid.clone())
97 .or_insert_with(|| State::new(&txid));
98 let old = Index::of(state);
99 let result = transition(&self.params, state);
100 let new = Index::of(state);
101 if old != new {
102 if let Some(h) = old.height {
103 remove_index(&mut self.mined, &h, &txid);
104 }
105 if let Some(hash) = old.header {
106 remove_index(&mut self.headers, &hash, &txid);
107 }
108 self.suspect.remove(&txid);
109 if let Some(h) = new.height {
110 self.mined.entry(h).or_default().insert(txid.clone());
111 }
112 if let Some(hash) = new.header {
113 self.headers.entry(hash).or_default().insert(txid.clone());
114 }
115 if new.suspect {
116 self.suspect.insert(txid);
117 }
118 }
119 result
120 }
121 pub fn on_chain<H: Headers + ?Sized>(
124 &mut self,
125 headers: &H,
126 event: ChainEvent,
127 ) -> Vec<ChainUpdate> {
128 let targets: Vec<TxId> = match &event {
129 ChainEvent::Fork { height, .. } => self.at_or_above(*height),
130 ChainEvent::Reorg { fork_height, .. } => self.at_or_above(*fork_height),
131 ChainEvent::Invalidated { block_hash } => self
132 .headers
133 .get(block_hash)
134 .map(|ids| ids.iter().cloned().collect())
135 .unwrap_or_default(),
136 ChainEvent::Tip { .. } => self.suspect.iter().cloned().collect(),
137 ChainEvent::Frozen { .. } | ChainEvent::TipAge { .. } => vec![],
138 };
139 targets
140 .into_iter()
141 .map(|txid| {
142 let result =
143 self.apply_transition(&txid, |_, state| state.on_chain(headers, &event));
144 let reask = self.states[&txid].reask().cloned();
145 ChainUpdate {
146 txid,
147 result,
148 reask,
149 }
150 })
151 .collect()
152 }
153 fn at_or_above(&self, height: Height) -> Vec<TxId> {
154 self.mined
155 .range(height..)
156 .flat_map(|(_, ids)| ids.iter().cloned())
157 .collect()
158 }
159}
160
161fn remove_index<K: Ord>(index: &mut BTreeMap<K, BTreeSet<TxId>>, key: &K, txid: &str) {
162 if let Some(ids) = index.get_mut(key) {
163 ids.remove(txid);
164 if ids.is_empty() {
165 index.remove(key);
166 }
167 }
168}