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,6 +165,10 @@ 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.
@@ -83,7 +176,7 @@ impl HistoryTree {
         &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,19 +190,31 @@ 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(())
     }
@@ -117,13 +222,7 @@ impl HistoryTree {
     /// 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;

zebra-chain/src/history_tree/tests/vectors.rs

@@ -0,0 +1,179 @@
+use std::sync::Arc;
+
+use crate::{
+    block::{
+        Block,
+        Commitment::{self, ChainHistoryActivationReserved},
+    },
+    history_tree::HistoryTree,
+    parameters::{Network, NetworkUpgrade},
+    sapling,
+    serialization::ZcashDeserializeInto,
+};
+
+use color_eyre::eyre;
+use eyre::Result;
+use zebra_test::vectors::{
+    MAINNET_BLOCKS, MAINNET_FINAL_SAPLING_ROOTS, TESTNET_BLOCKS, TESTNET_FINAL_SAPLING_ROOTS,
+};
+
+/// Test the history tree using the activation block of a network upgrade
+/// and its next block.
+///
+/// This test is very similar to the zcash_history test in
+/// zebra-chain/src/primitives/zcash_history/tests/vectors.rs, but with the
+/// higher level API.
+#[test]
+fn push_and_prune() -> Result<()> {
+    push_and_prune_for_network_upgrade(Network::Mainnet, NetworkUpgrade::Heartwood)?;
+    push_and_prune_for_network_upgrade(Network::Testnet, NetworkUpgrade::Heartwood)?;
+    push_and_prune_for_network_upgrade(Network::Mainnet, NetworkUpgrade::Canopy)?;
+    push_and_prune_for_network_upgrade(Network::Testnet, NetworkUpgrade::Canopy)?;
+    Ok(())
+}
+
+fn push_and_prune_for_network_upgrade(
+    network: Network,
+    network_upgrade: NetworkUpgrade,
+) -> Result<()> {
+    let (blocks, sapling_roots) = match network {
+        Network::Mainnet => (&*MAINNET_BLOCKS, &*MAINNET_FINAL_SAPLING_ROOTS),
+        Network::Testnet => (&*TESTNET_BLOCKS, &*TESTNET_FINAL_SAPLING_ROOTS),
+    };
+    let height = network_upgrade.activation_height(network).unwrap().0;
+
+    // Load first block (activation block of the given network upgrade)
+    let first_block = Arc::new(
+        blocks
+            .get(&height)
+            .expect("test vector exists")
+            .zcash_deserialize_into::<Block>()
+            .expect("block is structurally valid"),
+    );
+
+    // Check its commitment
+    let first_commitment = first_block.commitment(network)?;
+    if network_upgrade == NetworkUpgrade::Heartwood {
+        // Heartwood is the only upgrade that has a reserved value.
+        // (For other upgrades we could compare with the expected commitment,
+        // but we haven't calculated them.)
+        assert_eq!(first_commitment, ChainHistoryActivationReserved);
+    }
+
+    // Build initial history tree tree with only the first block
+    let first_sapling_root =
+        sapling::tree::Root(**sapling_roots.get(&height).expect("test vector exists"));
+    let mut tree = HistoryTree::from_block(
+        network,
+        first_block,
+        &first_sapling_root,
+        &Default::default(),
+    )?;
+
+    assert_eq!(tree.size(), 1);
+    assert_eq!(tree.peaks().len(), 1);
+    assert_eq!(tree.current_height().0, height);
+
+    // Compute root hash of the history tree, which will be included in the next block
+    let first_root = tree.hash();
+
+    // Load second block (activation + 1)
+    let second_block = Arc::new(
+        blocks
+            .get(&(height + 1))
+            .expect("test vector exists")
+            .zcash_deserialize_into::<Block>()
+            .expect("block is structurally valid"),
+    );
+
+    // Check its commitment
+    let second_commitment = second_block.commitment(network)?;
+    assert_eq!(second_commitment, Commitment::ChainHistoryRoot(first_root));
+
+    // Append second block to history tree
+    let second_sapling_root = sapling::tree::Root(
+        **sapling_roots
+            .get(&(height + 1))
+            .expect("test vector exists"),
+    );
+    tree.push(second_block, &second_sapling_root, &Default::default())
+        .unwrap();
+
+    // Adding a second block will produce a 3-node tree (one parent and two leafs).
+    assert_eq!(tree.size(), 3);
+    // The tree must have been pruned, resulting in a single peak (the parent).
+    assert_eq!(tree.peaks().len(), 1);
+    assert_eq!(tree.current_height().0, height + 1);
+
+    Ok(())
+}
+
+/// Test the history tree works during a network upgrade using the block
+/// of a network upgrade and the previous block from the previous upgrade.
+#[test]
+fn upgrade() -> Result<()> {
+    // The history tree only exists Hearwood-onward, and the only upgrade for which
+    // we have vectors since then is Canopy. Therefore, only test the Heartwood->Canopy upgrade.
+    upgrade_for_network_upgrade(Network::Mainnet, NetworkUpgrade::Canopy)?;
+    upgrade_for_network_upgrade(Network::Testnet, NetworkUpgrade::Canopy)?;
+    Ok(())
+}
+
+fn upgrade_for_network_upgrade(network: Network, network_upgrade: NetworkUpgrade) -> Result<()> {
+    let (blocks, sapling_roots) = match network {
+        Network::Mainnet => (&*MAINNET_BLOCKS, &*MAINNET_FINAL_SAPLING_ROOTS),
+        Network::Testnet => (&*TESTNET_BLOCKS, &*TESTNET_FINAL_SAPLING_ROOTS),
+    };
+    let height = network_upgrade.activation_height(network).unwrap().0;
+
+    // Load previous block (the block before the activation block of the given network upgrade)
+    let block_prev = Arc::new(
+        blocks
+            .get(&(height - 1))
+            .expect("test vector exists")
+            .zcash_deserialize_into::<Block>()
+            .expect("block is structurally valid"),
+    );
+
+    // Build a history tree with only the previous block (activation height - 1)
+    // This tree will not match the actual tree (which has all the blocks since the previous
+    // network upgrade), so we won't be able to check if its root is correct.
+    let sapling_root_prev =
+        sapling::tree::Root(**sapling_roots.get(&height).expect("test vector exists"));
+    let mut tree =
+        HistoryTree::from_block(network, block_prev, &sapling_root_prev, &Default::default())?;
+
+    assert_eq!(tree.size(), 1);
+    assert_eq!(tree.peaks().len(), 1);
+    assert_eq!(tree.current_height().0, height - 1);
+
+    // Load block of the activation height
+    let activation_block = Arc::new(
+        blocks
+            .get(&height)
+            .expect("test vector exists")
+            .zcash_deserialize_into::<Block>()
+            .expect("block is structurally valid"),
+    );
+
+    // Append block to history tree. This must trigger a upgrade of the tree,
+    // which should be recreated.
+    let activation_sapling_root = sapling::tree::Root(
+        **sapling_roots
+            .get(&(height + 1))
+            .expect("test vector exists"),
+    );
+    tree.push(
+        activation_block,
+        &activation_sapling_root,
+        &Default::default(),
+    )
+    .unwrap();
+
+    // Check if the tree has a single node, i.e. it has been recreated.
+    assert_eq!(tree.size(), 1);
+    assert_eq!(tree.peaks().len(), 1);
+    assert_eq!(tree.current_height().0, height);
+
+    Ok(())
+}

zebra-chain/src/primitives/zcash_history.rs

@@ -8,20 +8,34 @@ mod tests;
 
 use std::{collections::BTreeMap, convert::TryInto, io, sync::Arc};
 
+pub use zcash_history::{V1, V2};
+
 use crate::{
     block::{Block, ChainHistoryMmrRootHash},
-    parameters::{ConsensusBranchId, Network, NetworkUpgrade},
+    orchard,
+    parameters::{Network, NetworkUpgrade},
     sapling,
 };
 
+/// A trait to represent a version of `Tree`.
+pub trait Version: zcash_history::Version {
+    /// Convert a Block into the NodeData for this version.
+    fn block_to_history_node(
+        block: Arc<Block>,
+        network: Network,
+        sapling_root: &sapling::tree::Root,
+        orchard_root: &orchard::tree::Root,
+    ) -> Self::NodeData;
+}
+
 /// A MMR Tree using zcash_history::Tree.
 ///
 /// Currently it should not be used as a long-term data structure because it
 /// may grow without limits.
-pub struct Tree {
+pub struct Tree<V: zcash_history::Version> {
     network: Network,
     network_upgrade: NetworkUpgrade,
-    inner: zcash_history::Tree<zcash_history::V1>,
+    inner: zcash_history::Tree<V>,
 }
 
 /// An encoded tree node data.
@@ -50,9 +64,21 @@ pub struct Entry {
     inner: [u8; zcash_history::MAX_ENTRY_SIZE],
 }
 
-impl From<zcash_history::Entry<zcash_history::V1>> for Entry {
-    /// Convert from librustzcash.
-    fn from(inner_entry: zcash_history::Entry<zcash_history::V1>) -> Self {
+impl Entry {
+    /// Create a leaf Entry for the given block, its network, and the root of its
+    /// note commitment trees.
+    ///
+    /// `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 V1 trees).
+    fn new_leaf<V: Version>(
+        block: Arc<Block>,
+        network: Network,
+        sapling_root: &sapling::tree::Root,
+        orchard_root: &orchard::tree::Root,
+    ) -> Self {
+        let node_data = V::block_to_history_node(block, network, sapling_root, orchard_root);
+        let inner_entry = zcash_history::Entry::<V>::new_leaf(node_data);
         let mut entry = Entry {
             inner: [0; zcash_history::MAX_ENTRY_SIZE],
         };
@@ -63,34 +89,7 @@ impl From<zcash_history::Entry<zcash_history::V1>> for Entry {
     }
 }
 
-impl Entry {
-    /// Create a leaf Entry for the given block, its network, and the root of its
-    /// Sapling note commitment tree.
-    fn new_leaf(block: Arc<Block>, network: Network, sapling_root: &sapling::tree::Root) -> Self {
-        let node_data = block_to_history_node(block, network, sapling_root);
-        let inner_entry = zcash_history::Entry::<zcash_history::V1>::new_leaf(node_data);
-        inner_entry.into()
-    }
-
-    /// Create a node (non-leaf) Entry from the encoded node data and the indices of
-    /// its children (in the array representation of the MMR tree).
-    fn new_node(
-        branch_id: ConsensusBranchId,
-        data: NodeData,
-        left_idx: u32,
-        right_idx: u32,
-    ) -> Result<Self, io::Error> {
-        let node_data = zcash_history::NodeData::from_bytes(branch_id.into(), data.inner)?;
-        let inner_entry = zcash_history::Entry::new(
-            node_data,
-            zcash_history::EntryLink::Stored(left_idx),
-            zcash_history::EntryLink::Stored(right_idx),
-        );
-        Ok(inner_entry.into())
-    }
-}
-
-impl Tree {
+impl<V: Version> Tree<V> {
     /// Create a MMR tree with the given length from the given cache of nodes.
     ///
     /// The `peaks` are the peaks of the MMR tree to build and their position in the
@@ -134,16 +133,19 @@ 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.
+    /// `orchard_root` is the root of the Orchard note commitment tree of the block;
+    ///  (ignored for V1 trees).
     pub fn new_from_block(
         network: Network,
         block: Arc<Block>,
         sapling_root: &sapling::tree::Root,
+        orchard_root: &orchard::tree::Root,
     ) -> 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 entry0 = Entry::new_leaf::<V>(block, network, sapling_root, orchard_root);
         let mut peaks = BTreeMap::new();
         peaks.insert(0u32, entry0);
         Ok((
@@ -157,6 +159,8 @@ impl Tree {
     /// Append a new block to the tree, as a new leaf.
     ///
     /// `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 V1 trees).
     ///
     /// Returns a vector of nodes added to the tree (leaf + internal nodes).
     ///
@@ -168,6 +172,7 @@ impl Tree {
         &mut self,
         block: Arc<Block>,
         sapling_root: &sapling::tree::Root,
+        orchard_root: &orchard::tree::Root,
     ) -> Result<Vec<Entry>, zcash_history::Error> {
         let height = block
             .coinbase_height()
@@ -180,7 +185,7 @@ impl Tree {
             );
         }
 
-        let node_data = block_to_history_node(block, self.network, sapling_root);
+        let node_data = V::block_to_history_node(block, self.network, sapling_root, orchard_root);
         let appended = self.inner.append_leaf(node_data)?;
 
         let mut new_nodes = Vec::new();
@@ -202,11 +207,11 @@ impl Tree {
     /// Append multiple blocks to the tree.
     fn append_leaf_iter(
         &mut self,
-        vals: impl Iterator<Item = (Arc<Block>, sapling::tree::Root)>,
+        vals: impl Iterator<Item = (Arc<Block>, sapling::tree::Root, orchard::tree::Root)>,
     ) -> 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)?);
+        for (block, sapling_root, orchard_root) in vals {
+            new_nodes.append(&mut self.append_leaf(block, &sapling_root, &orchard_root)?);
         }
         Ok(new_nodes)
     }
@@ -222,72 +227,96 @@ impl Tree {
     pub fn hash(&self) -> ChainHistoryMmrRootHash {
         // Both append_leaf() and truncate_leaf() leave a root node, so it should
         // always exist.
-        self.inner
-            .root_node()
-            .expect("must have root node")
-            .data()
-            .hash()
-            .into()
+        V::hash(self.inner.root_node().expect("must have root node").data()).into()
     }
 }
 
-/// Convert a Block into a zcash_history::NodeData used in the MMR tree.
-///
-/// `sapling_root` is the root of the Sapling note commitment tree of the block.
-fn block_to_history_node(
-    block: Arc<Block>,
-    network: Network,
-    sapling_root: &sapling::tree::Root,
-) -> zcash_history::NodeData {
-    let height = block
-        .coinbase_height()
-        .expect("block must have coinbase height during contextual verification");
-    let branch_id = ConsensusBranchId::current(network, height)
-        .expect("must have branch ID for chain history network upgrades");
-    let block_hash = block.hash().0;
-    let time: u32 = block
-        .header
-        .time
-        .timestamp()
-        .try_into()
-        .expect("deserialized and generated timestamps are u32 values");
-    let target = block.header.difficulty_threshold.0;
-    let sapling_root: [u8; 32] = sapling_root.into();
-    let work = block
-        .header
-        .difficulty_threshold
-        .to_work()
-        .expect("work must be valid during contextual verification");
-    // There is no direct `std::primitive::u128` to `bigint::U256` conversion
-    let work = bigint::U256::from_big_endian(&work.as_u128().to_be_bytes());
+impl Version for zcash_history::V1 {
+    /// Convert a Block into a V1::NodeData used in the MMR tree.
+    ///
+    /// `sapling_root` is the root of the Sapling note commitment tree of the block.
+    /// `orchard_root` is ignored.
+    fn block_to_history_node(
+        block: Arc<Block>,
+        network: Network,
+        sapling_root: &sapling::tree::Root,
+        _orchard_root: &orchard::tree::Root,
+    ) -> Self::NodeData {
+        let height = block
+            .coinbase_height()
+            .expect("block must have coinbase height during contextual verification");
+        let network_upgrade = NetworkUpgrade::current(network, height);
+        let branch_id = network_upgrade
+            .branch_id()
+            .expect("must have branch ID for chain history network upgrades");
+        let block_hash = block.hash().0;
+        let time: u32 = block
+            .header
+            .time
+            .timestamp()
+            .try_into()
+            .expect("deserialized and generated timestamps are u32 values");
+        let target = block.header.difficulty_threshold.0;
+        let sapling_root: [u8; 32] = sapling_root.into();
+        let work = block
+            .header
+            .difficulty_threshold
+            .to_work()
+            .expect("work must be valid during contextual verification");
+        // There is no direct `std::primitive::u128` to `bigint::U256` conversion
+        let work = bigint::U256::from_big_endian(&work.as_u128().to_be_bytes());
 
-    let sapling_tx_count = count_sapling_transactions(block);
+        let sapling_tx_count = block.sapling_transactions_count();
 
-    zcash_history::NodeData {
-        consensus_branch_id: branch_id.into(),
-        subtree_commitment: block_hash,
-        start_time: time,
-        end_time: time,
-        start_target: target,
-        end_target: target,
-        start_sapling_root: sapling_root,
-        end_sapling_root: sapling_root,
-        subtree_total_work: work,
-        start_height: height.0 as u64,
-        end_height: height.0 as u64,
-        sapling_tx: sapling_tx_count,
+        match network_upgrade {
+            NetworkUpgrade::Genesis
+            | NetworkUpgrade::BeforeOverwinter
+            | NetworkUpgrade::Overwinter
+            | NetworkUpgrade::Sapling
+            | NetworkUpgrade::Blossom => {
+                panic!("HistoryTree does not exist for pre-Heartwood upgrades")
+            }
+            // Nu5 is included because this function is called by the V2 implementation
+            // since the V1::NodeData is included inside the V2::NodeData.
+            NetworkUpgrade::Heartwood | NetworkUpgrade::Canopy | NetworkUpgrade::Nu5 => {
+                zcash_history::NodeData {
+                    consensus_branch_id: branch_id.into(),
+                    subtree_commitment: block_hash,
+                    start_time: time,
+                    end_time: time,
+                    start_target: target,
+                    end_target: target,
+                    start_sapling_root: sapling_root,
+                    end_sapling_root: sapling_root,
+                    subtree_total_work: work,
+                    start_height: height.0 as u64,
+                    end_height: height.0 as u64,
+                    sapling_tx: sapling_tx_count,
+                }
+            }
+        }
     }
 }
 
-/// 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).
-fn count_sapling_transactions(block: Arc<Block>) -> u64 {
-    block
-        .transactions
-        .iter()
-        .filter(|tx| tx.has_sapling_shielded_data())
-        .count()
-        .try_into()
-        .expect("number of transactions must fit u64")
+impl Version for V2 {
+    /// Convert a Block into a V1::NodeData used in the MMR 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.
+    fn block_to_history_node(
+        block: Arc<Block>,
+        network: Network,
+        sapling_root: &sapling::tree::Root,
+        orchard_root: &orchard::tree::Root,
+    ) -> Self::NodeData {
+        let orchard_tx_count = block.orchard_transactions_count();
+        let node_data_v1 = V1::block_to_history_node(block, network, sapling_root, orchard_root);
+        let orchard_root: [u8; 32] = orchard_root.into();
+        Self::NodeData {
+            v1: node_data_v1,
+            start_orchard_root: orchard_root,
+            end_orchard_root: orchard_root,
+            orchard_tx: orchard_tx_count,
+        }
+    }
 }

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

@@ -49,7 +49,8 @@ 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::<V1>::new_from_block(network, block0, &sapling_root0, &Default::default())?;
 
     // Compute root hash of the MMR tree, which will be included in the next block
     let hash0 = tree.hash();
@@ -73,7 +74,9 @@ fn tree_for_network_upgrade(network: Network, network_upgrade: NetworkUpgrade) -
             .get(&(height + 1))
             .expect("test vector exists"),
     );
-    let append = tree.append_leaf(block1, &sapling_root1).unwrap();
+    let append = tree
+        .append_leaf(block1, &sapling_root1, &Default::default())
+        .unwrap();
 
     // Tree how has 3 nodes: two leafs for each block, and one parent node
     // which is the new root

zebra-chain/src/transaction.rs

@@ -690,6 +690,11 @@ impl Transaction {
             .map(|orchard_shielded_data| orchard_shielded_data.flags)
     }
 
+    /// Return if the transaction has any Orchard shielded data.
+    pub fn has_orchard_shielded_data(&self) -> bool {
+        self.orchard_shielded_data().is_some()
+    }
+
     // value balances
 
     /// Return the transparent value balance.