1mod tests;
5
6use std::{
7 collections::{BTreeMap, HashSet},
8 io,
9 ops::Deref,
10 sync::Arc,
11};
12
13use thiserror::Error;
14
15use crate::{
16 block::{Block, ChainHistoryMmrRootHash, Height},
17 fmt::SummaryDebug,
18 orchard,
19 parameters::{Network, NetworkUpgrade},
20 primitives::zcash_history::{Entry, Tree, V1 as PreOrchard, V2 as OrchardOnward},
21 sapling,
22};
23
24#[derive(Debug, Error)]
26#[non_exhaustive]
27#[allow(missing_docs)]
28pub enum HistoryTreeError {
29 #[error("zcash_history error: {inner:?}")]
30 #[non_exhaustive]
31 InnerError { inner: zcash_history::Error },
32
33 #[error("I/O error: {0}")]
34 IOError(#[from] io::Error),
35}
36
37impl PartialEq for HistoryTreeError {
38 fn eq(&self, other: &Self) -> bool {
39 format!("{self:?}") == format!("{other:?}")
42 }
43}
44
45impl Eq for HistoryTreeError {}
46
47#[derive(Debug)]
49enum InnerHistoryTree {
50 PreOrchard(Tree<PreOrchard>),
52 OrchardOnward(Tree<OrchardOnward>),
54}
55
56#[derive(Debug)]
59pub struct NonEmptyHistoryTree {
60 network: Network,
61 network_upgrade: NetworkUpgrade,
62 inner: InnerHistoryTree,
66 size: u32,
68 peaks: SummaryDebug<BTreeMap<u32, Entry>>,
71 current_height: Height,
73}
74
75impl NonEmptyHistoryTree {
76 pub fn from_cache(
81 network: &Network,
82 size: u32,
83 peaks: BTreeMap<u32, Entry>,
84 current_height: Height,
85 ) -> Result<Self, HistoryTreeError> {
86 let network_upgrade = NetworkUpgrade::current(network, current_height);
87 let inner = match network_upgrade {
88 NetworkUpgrade::Genesis
89 | NetworkUpgrade::BeforeOverwinter
90 | NetworkUpgrade::Overwinter
91 | NetworkUpgrade::Sapling
92 | NetworkUpgrade::Blossom => {
93 panic!("HistoryTree does not exist for pre-Heartwood upgrades")
94 }
95 NetworkUpgrade::Heartwood | NetworkUpgrade::Canopy => {
96 let tree = Tree::<PreOrchard>::new_from_cache(
97 network,
98 network_upgrade,
99 size,
100 &peaks,
101 &Default::default(),
102 )?;
103 InnerHistoryTree::PreOrchard(tree)
104 }
105 NetworkUpgrade::Nu5
106 | NetworkUpgrade::Nu6
107 | NetworkUpgrade::Nu6_1
108 | NetworkUpgrade::Nu7 => {
109 let tree = Tree::<OrchardOnward>::new_from_cache(
110 network,
111 network_upgrade,
112 size,
113 &peaks,
114 &Default::default(),
115 )?;
116 InnerHistoryTree::OrchardOnward(tree)
117 }
118 };
119 Ok(Self {
120 network: network.clone(),
121 network_upgrade,
122 inner,
123 size,
124 peaks: peaks.into(),
125 current_height,
126 })
127 }
128
129 #[allow(clippy::unwrap_in_result)]
135 pub fn from_block(
136 network: &Network,
137 block: Arc<Block>,
138 sapling_root: &sapling::tree::Root,
139 orchard_root: &orchard::tree::Root,
140 ) -> Result<Self, HistoryTreeError> {
141 let height = block
142 .coinbase_height()
143 .expect("block must have coinbase height during contextual verification");
144 let network_upgrade = NetworkUpgrade::current(network, height);
145 let (tree, entry) = match network_upgrade {
146 NetworkUpgrade::Genesis
147 | NetworkUpgrade::BeforeOverwinter
148 | NetworkUpgrade::Overwinter
149 | NetworkUpgrade::Sapling
150 | NetworkUpgrade::Blossom => {
151 panic!("HistoryTree does not exist for pre-Heartwood upgrades")
152 }
153 NetworkUpgrade::Heartwood | NetworkUpgrade::Canopy => {
154 let (tree, entry) = Tree::<PreOrchard>::new_from_block(
155 network,
156 block,
157 sapling_root,
158 &Default::default(),
159 )?;
160 (InnerHistoryTree::PreOrchard(tree), entry)
161 }
162 NetworkUpgrade::Nu5
163 | NetworkUpgrade::Nu6
164 | NetworkUpgrade::Nu6_1
165 | NetworkUpgrade::Nu7 => {
166 let (tree, entry) = Tree::<OrchardOnward>::new_from_block(
167 network,
168 block,
169 sapling_root,
170 orchard_root,
171 )?;
172 (InnerHistoryTree::OrchardOnward(tree), entry)
173 }
174 };
175 let mut peaks = BTreeMap::new();
176 peaks.insert(0u32, entry);
177 Ok(NonEmptyHistoryTree {
178 network: network.clone(),
179 network_upgrade,
180 inner: tree,
181 size: 1,
182 peaks: peaks.into(),
183 current_height: height,
184 })
185 }
186
187 #[allow(clippy::unwrap_in_result)]
197 pub fn push(
198 &mut self,
199 block: Arc<Block>,
200 sapling_root: &sapling::tree::Root,
201 orchard_root: &orchard::tree::Root,
202 ) -> Result<(), HistoryTreeError> {
203 let height = block
207 .coinbase_height()
208 .expect("block must have coinbase height during contextual verification");
209
210 assert!(
211 Some(height) == self.current_height + 1,
212 "added block with height {:?} but it must be {:?}+1",
213 height,
214 self.current_height
215 );
216
217 let network_upgrade = NetworkUpgrade::current(&self.network, height);
218 if network_upgrade != self.network_upgrade {
219 let new_tree = Self::from_block(&self.network, block, sapling_root, orchard_root)?;
222 *self = new_tree;
224 assert_eq!(self.network_upgrade, network_upgrade);
225 return Ok(());
226 }
227
228 let new_entries = match &mut self.inner {
229 InnerHistoryTree::PreOrchard(tree) => tree
230 .append_leaf(block, sapling_root, orchard_root)
231 .map_err(|e| HistoryTreeError::InnerError { inner: e })?,
232 InnerHistoryTree::OrchardOnward(tree) => tree
233 .append_leaf(block, sapling_root, orchard_root)
234 .map_err(|e| HistoryTreeError::InnerError { inner: e })?,
235 };
236 for entry in new_entries {
237 self.peaks.insert(self.size, entry);
239 self.size += 1;
240 }
241 self.prune()?;
242 self.current_height = height;
243 Ok(())
244 }
245
246 pub fn try_extend<
248 'a,
249 T: IntoIterator<Item = (Arc<Block>, &'a sapling::tree::Root, &'a orchard::tree::Root)>,
250 >(
251 &mut self,
252 iter: T,
253 ) -> Result<(), HistoryTreeError> {
254 for (block, sapling_root, orchard_root) in iter {
255 self.push(block, sapling_root, orchard_root)?;
256 }
257 Ok(())
258 }
259
260 fn prune(&mut self) -> Result<(), io::Error> {
262 let mut peak_pos_set = HashSet::new();
269
270 let mut alt = (32 - (self.size + 1).leading_zeros() - 1) - 1;
285
286 let mut peak_pos = (1 << (alt + 1)) - 2;
291
292 loop {
309 if peak_pos >= self.size {
312 peak_pos -= 1 << alt;
314 alt -= 1;
315 }
316
317 if peak_pos < self.size {
319 peak_pos_set.insert(peak_pos);
321
322 peak_pos = peak_pos + (1 << (alt + 1)) - 1;
324 }
325
326 if alt == 0 {
327 break;
328 }
329 }
330
331 self.peaks.retain(|k, _| peak_pos_set.contains(k));
333 self.inner = match self.inner {
335 InnerHistoryTree::PreOrchard(_) => {
336 InnerHistoryTree::PreOrchard(Tree::<PreOrchard>::new_from_cache(
337 &self.network,
338 self.network_upgrade,
339 self.size,
340 &self.peaks,
341 &Default::default(),
342 )?)
343 }
344 InnerHistoryTree::OrchardOnward(_) => {
345 InnerHistoryTree::OrchardOnward(Tree::<OrchardOnward>::new_from_cache(
346 &self.network,
347 self.network_upgrade,
348 self.size,
349 &self.peaks,
350 &Default::default(),
351 )?)
352 }
353 };
354 Ok(())
355 }
356
357 pub fn hash(&self) -> ChainHistoryMmrRootHash {
359 match &self.inner {
360 InnerHistoryTree::PreOrchard(tree) => tree.hash(),
361 InnerHistoryTree::OrchardOnward(tree) => tree.hash(),
362 }
363 }
364
365 pub fn peaks(&self) -> &BTreeMap<u32, Entry> {
367 &self.peaks
368 }
369
370 pub fn size(&self) -> u32 {
372 self.size
373 }
374
375 pub fn current_height(&self) -> Height {
377 self.current_height
378 }
379
380 pub fn network(&self) -> &Network {
382 &self.network
383 }
384}
385
386impl Clone for NonEmptyHistoryTree {
387 fn clone(&self) -> Self {
388 let tree = match self.inner {
389 InnerHistoryTree::PreOrchard(_) => InnerHistoryTree::PreOrchard(
390 Tree::<PreOrchard>::new_from_cache(
391 &self.network,
392 self.network_upgrade,
393 self.size,
394 &self.peaks,
395 &Default::default(),
396 )
397 .expect("rebuilding an existing tree should always work"),
398 ),
399 InnerHistoryTree::OrchardOnward(_) => InnerHistoryTree::OrchardOnward(
400 Tree::<OrchardOnward>::new_from_cache(
401 &self.network,
402 self.network_upgrade,
403 self.size,
404 &self.peaks,
405 &Default::default(),
406 )
407 .expect("rebuilding an existing tree should always work"),
408 ),
409 };
410 NonEmptyHistoryTree {
411 network: self.network.clone(),
412 network_upgrade: self.network_upgrade,
413 inner: tree,
414 size: self.size,
415 peaks: self.peaks.clone(),
416 current_height: self.current_height,
417 }
418 }
419}
420
421#[derive(Debug, Default, Clone)]
424pub struct HistoryTree(Option<NonEmptyHistoryTree>);
425
426impl HistoryTree {
427 #[allow(clippy::unwrap_in_result)]
431 pub fn from_block(
432 network: &Network,
433 block: Arc<Block>,
434 sapling_root: &sapling::tree::Root,
435 orchard_root: &orchard::tree::Root,
436 ) -> Result<Self, HistoryTreeError> {
437 let Some(heartwood_height) = NetworkUpgrade::Heartwood.activation_height(network) else {
438 return Ok(HistoryTree(None));
440 };
441
442 match block
443 .coinbase_height()
444 .expect("must have height")
445 .cmp(&heartwood_height)
446 {
447 std::cmp::Ordering::Less => Ok(HistoryTree(None)),
448 _ => Ok(
449 NonEmptyHistoryTree::from_block(network, block, sapling_root, orchard_root)?.into(),
450 ),
451 }
452 }
453
454 #[allow(clippy::unwrap_in_result)]
459 pub fn push(
460 &mut self,
461 network: &Network,
462 block: Arc<Block>,
463 sapling_root: &sapling::tree::Root,
464 orchard_root: &orchard::tree::Root,
465 ) -> Result<(), HistoryTreeError> {
466 let Some(heartwood_height) = NetworkUpgrade::Heartwood.activation_height(network) else {
467 assert!(
468 self.0.is_none(),
469 "history tree must not exist pre-Heartwood"
470 );
471
472 return Ok(());
473 };
474
475 match block
476 .coinbase_height()
477 .expect("must have height")
478 .cmp(&heartwood_height)
479 {
480 std::cmp::Ordering::Less => {
481 assert!(
482 self.0.is_none(),
483 "history tree must not exist pre-Heartwood"
484 );
485 }
486 std::cmp::Ordering::Equal => {
487 let tree = Some(NonEmptyHistoryTree::from_block(
488 network,
489 block,
490 sapling_root,
491 orchard_root,
492 )?);
493 *self = HistoryTree(tree);
495 }
496 std::cmp::Ordering::Greater => {
497 self.0
498 .as_mut()
499 .expect("history tree must exist Heartwood-onward")
500 .push(block.clone(), sapling_root, orchard_root)?;
501 }
502 };
503 Ok(())
504 }
505
506 pub fn hash(&self) -> Option<ChainHistoryMmrRootHash> {
508 Some(self.0.as_ref()?.hash())
509 }
510}
511
512impl From<NonEmptyHistoryTree> for HistoryTree {
513 fn from(tree: NonEmptyHistoryTree) -> Self {
514 HistoryTree(Some(tree))
515 }
516}
517
518impl From<Option<NonEmptyHistoryTree>> for HistoryTree {
519 fn from(tree: Option<NonEmptyHistoryTree>) -> Self {
520 HistoryTree(tree)
521 }
522}
523
524impl Deref for HistoryTree {
525 type Target = Option<NonEmptyHistoryTree>;
526 fn deref(&self) -> &Self::Target {
527 &self.0
528 }
529}
530
531impl PartialEq for HistoryTree {
532 fn eq(&self, other: &Self) -> bool {
533 self.as_ref().map(|tree| tree.hash()) == other.as_ref().map(|other_tree| other_tree.hash())
534 }
535}
536
537impl Eq for HistoryTree {}