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)]
22pub struct Tracker {
23 params: Params,
24 states: BTreeMap<TxId, State>,
25 mined: BTreeMap<Height, BTreeSet<TxId>>,
26 headers: BTreeMap<Hash, BTreeSet<TxId>>,
27 suspect: BTreeSet<TxId>,
28}
29
30#[derive(Debug, PartialEq, Eq)]
31struct Index {
32 height: Option<Height>,
33 header: Option<Hash>,
34 suspect: bool,
35}
36impl Index {
37 fn of(state: &State) -> Self {
38 Self {
39 height: state.word().height(),
40 header: match state.word() {
41 Word::Mined(m) => Some(m.checked().header().hash.clone()),
42 _ => None,
43 },
44 suspect: state.suspect(),
45 }
46 }
47}
48
49impl Tracker {
50 #[must_use]
52 pub fn new(params: Params) -> Self {
53 Self {
54 params,
55 states: BTreeMap::new(),
56 mined: BTreeMap::new(),
57 headers: BTreeMap::new(),
58 suspect: BTreeSet::new(),
59 }
60 }
61 pub fn track(&mut self, txid: impl Into<TxId>) -> &State {
63 let txid = txid.into().to_ascii_lowercase();
64 self.states
65 .entry(txid.clone())
66 .or_insert_with(|| State::new(txid))
67 }
68 pub fn restore(&mut self, state: State) -> Option<State> {
72 let txid = state.txid().to_owned();
73 let tracked = self.states.contains_key(&txid);
74 let mut previous = None;
75 let _ = self.apply_transition(&txid, |_, slot| {
77 previous = Some(std::mem::replace(slot, state));
78 Ok(())
79 });
80 previous.filter(|_| tracked)
81 }
82 #[must_use]
84 pub fn get(&self, txid: &str) -> Option<&State> {
85 self.states.get(&txid.to_ascii_lowercase())
86 }
87 #[must_use]
89 pub fn len(&self) -> usize {
90 self.states.len()
91 }
92 #[must_use]
94 pub fn is_empty(&self) -> bool {
95 self.states.is_empty()
96 }
97 pub fn apply<H: Headers + ?Sized>(
104 &mut self,
105 txid: &str,
106 headers: &H,
107 input: Input,
108 ) -> Result<(), CheckError> {
109 self.apply_transition(txid, |params, state| state.step(params, headers, input))
110 }
111 fn apply_transition(
112 &mut self,
113 txid: &str,
114 transition: impl FnOnce(&Params, &mut State) -> Result<(), CheckError>,
115 ) -> Result<(), CheckError> {
116 let txid = txid.to_ascii_lowercase();
117 let state = self
118 .states
119 .entry(txid.clone())
120 .or_insert_with(|| State::new(&txid));
121 let old = Index::of(state);
122 let result = transition(&self.params, state);
123 let new = Index::of(state);
124 if old != new {
125 if let Some(h) = old.height {
126 remove_index(&mut self.mined, &h, &txid);
127 }
128 if let Some(hash) = old.header {
129 remove_index(&mut self.headers, &hash, &txid);
130 }
131 self.suspect.remove(&txid);
132 if let Some(h) = new.height {
133 self.mined.entry(h).or_default().insert(txid.clone());
134 }
135 if let Some(hash) = new.header {
136 self.headers.entry(hash).or_default().insert(txid.clone());
137 }
138 if new.suspect {
139 self.suspect.insert(txid);
140 }
141 }
142 result
143 }
144 #[allow(clippy::needless_pass_by_value)]
155 pub fn on_chain<H: Headers + ?Sized>(
156 &mut self,
157 headers: &H,
158 event: ChainEvent,
159 ) -> Vec<ChainUpdate> {
160 let targets: Vec<TxId> = match &event {
161 ChainEvent::Fork { height, .. } => self.at_or_above(*height),
162 ChainEvent::Reorg { fork_height, .. }
163 | ChainEvent::Split { fork_height, .. }
164 | ChainEvent::SplitResolved { fork_height, .. } => self.at_or_above(*fork_height),
165 ChainEvent::Invalidated { block_hash } => self
166 .headers
167 .get(block_hash)
168 .map(|ids| ids.iter().cloned().collect())
169 .unwrap_or_default(),
170 ChainEvent::Tip { .. } => self.suspect.iter().cloned().collect(),
171 ChainEvent::Frozen { .. } | ChainEvent::TipAge { .. } => vec![],
172 };
173 targets
174 .into_iter()
175 .map(|txid| {
176 let result =
177 self.apply_transition(&txid, |_, state| state.on_chain(headers, &event));
178 let reask = self.states[&txid].reask().cloned();
179 ChainUpdate {
180 txid,
181 result,
182 reask,
183 }
184 })
185 .collect()
186 }
187 fn at_or_above(&self, height: Height) -> Vec<TxId> {
188 self.mined
189 .range(height..)
190 .flat_map(|(_, ids)| ids.iter().cloned())
191 .collect()
192 }
193}
194
195fn remove_index<K: Ord>(index: &mut BTreeMap<K, BTreeSet<TxId>>, key: &K, txid: &str) {
196 if let Some(ids) = index.get_mut(key) {
197 ids.remove(txid);
198 if ids.is_empty() {
199 index.remove(key);
200 }
201 }
202}