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)]
148 pub fn on_chain<H: Headers + ?Sized>(
149 &mut self,
150 headers: &H,
151 event: ChainEvent,
152 ) -> Vec<ChainUpdate> {
153 let targets: Vec<TxId> = match &event {
154 ChainEvent::Fork { height, .. } => self.at_or_above(*height),
155 ChainEvent::Reorg { fork_height, .. }
156 | ChainEvent::Split { fork_height, .. }
157 | ChainEvent::SplitResolved { fork_height, .. } => self.at_or_above(*fork_height),
158 ChainEvent::Invalidated { block_hash } => self
159 .headers
160 .get(block_hash)
161 .map(|ids| ids.iter().cloned().collect())
162 .unwrap_or_default(),
163 ChainEvent::Tip { .. } => self.suspect.iter().cloned().collect(),
164 ChainEvent::Frozen { .. } | ChainEvent::TipAge { .. } => vec![],
165 };
166 targets
167 .into_iter()
168 .map(|txid| {
169 let result =
170 self.apply_transition(&txid, |_, state| state.on_chain(headers, &event));
171 let reask = self.states[&txid].reask().cloned();
172 ChainUpdate {
173 txid,
174 result,
175 reask,
176 }
177 })
178 .collect()
179 }
180 fn at_or_above(&self, height: Height) -> Vec<TxId> {
181 self.mined
182 .range(height..)
183 .flat_map(|(_, ids)| ids.iter().cloned())
184 .collect()
185 }
186}
187
188fn remove_index<K: Ord>(index: &mut BTreeMap<K, BTreeSet<TxId>>, key: &K, txid: &str) {
189 if let Some(ids) = index.get_mut(key) {
190 ids.remove(txid);
191 if ids.is_empty() {
192 index.remove(key);
193 }
194 }
195}