Add HistoryTree struct (#2396)

* Add HistoryTree struct

* Update zebra-chain/src/primitives/zcash_history.rs

Co-authored-by: Deirdre Connolly <deirdre@zfnd.org>

* fix formatting

* Apply suggestions from code review

Co-authored-by: Deirdre Connolly <deirdre@zfnd.org>

* history_tree.rs: fixes from code review

Co-authored-by: Deirdre Connolly <deirdre@zfnd.org>

zebra-chain/src/history_tree.rs

@@ -0,0 +1,246 @@
+//! History tree (Merkle mountain range) structure that contains information about
+//! the block history as specified in ZIP-221.
+
+use std::{
+    collections::{BTreeMap, HashSet},
+    io,
+    sync::Arc,
+};
+
+use thiserror::Error;
+
+use crate::{
+    block::{Block, ChainHistoryMmrRootHash, Height},
+    orchard,
+    parameters::{Network, NetworkUpgrade},
+    primitives::zcash_history::{Entry, Tree as InnerHistoryTree},
+    sapling,
+};
+
+/// An error describing why a history tree operation failed.
+#[derive(Debug, Error)]
+#[non_exhaustive]
+#[allow(missing_docs)]
+pub enum HistoryTreeError {
+    #[error("zcash_history error: {inner:?}")]
+    #[non_exhaustive]
+    InnerError { inner: zcash_history::Error },
+
+    #[error("I/O error")]
+    IOError(#[from] io::Error),
+}
+
+/// History tree (Merkle mountain range) structure that contains information about
+// the block history, as specified in [ZIP-221][https://zips.z.cash/zip-0221].
+pub struct HistoryTree {
+    network: Network,
+    network_upgrade: NetworkUpgrade,
+    /// Merkle mountain range tree from `zcash_history`.
+    /// This is a "runtime" structure used to add / remove nodes, and it's not
+    /// persistent.
+    inner: InnerHistoryTree,
+    /// The number of nodes in the tree.
+    size: u32,
+    /// The peaks of the tree, indexed by their position in the array representation
+    /// of the tree. This can be persisted to save the tree.
+    peaks: BTreeMap<u32, Entry>,
+    /// The height of the most recent block added to the tree.
+    current_height: Height,
+}
+
+impl HistoryTree {
+    /// Create a new history tree with a single block.
+    pub fn from_block(
+        network: Network,
+        block: Arc<Block>,
+        sapling_root: &sapling::tree::Root,
+        _orchard_root: Option<&orchard::tree::Root>,
+    ) -> Result<Self, io::Error> {
+        let height = block
+            .coinbase_height()
+            .expect("block must have coinbase height during contextual verification");
+        let network_upgrade = NetworkUpgrade::current(network, height);
+        // TODO: handle Orchard root, see https://github.com/ZcashFoundation/zebra/issues/2283
+        let (tree, entry) = InnerHistoryTree::new_from_block(network, block, sapling_root)?;
+        let mut peaks = BTreeMap::new();
+        peaks.insert(0u32, entry);
+        Ok(HistoryTree {
+            network,
+            network_upgrade,
+            inner: tree,
+            size: 1,
+            peaks,
+            current_height: height,
+        })
+    }
+
+    /// Add block data to the tree.
+    ///
+    /// # Panics
+    ///
+    /// If the block height is not one more than the previously pushed block.
+    pub fn push(
+        &mut self,
+        block: Arc<Block>,
+        sapling_root: &sapling::tree::Root,
+        _orchard_root: Option<&orchard::tree::Root>,
+    ) -> Result<(), HistoryTreeError> {
+        // Check if the block has the expected height.
+        // librustzcash assumes the heights are correct and corrupts the tree if they are wrong,
+        // resulting in a confusing error, which we prevent here.
+        let height = block
+            .coinbase_height()
+            .expect("block must have coinbase height during contextual verification");
+        if height - self.current_height != 1 {
+            panic!(
+                "added block with height {:?} but it must be {:?}+1",
+                height, self.current_height
+            );
+        }
+
+        // TODO: handle orchard root
+        let new_entries = self
+            .inner
+            .append_leaf(block, sapling_root)
+            .map_err(|e| HistoryTreeError::InnerError { inner: e })?;
+        for entry in new_entries {
+            // Not every entry is a peak; those will be trimmed later
+            self.peaks.insert(self.size, entry);
+            self.size += 1;
+        }
+        self.prune()?;
+        // TODO: implement network upgrade logic: drop previous history, start new history
+        self.current_height = height;
+        Ok(())
+    }
+
+    /// Extend the history tree with the given blocks.
+    pub fn try_extend<
+        'a,
+        T: IntoIterator<
+            Item = (
+                Arc<Block>,
+                &'a sapling::tree::Root,
+                Option<&'a orchard::tree::Root>,
+            ),
+        >,
+    >(
+        &mut self,
+        iter: T,
+    ) -> Result<(), HistoryTreeError> {
+        for (block, sapling_root, orchard_root) in iter {
+            self.push(block, sapling_root, orchard_root)?;
+        }
+        Ok(())
+    }
+
+    /// Prune tree, removing all non-peak entries.
+    fn prune(&mut self) -> Result<(), io::Error> {
+        // Go through all the peaks of the tree.
+        // This code is based on a librustzcash example:
+        // https://github.com/zcash/librustzcash/blob/02052526925fba9389f1428d6df254d4dec967e6/zcash_history/examples/long.rs
+        // The explanation of how it works is from zcashd:
+        // https://github.com/zcash/zcash/blob/0247c0c682d59184a717a6536edb0d18834be9a7/src/coins.cpp#L351
+
+        let mut peak_pos_set = HashSet::new();
+
+        // Assume the following example peak layout with 14 leaves, and 25 stored nodes in
+        // total (the "tree length"):
+        //
+        //             P
+        //            /\
+        //           /  \
+        //          / \  \
+        //        /    \  \  Altitude
+        //     _A_      \  \    3
+        //   _/   \_     B  \   2
+        //  / \   / \   / \  C  1
+        // /\ /\ /\ /\ /\ /\ /\ 0
+        //
+        // We start by determining the altitude of the highest peak (A).
+        let mut alt = (32 - ((self.size + 1) as u32).leading_zeros() - 1) - 1;
+
+        // We determine the position of the highest peak (A) by pretending it is the right
+        // sibling in a tree, and its left-most leaf has position 0. Then the left sibling
+        // of (A) has position -1, and so we can "jump" to the peak's position by computing
+        // -1 + 2^(alt + 1) - 1.
+        let mut peak_pos = (1 << (alt + 1)) - 2;
+
+        // Now that we have the position and altitude of the highest peak (A), we collect
+        // the remaining peaks (B, C). We navigate the peaks as if they were nodes in this
+        // Merkle tree (with additional imaginary nodes 1 and 2, that have positions beyond
+        // the MMR's length):
+        //
+        //             / \
+        //            /   \
+        //           /     \
+        //         /         \
+        //       A ==========> 1
+        //      / \          //  \
+        //    _/   \_       B ==> 2
+        //   /\     /\     /\    //
+        //  /  \   /  \   /  \   C
+        // /\  /\ /\  /\ /\  /\ /\
+        //
+        loop {
+            // If peak_pos is out of bounds of the tree, we compute the position of its left
+            // child, and drop down one level in the tree.
+            if peak_pos >= self.size {
+                // left child, -2^alt
+                peak_pos -= 1 << alt;
+                alt -= 1;
+            }
+
+            // If the peak exists, we take it and then continue with its right sibling.
+            if peak_pos < self.size {
+                // There is a peak at index `peak_pos`
+                peak_pos_set.insert(peak_pos);
+
+                // right sibling
+                peak_pos = peak_pos + (1 << (alt + 1)) - 1;
+            }
+
+            if alt == 0 {
+                break;
+            }
+        }
+
+        // Remove all non-peak entries
+        self.peaks.retain(|k, _| peak_pos_set.contains(k));
+        // Rebuild tree
+        self.inner = InnerHistoryTree::new_from_cache(
+            self.network,
+            self.network_upgrade,
+            self.size,
+            &self.peaks,
+            &Default::default(),
+        )?;
+        Ok(())
+    }
+
+    /// Return the hash of the tree root.
+    pub fn hash(&self) -> ChainHistoryMmrRootHash {
+        self.inner.hash()
+    }
+}
+
+impl Clone for HistoryTree {
+    fn clone(&self) -> Self {
+        let tree = InnerHistoryTree::new_from_cache(
+            self.network,
+            self.network_upgrade,
+            self.size,
+            &self.peaks,
+            &Default::default(),
+        )
+        .expect("rebuilding an existing tree should always work");
+        HistoryTree {
+            network: self.network,
+            network_upgrade: self.network_upgrade,
+            inner: tree,
+            size: self.size,
+            peaks: self.peaks.clone(),
+            current_height: self.current_height,
+        }
+    }
+}

zebra-chain/src/lib.rs

@@ -23,6 +23,7 @@ extern crate bitflags;
 pub mod amount;
 pub mod block;
 pub mod fmt;
+pub mod history_tree;
 pub mod orchard;
 pub mod parameters;
 pub mod primitives;

zebra-chain/src/primitives/zcash_history.rs

@@ -6,7 +6,7 @@
 
 mod tests;
 
-use std::{collections::HashMap, convert::TryInto, io, sync::Arc};
+use std::{collections::BTreeMap, convert::TryInto, io, sync::Arc};
 
 use crate::{
     block::{Block, ChainHistoryMmrRootHash},
@@ -43,7 +43,9 @@ impl From<&zcash_history::NodeData> for NodeData {
 }
 
 /// An encoded entry in the tree.
+///
 /// Contains the node data and information about its position in the tree.
+#[derive(Clone)]
 pub struct Entry {
     inner: [u8; zcash_history::MAX_ENTRY_SIZE],
 }
@@ -101,12 +103,12 @@ impl Tree {
     /// # Panics
     ///
     /// Will panic if `peaks` is empty.
-    fn new_from_cache(
+    pub fn new_from_cache(
         network: Network,
         network_upgrade: NetworkUpgrade,
         length: u32,
-        peaks: &HashMap<u32, &Entry>,
+        peaks: &BTreeMap<u32, Entry>,
-        extra: &HashMap<u32, &Entry>,
+        extra: &BTreeMap<u32, Entry>,
     ) -> Result<Self, io::Error> {
         let branch_id = network_upgrade
             .branch_id()
@@ -132,19 +134,24 @@ impl Tree {
     /// Create a single-node MMR tree for the given block.
     ///
     /// `sapling_root` is the root of the Sapling note commitment tree of the block.
-    fn new_from_block(
+    pub fn new_from_block(
         network: Network,
         block: Arc<Block>,
         sapling_root: &sapling::tree::Root,
-    ) -> Result<Self, io::Error> {
+    ) -> Result<(Self, Entry), io::Error> {
         let height = block
             .coinbase_height()
             .expect("block must have coinbase height during contextual verification");
         let network_upgrade = NetworkUpgrade::current(network, height);
         let entry0 = Entry::new_leaf(block, network, sapling_root);
-        let mut peaks = HashMap::new();
+        let mut peaks = BTreeMap::new();
-        peaks.insert(0u32, &entry0);
+        peaks.insert(0u32, entry0);
-        Tree::new_from_cache(network, network_upgrade, 1, &peaks, &HashMap::new())
+        Ok((
+            Tree::new_from_cache(network, network_upgrade, 1, &peaks, &BTreeMap::new())?,
+            peaks
+                .remove(&0u32)
+                .expect("must work since it was just added"),
+        ))
     }
 
     /// Append a new block to the tree, as a new leaf.
@@ -157,18 +164,18 @@ impl Tree {
     ///
     /// Panics if the network upgrade of the given block is different from
     /// the network upgrade of the other blocks in the tree.
-    fn append_leaf(
+    pub fn append_leaf(
         &mut self,
         block: Arc<Block>,
         sapling_root: &sapling::tree::Root,
-    ) -> Result<Vec<NodeData>, zcash_history::Error> {
+    ) -> Result<Vec<Entry>, zcash_history::Error> {
         let height = block
             .coinbase_height()
             .expect("block must have coinbase height during contextual verification");
         let network_upgrade = NetworkUpgrade::current(self.network, height);
         if self.network_upgrade != network_upgrade {
             panic!(
-                "added block from network upgrade {:?} but MMR tree is restricted to {:?}",
+                "added block from network upgrade {:?} but history tree is restricted to {:?}",
                 network_upgrade, self.network_upgrade
             );
         }
@@ -177,17 +184,17 @@ impl Tree {
         let appended = self.inner.append_leaf(node_data)?;
 
         let mut new_nodes = Vec::new();
-        for entry in appended {
+        for entry_link in appended {
-            let mut node = NodeData {
+            let mut entry = Entry {
-                inner: [0; zcash_history::MAX_NODE_DATA_SIZE],
+                inner: [0; zcash_history::MAX_ENTRY_SIZE],
             };
             self.inner
-                .resolve_link(entry)
+                .resolve_link(entry_link)
                 .expect("entry was just generated so it must be valid")
-                .data()
+                .node()
-                .write(&mut &mut node.inner[..])
+                .write(&mut &mut entry.inner[..])
                 .expect("buffer was created with enough capacity");
-            new_nodes.push(node);
+            new_nodes.push(entry);
         }
         Ok(new_nodes)
     }
@@ -196,7 +203,7 @@ impl Tree {
     fn append_leaf_iter(
         &mut self,
         vals: impl Iterator<Item = (Arc<Block>, sapling::tree::Root)>,
-    ) -> Result<Vec<NodeData>, zcash_history::Error> {
+    ) -> Result<Vec<Entry>, zcash_history::Error> {
         let mut new_nodes = Vec::new();
         for (block, root) in vals {
             new_nodes.append(&mut self.append_leaf(block, &root)?);
@@ -212,7 +219,7 @@ impl Tree {
     }
 
     /// Return the root hash of the tree, i.e. `hashChainHistoryRoot`.
-    fn hash(&self) -> ChainHistoryMmrRootHash {
+    pub fn hash(&self) -> ChainHistoryMmrRootHash {
         // Both append_leaf() and truncate_leaf() leave a root node, so it should
         // always exist.
         self.inner

zebra-chain/src/primitives/zcash_history/tests/vectors.rs

@@ -49,7 +49,7 @@ fn tree_for_network_upgrade(network: Network, network_upgrade: NetworkUpgrade) -
     // Build initial MMR tree with only Block 0
     let sapling_root0 =
         sapling::tree::Root(**sapling_roots.get(&height).expect("test vector exists"));
-    let mut tree = Tree::new_from_block(network, block0, &sapling_root0)?;
+    let (mut tree, _) = Tree::new_from_block(network, block0, &sapling_root0)?;
 
     // Compute root hash of the MMR tree, which will be included in the next block
     let hash0 = tree.hash();