zebra_chain/
history_tree.rs

1//! History tree (Merkle mountain range) structure that contains information about
2//! the block history as specified in ZIP-221.
3
4mod 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/// An error describing why a history tree operation failed.
25#[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        // Workaround since subtypes do not implement Eq.
40        // This is only used for tests anyway.
41        format!("{self:?}") == format!("{other:?}")
42    }
43}
44
45impl Eq for HistoryTreeError {}
46
47/// The inner [Tree] in one of its supported versions.
48#[derive(Debug)]
49enum InnerHistoryTree {
50    /// A pre-Orchard tree.
51    PreOrchard(Tree<PreOrchard>),
52    /// An Orchard-onward tree.
53    OrchardOnward(Tree<OrchardOnward>),
54}
55
56/// History tree (Merkle mountain range) structure that contains information about
57/// the block history, as specified in [ZIP-221](https://zips.z.cash/zip-0221).
58#[derive(Debug)]
59pub struct NonEmptyHistoryTree {
60    network: Network,
61    network_upgrade: NetworkUpgrade,
62    /// Merkle mountain range tree from `zcash_history`.
63    /// This is a "runtime" structure used to add / remove nodes, and it's not
64    /// persistent.
65    inner: InnerHistoryTree,
66    /// The number of nodes in the tree.
67    size: u32,
68    /// The peaks of the tree, indexed by their position in the array representation
69    /// of the tree. This can be persisted to save the tree.
70    peaks: SummaryDebug<BTreeMap<u32, Entry>>,
71    /// The height of the most recent block added to the tree.
72    current_height: Height,
73}
74
75impl NonEmptyHistoryTree {
76    /// Recreate a [`HistoryTree`] from previously saved data.
77    ///
78    /// The parameters must come from the values of [`NonEmptyHistoryTree::size`],
79    /// [`NonEmptyHistoryTree::peaks`] and [`NonEmptyHistoryTree::current_height`] of a HistoryTree.
80    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    /// Create a new history tree with a single block.
130    ///
131    /// `sapling_root` is the root of the Sapling note commitment tree of the block.
132    /// `orchard_root` is the root of the Orchard note commitment tree of the block;
133    ///  (ignored for pre-Orchard blocks).
134    #[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    /// Add block data to the tree.
188    ///
189    /// `sapling_root` is the root of the Sapling note commitment tree of the block.
190    /// `orchard_root` is the root of the Orchard note commitment tree of the block;
191    ///  (ignored for pre-Orchard blocks).
192    ///
193    /// # Panics
194    ///
195    /// If the block height is not one more than the previously pushed block.
196    #[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        // Check if the block has the expected height.
204        // librustzcash assumes the heights are correct and corrupts the tree if they are wrong,
205        // resulting in a confusing error, which we prevent here.
206        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            // This is the activation block of a network upgrade.
220            // Create a new tree.
221            let new_tree = Self::from_block(&self.network, block, sapling_root, orchard_root)?;
222            // Replaces self with the new tree
223            *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            // Not every entry is a peak; those will be trimmed later
238            self.peaks.insert(self.size, entry);
239            self.size += 1;
240        }
241        self.prune()?;
242        self.current_height = height;
243        Ok(())
244    }
245
246    /// Extend the history tree with the given blocks.
247    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    /// Prune tree, removing all non-peak entries.
261    fn prune(&mut self) -> Result<(), io::Error> {
262        // Go through all the peaks of the tree.
263        // This code is based on a librustzcash example:
264        // https://github.com/zcash/librustzcash/blob/02052526925fba9389f1428d6df254d4dec967e6/zcash_history/examples/long.rs
265        // The explanation of how it works is from zcashd:
266        // https://github.com/zcash/zcash/blob/0247c0c682d59184a717a6536edb0d18834be9a7/src/coins.cpp#L351
267
268        let mut peak_pos_set = HashSet::new();
269
270        // Assume the following example peak layout with 14 leaves, and 25 stored nodes in
271        // total (the "tree length"):
272        //
273        //             P
274        //            /\
275        //           /  \
276        //          / \  \
277        //        /    \  \  Altitude
278        //     _A_      \  \    3
279        //   _/   \_     B  \   2
280        //  / \   / \   / \  C  1
281        // /\ /\ /\ /\ /\ /\ /\ 0
282        //
283        // We start by determining the altitude of the highest peak (A).
284        let mut alt = (32 - (self.size + 1).leading_zeros() - 1) - 1;
285
286        // We determine the position of the highest peak (A) by pretending it is the right
287        // sibling in a tree, and its left-most leaf has position 0. Then the left sibling
288        // of (A) has position -1, and so we can "jump" to the peak's position by computing
289        // -1 + 2^(alt + 1) - 1.
290        let mut peak_pos = (1 << (alt + 1)) - 2;
291
292        // Now that we have the position and altitude of the highest peak (A), we collect
293        // the remaining peaks (B, C). We navigate the peaks as if they were nodes in this
294        // Merkle tree (with additional imaginary nodes 1 and 2, that have positions beyond
295        // the MMR's length):
296        //
297        //             / \
298        //            /   \
299        //           /     \
300        //         /         \
301        //       A ==========> 1
302        //      / \          //  \
303        //    _/   \_       B ==> 2
304        //   /\     /\     /\    //
305        //  /  \   /  \   /  \   C
306        // /\  /\ /\  /\ /\  /\ /\
307        //
308        loop {
309            // If peak_pos is out of bounds of the tree, we compute the position of its left
310            // child, and drop down one level in the tree.
311            if peak_pos >= self.size {
312                // left child, -2^alt
313                peak_pos -= 1 << alt;
314                alt -= 1;
315            }
316
317            // If the peak exists, we take it and then continue with its right sibling.
318            if peak_pos < self.size {
319                // There is a peak at index `peak_pos`
320                peak_pos_set.insert(peak_pos);
321
322                // right sibling
323                peak_pos = peak_pos + (1 << (alt + 1)) - 1;
324            }
325
326            if alt == 0 {
327                break;
328            }
329        }
330
331        // Remove all non-peak entries
332        self.peaks.retain(|k, _| peak_pos_set.contains(k));
333        // Rebuild tree
334        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    /// Return the hash of the tree root.
358    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    /// Return the peaks of the tree.
366    pub fn peaks(&self) -> &BTreeMap<u32, Entry> {
367        &self.peaks
368    }
369
370    /// Return the (total) number of nodes in the tree.
371    pub fn size(&self) -> u32 {
372        self.size
373    }
374
375    /// Return the height of the last added block.
376    pub fn current_height(&self) -> Height {
377        self.current_height
378    }
379
380    /// Return the network where this tree is used.
381    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/// A History Tree that keeps track of its own creation in the Heartwood
422/// activation block, being empty beforehand.
423#[derive(Debug, Default, Clone)]
424pub struct HistoryTree(Option<NonEmptyHistoryTree>);
425
426impl HistoryTree {
427    /// Create a HistoryTree from a block.
428    ///
429    /// If the block is pre-Heartwood, it returns an empty history tree.
430    #[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 early if there is no Heartwood activation height.
439            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    /// Push a block to a maybe-existing HistoryTree, handling network upgrades.
455    ///
456    /// The tree is updated in-place. It is created when pushing the Heartwood
457    /// activation block.
458    #[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                // Replace the current object with the new tree
494                *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    /// Return the hash of the tree root if the tree is not empty.
507    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 {}