ZIP-221: Add Orchard support to history tree (#2531)

* Add Orchard support to HistoryTree

* Handle network upgrades in HistoryTree

* Add additional methods to save/load HistoryTree

* Apply suggestions from code review

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

* Clarification of Entry documentation

* Improvements from code review

* Add HistoryTree tests

* Improved test comments and variable names based on feedback from #2458 on similar test

* Update zebra-chain/src/history_tree.rs

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

* Use type aliases for V1 and V2 history trees

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

zebra-chain/src/block.rs

@@ -14,7 +14,7 @@ pub mod arbitrary;
 #[cfg(any(test, feature = "bench"))]
 pub mod tests;
 
-use std::{collections::HashMap, fmt};
+use std::{collections::HashMap, convert::TryInto, fmt};
 
 pub use commitment::{ChainHistoryMmrRootHash, Commitment, CommitmentError};
 pub use hash::Hash;
@@ -146,6 +146,30 @@ impl Block {
             .flatten()
     }
 
+    /// Count how many Sapling transactions exist in a block,
+    /// i.e. transactions "where either of vSpendsSapling or vOutputsSapling is non-empty"
+    /// (https://zips.z.cash/zip-0221#tree-node-specification).
+    pub fn sapling_transactions_count(&self) -> u64 {
+        self.transactions
+            .iter()
+            .filter(|tx| tx.has_sapling_shielded_data())
+            .count()
+            .try_into()
+            .expect("number of transactions must fit u64")
+    }
+
+    /// Count how many Orchard transactions exist in a block,
+    /// i.e. transactions "where vActionsOrchard is non-empty."
+    /// (https://zips.z.cash/zip-0221#tree-node-specification).
+    pub fn orchard_transactions_count(&self) -> u64 {
+        self.transactions
+            .iter()
+            .filter(|tx| tx.has_orchard_shielded_data())
+            .count()
+            .try_into()
+            .expect("number of transactions must fit u64")
+    }
+
     /// Get all the value balances from this block by summing all the value balances
     /// in each transaction the block has.
     ///

zebra-chain/src/history_tree.rs

@@ -1,6 +1,8 @@
 //! History tree (Merkle mountain range) structure that contains information about
 //! the block history as specified in ZIP-221.
 
+mod tests;
+
 use std::{
     collections::{BTreeMap, HashSet},
     io,
@@ -13,7 +15,7 @@ use crate::{
     block::{Block, ChainHistoryMmrRootHash, Height},
     orchard,
     parameters::{Network, NetworkUpgrade},
-    primitives::zcash_history::{Entry, Tree as InnerHistoryTree},
+    primitives::zcash_history::{Entry, Tree, V1 as PreOrchard, V2 as OrchardOnward},
     sapling,
 };
 
@@ -30,6 +32,14 @@ pub enum HistoryTreeError {
     IOError(#[from] io::Error),
 }
 
+/// The inner [Tree] in one of its supported versions.
+enum InnerHistoryTree {
+    /// A pre-Orchard tree.
+    PreOrchard(Tree<PreOrchard>),
+    /// An Orchard-onward tree.
+    OrchardOnward(Tree<OrchardOnward>),
+}
+
 /// 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 {
@@ -49,19 +59,98 @@ pub struct HistoryTree {
 }
 
 impl HistoryTree {
+    /// Recreate a [`HistoryTree`] from previously saved data.
+    ///
+    /// The parameters must come from the values of [HistoryTree::size],
+    /// [HistoryTree::peaks] and [HistoryTree::current_height] of a HistoryTree.
+    pub fn from_cache(
+        network: Network,
+        size: u32,
+        peaks: BTreeMap<u32, Entry>,
+        current_height: Height,
+    ) -> Result<Self, io::Error> {
+        let network_upgrade = NetworkUpgrade::current(network, current_height);
+        let inner = match network_upgrade {
+            NetworkUpgrade::Genesis
+            | NetworkUpgrade::BeforeOverwinter
+            | NetworkUpgrade::Overwinter
+            | NetworkUpgrade::Sapling
+            | NetworkUpgrade::Blossom => {
+                panic!("HistoryTree does not exist for pre-Heartwood upgrades")
+            }
+            NetworkUpgrade::Heartwood | NetworkUpgrade::Canopy => {
+                let tree = Tree::<PreOrchard>::new_from_cache(
+                    network,
+                    network_upgrade,
+                    size,
+                    &peaks,
+                    &Default::default(),
+                )?;
+                InnerHistoryTree::PreOrchard(tree)
+            }
+            NetworkUpgrade::Nu5 => {
+                let tree = Tree::<OrchardOnward>::new_from_cache(
+                    network,
+                    network_upgrade,
+                    size,
+                    &peaks,
+                    &Default::default(),
+                )?;
+                InnerHistoryTree::OrchardOnward(tree)
+            }
+        };
+        Ok(Self {
+            network,
+            network_upgrade,
+            inner,
+            size,
+            peaks,
+            current_height,
+        })
+    }
+
     /// Create a new history tree with a single block.
+    ///
+    /// `sapling_root` is the root of the Sapling note commitment tree of the block.
+    /// `orchard_root` is the root of the Orchard note commitment tree of the block;
+    ///  (ignored for pre-Orchard blocks).
     pub fn from_block(
         network: Network,
         block: Arc<Block>,
         sapling_root: &sapling::tree::Root,
-        _orchard_root: Option<&orchard::tree::Root>,
+        orchard_root: &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 (tree, entry) = match network_upgrade {
+            NetworkUpgrade::Genesis
+            | NetworkUpgrade::BeforeOverwinter
+            | NetworkUpgrade::Overwinter
+            | NetworkUpgrade::Sapling
+            | NetworkUpgrade::Blossom => {
+                panic!("HistoryTree does not exist for pre-Heartwood upgrades")
+            }
+            NetworkUpgrade::Heartwood | NetworkUpgrade::Canopy => {
+                let (tree, entry) = Tree::<PreOrchard>::new_from_block(
+                    network,
+                    block,
+                    sapling_root,
+                    &Default::default(),
+                )?;
+                (InnerHistoryTree::PreOrchard(tree), entry)
+            }
+            NetworkUpgrade::Nu5 => {
+                let (tree, entry) = Tree::<OrchardOnward>::new_from_block(
+                    network,
+                    block,
+                    sapling_root,
+                    orchard_root,
+                )?;
+                (InnerHistoryTree::OrchardOnward(tree), entry)
+            }
+        };
         let mut peaks = BTreeMap::new();
         peaks.insert(0u32, entry);
         Ok(HistoryTree {
@@ -76,14 +165,18 @@ impl HistoryTree {
 
     /// Add block data to the tree.
     ///
+    /// `sapling_root` is the root of the Sapling note commitment tree of the block.
+    /// `orchard_root` is the root of the Orchard note commitment tree of the block;
+    ///  (ignored for pre-Orchard blocks).
+    ///
     /// # 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>,
+        orchard_root: &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,
@@ -97,33 +190,39 @@ impl HistoryTree {
                 height, self.current_height
             );
         }
+        let network_upgrade = NetworkUpgrade::current(self.network, height);
+        if network_upgrade != self.network_upgrade {
+            // This is the activation block of a network upgrade.
+            // Create a new tree.
+            let new_tree = Self::from_block(self.network, block, sapling_root, orchard_root)?;
+            // Replaces self with the new tree
+            *self = new_tree;
+            assert_eq!(self.network_upgrade, network_upgrade);
+            return Ok(());
+        }
 
-        // TODO: handle orchard root
-        let new_entries = self
-            .inner
-            .append_leaf(block, sapling_root)
-            .map_err(|e| HistoryTreeError::InnerError { inner: e })?;
+        let new_entries = match &mut self.inner {
+            InnerHistoryTree::PreOrchard(tree) => tree
+                .append_leaf(block, sapling_root, orchard_root)
+                .map_err(|e| HistoryTreeError::InnerError { inner: e })?,
+            InnerHistoryTree::OrchardOnward(tree) => tree
+                .append_leaf(block, sapling_root, orchard_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>,
-            ),
-        >,
+        T: IntoIterator<Item = (Arc<Block>, &'a sapling::tree::Root, &'a orchard::tree::Root)>,
     >(
         &mut self,
         iter: T,
@@ -208,32 +307,77 @@ impl HistoryTree {
         // 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(),
-        )?;
+        self.inner = match self.inner {
+            InnerHistoryTree::PreOrchard(_) => {
+                InnerHistoryTree::PreOrchard(Tree::<PreOrchard>::new_from_cache(
+                    self.network,
+                    self.network_upgrade,
+                    self.size,
+                    &self.peaks,
+                    &Default::default(),
+                )?)
+            }
+            InnerHistoryTree::OrchardOnward(_) => {
+                InnerHistoryTree::OrchardOnward(Tree::<OrchardOnward>::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()
+        match &self.inner {
+            InnerHistoryTree::PreOrchard(tree) => tree.hash(),
+            InnerHistoryTree::OrchardOnward(tree) => tree.hash(),
+        }
+    }
+
+    /// Return the peaks of the tree.
+    pub fn peaks(&self) -> &BTreeMap<u32, Entry> {
+        &self.peaks
+    }
+
+    /// Return the (total) number of nodes in the tree.
+    pub fn size(&self) -> u32 {
+        self.size
+    }
+
+    /// Return the height of the last added block.
+    pub fn current_height(&self) -> Height {
+        self.current_height
     }
 }
 
 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");
+        let tree = match self.inner {
+            InnerHistoryTree::PreOrchard(_) => InnerHistoryTree::PreOrchard(
+                Tree::<PreOrchard>::new_from_cache(
+                    self.network,
+                    self.network_upgrade,
+                    self.size,
+                    &self.peaks,
+                    &Default::default(),
+                )
+                .expect("rebuilding an existing tree should always work"),
+            ),
+            InnerHistoryTree::OrchardOnward(_) => InnerHistoryTree::OrchardOnward(
+                Tree::<OrchardOnward>::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,

zebra-chain/src/history_tree/tests.rs

@@ -0,0 +1,4 @@
+//! Tests for history trees
+
+#[cfg(test)]
+mod vectors;