Add ZIP-221 history tree to non-finalized state (#2583)

* Refactor HistoryTree into NonEmptyHistoryTree and HistoryTree

* HistoryTree: use Deref instead of AsRef; remove unneeded PartialEq

* ZIP-221: Validate chain history commitments in the non-finalized state (#2301)

* sketch of implementation

* refined implementation; still incomplete

* update librustzcash, change zcash_history to work with it

* simplified code per review; renamed MMR to HistoryTree

* expand HistoryTree implementation

* handle and propagate errors

* simplify check.rs tracing

* add suggested TODO

* add HistoryTree::prune

* fix bug in pruning

* fix compilation of tests; still need to make them pass

* Apply suggestions from code review

Co-authored-by: teor <teor@riseup.net>

* Apply suggestions from code review

Co-authored-by: teor <teor@riseup.net>

* improvements from code review

* improve check.rs comments and variable names

* fix HistoryTree which should use BTreeMap and not HashMap; fix non_finalized_state prop tests

* fix finalized_state proptest

* fix non_finalized_state tests by setting the correct commitments

* renamed mmr.rs to history_tree.rs

* Add HistoryTree struct

* expand non_finalized_state protest

* fix typo

* 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

* fixes to work with updated HistoryTree

* Improvements from code review

* Add Debug implementations to allow comparing Chains with proptest_assert_eq

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

* Apply suggestions from code review

Co-authored-by: teor <teor@riseup.net>

* Improvements from code review

* Restore blocks returned by PreparedChain since other tests broken; adjust tests with history trees

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

zebra-chain/src/block/commitment.rs

@@ -149,6 +149,12 @@ impl From<[u8; 32]> for ChainHistoryMmrRootHash {
     }
 }
 
+impl From<ChainHistoryMmrRootHash> for [u8; 32] {
+    fn from(hash: ChainHistoryMmrRootHash) -> Self {
+        hash.0
+    }
+}
+
 /// A block commitment to chain history and transaction auth.
 /// - the chain history tree for all ancestors in the current network upgrade,
 ///   and
@@ -170,7 +176,7 @@ pub struct ChainHistoryBlockTxAuthCommitmentHash([u8; 32]);
 /// implement, and ensures that we don't reject blocks or transactions
 /// for a non-enumerated reason.
 #[allow(dead_code, missing_docs)]
-#[derive(Error, Debug, PartialEq)]
+#[derive(Error, Debug, PartialEq, Eq)]
 pub enum CommitmentError {
     #[error("invalid final sapling root: expected {expected:?}, actual: {actual:?}")]
     InvalidFinalSaplingRoot {

zebra-chain/src/history_tree.rs

@@ -33,6 +33,16 @@ pub enum HistoryTreeError {
     IOError(#[from] io::Error),
 }
 
+impl PartialEq for HistoryTreeError {
+    fn eq(&self, other: &Self) -> bool {
+        // Workaround since subtypes do not implement Eq.
+        // This is only used for tests anyway.
+        format!("{:?}", self) == format!("{:?}", other)
+    }
+}
+
+impl Eq for HistoryTreeError {}
+
 /// The inner [Tree] in one of its supported versions.
 #[derive(Debug)]
 enum InnerHistoryTree {
@@ -403,6 +413,30 @@ impl Clone for NonEmptyHistoryTree {
 pub struct HistoryTree(Option<NonEmptyHistoryTree>);
 
 impl HistoryTree {
+    /// Create a HistoryTree from a block.
+    ///
+    /// If the block is pre-Heartwood, it returns an empty history tree.
+    pub fn from_block(
+        network: Network,
+        block: Arc<Block>,
+        sapling_root: &sapling::tree::Root,
+        orchard_root: &orchard::tree::Root,
+    ) -> Result<Self, HistoryTreeError> {
+        let heartwood_height = NetworkUpgrade::Heartwood
+            .activation_height(network)
+            .expect("Heartwood height is known");
+        match block
+            .coinbase_height()
+            .expect("must have height")
+            .cmp(&heartwood_height)
+        {
+            std::cmp::Ordering::Less => Ok(HistoryTree(None)),
+            _ => Ok(
+                NonEmptyHistoryTree::from_block(network, block, sapling_root, orchard_root)?.into(),
+            ),
+        }
+    }
+
     /// Push a block to a maybe-existing HistoryTree, handling network upgrades.
     ///
     /// The tree is updated in-place. It is created when pushing the Heartwood

zebra-state/src/error.rs

@@ -4,8 +4,8 @@ use chrono::{DateTime, Utc};
 use thiserror::Error;
 
 use zebra_chain::{
-    amount, block, orchard, sapling, sprout, transparent, value_balance::ValueBalanceError,
+    amount, block, history_tree::HistoryTreeError, orchard, sapling, sprout, transparent,
-    work::difficulty::CompactDifficulty,
+    value_balance::ValueBalanceError, work::difficulty::CompactDifficulty,
 };
 
 use crate::constants::MIN_TRANSPARENT_COINBASE_MATURITY;
@@ -36,12 +36,12 @@ impl From<BoxError> for CloneError {
 pub type BoxError = Box<dyn std::error::Error + Send + Sync + 'static>;
 
 /// An error describing the reason a block could not be committed to the state.
-#[derive(Debug, Error, Clone, PartialEq, Eq)]
+#[derive(Debug, Error, PartialEq, Eq)]
 #[error("block is not contextually valid")]
 pub struct CommitBlockError(#[from] ValidateContextError);
 
 /// An error describing why a block failed contextual validation.
-#[derive(Debug, Error, Clone, PartialEq, Eq)]
+#[derive(Debug, Error, PartialEq, Eq)]
 #[non_exhaustive]
 #[allow(missing_docs)]
 pub enum ValidateContextError {
@@ -185,6 +185,9 @@ pub enum ValidateContextError {
 
     #[error("error in Orchard note commitment tree")]
     OrchardNoteCommitmentTreeError(#[from] zebra_chain::orchard::tree::NoteCommitmentTreeError),
+
+    #[error("error building the history tree")]
+    HistoryTreeError(#[from] HistoryTreeError),
 }
 
 /// Trait for creating the corresponding duplicate nullifier error from a nullifier.

zebra-state/src/service/arbitrary.rs

@@ -10,6 +10,7 @@ use proptest::{
 use zebra_chain::{
     block::{self, Block},
     fmt::SummaryDebug,
+    history_tree::HistoryTree,
     parameters::NetworkUpgrade,
     LedgerState,
 };
@@ -34,6 +35,7 @@ pub struct PreparedChainTree {
     chain: Arc<SummaryDebug<Vec<PreparedBlock>>>,
     count: BinarySearch,
     network: Network,
+    history_tree: HistoryTree,
 }
 
 impl ValueTree for PreparedChainTree {
@@ -41,10 +43,16 @@ impl ValueTree for PreparedChainTree {
         Arc<SummaryDebug<Vec<PreparedBlock>>>,
         <BinarySearch as ValueTree>::Value,
         Network,
+        HistoryTree,
     );
 
     fn current(&self) -> Self::Value {
-        (self.chain.clone(), self.count.current(), self.network)
+        (
+            self.chain.clone(),
+            self.count.current(),
+            self.network,
+            self.history_tree.clone(),
+        )
     }
 
     fn simplify(&mut self) -> bool {
@@ -59,7 +67,36 @@ impl ValueTree for PreparedChainTree {
 #[derive(Debug, Default)]
 pub struct PreparedChain {
     // the proptests are threaded (not async), so we want to use a threaded mutex here
-    chain: std::sync::Mutex<Option<(Network, Arc<SummaryDebug<Vec<PreparedBlock>>>)>>,
+    chain: std::sync::Mutex<Option<(Network, Arc<SummaryDebug<Vec<PreparedBlock>>>, HistoryTree)>>,
+    // the strategy for generating LedgerStates. If None, it calls [`LedgerState::genesis_strategy`].
+    ledger_strategy: Option<BoxedStrategy<LedgerState>>,
+}
+
+impl PreparedChain {
+    /// Create a PreparedChain strategy with Heartwood-onward blocks.
+    #[cfg(test)]
+    pub(crate) fn new_heartwood() -> Self {
+        // The history tree only works with Heartwood onward.
+        // Since the network will be chosen later, we pick the larger
+        // between the mainnet and testnet Heartwood activation heights.
+        let main_height = NetworkUpgrade::Heartwood
+            .activation_height(Network::Mainnet)
+            .expect("must have height");
+        let test_height = NetworkUpgrade::Heartwood
+            .activation_height(Network::Testnet)
+            .expect("must have height");
+        let height = std::cmp::max(main_height, test_height);
+
+        PreparedChain {
+            ledger_strategy: Some(LedgerState::height_strategy(
+                height,
+                NetworkUpgrade::Nu5,
+                None,
+                false,
+            )),
+            ..Default::default()
+        }
+    }
 }
 
 impl Strategy for PreparedChain {
@@ -70,7 +107,12 @@ impl Strategy for PreparedChain {
         let mut chain = self.chain.lock().unwrap();
         if chain.is_none() {
             // TODO: use the latest network upgrade (#1974)
-            let ledger_strategy = LedgerState::genesis_strategy(NetworkUpgrade::Nu5, None, false);
+            let default_ledger_strategy =
+                LedgerState::genesis_strategy(NetworkUpgrade::Nu5, None, false);
+            let ledger_strategy = self
+                .ledger_strategy
+                .as_ref()
+                .unwrap_or(&default_ledger_strategy);
 
             let (network, blocks) = ledger_strategy
                 .prop_flat_map(|ledger| {
@@ -93,7 +135,16 @@ impl Strategy for PreparedChain {
                 })
                 .new_tree(runner)?
                 .current();
-            *chain = Some((network, Arc::new(SummaryDebug(blocks))));
+            // Generate a history tree from the first block
+            let history_tree = HistoryTree::from_block(
+                network,
+                blocks[0].block.clone(),
+                // Dummy roots since this is only used for tests
+                &Default::default(),
+                &Default::default(),
+            )
+            .expect("history tree should be created");
+            *chain = Some((network, Arc::new(SummaryDebug(blocks)), history_tree));
         }
 
         let chain = chain.clone().expect("should be generated");
@@ -102,6 +153,7 @@ impl Strategy for PreparedChain {
             chain: chain.1,
             count,
             network: chain.0,
+            history_tree: chain.2,
         })
     }
 }

zebra-state/src/service/finalized_state/tests/prop.rs

@@ -21,7 +21,7 @@ fn blocks_with_v5_transactions() -> Result<()> {
         .ok()
         .and_then(|v| v.parse().ok())
         .unwrap_or(DEFAULT_PARTIAL_CHAIN_PROPTEST_CASES)),
-        |((chain, count, network) in PreparedChain::default())| {
+        |((chain, count, network, _history_tree) in PreparedChain::default())| {
             let mut state = FinalizedState::new(&Config::ephemeral(), network);
             let mut height = Height(0);
             // use `count` to minimize test failures, so they are easier to diagnose

zebra-state/src/service/non_finalized_state.rs

@@ -14,6 +14,7 @@ use std::{collections::BTreeSet, mem, ops::Deref, sync::Arc};
 
 use zebra_chain::{
     block::{self, Block},
+    history_tree::HistoryTree,
     orchard,
     parameters::Network,
     sapling,
@@ -129,6 +130,7 @@ impl NonFinalizedState {
             parent_hash,
             finalized_state.sapling_note_commitment_tree(),
             finalized_state.orchard_note_commitment_tree(),
+            finalized_state.history_tree(),
         )?;
 
         // We might have taken a chain, so all validation must happen within
@@ -161,8 +163,10 @@ impl NonFinalizedState {
         finalized_state: &FinalizedState,
     ) -> Result<(), ValidateContextError> {
         let chain = Chain::new(
+            self.network,
             finalized_state.sapling_note_commitment_tree(),
             finalized_state.orchard_note_commitment_tree(),
+            finalized_state.history_tree(),
         );
         let (height, hash) = (prepared.height, prepared.hash);
 
@@ -355,13 +359,14 @@ impl NonFinalizedState {
     /// The chain can be an existing chain in the non-finalized state or a freshly
     /// created fork, if needed.
     ///
-    /// The note commitment trees must be the trees of the finalized tip.
+    /// The trees must be the trees of the finalized tip.
     /// They are used to recreate the trees if a fork is needed.
     fn parent_chain(
         &mut self,
         parent_hash: block::Hash,
         sapling_note_commitment_tree: sapling::tree::NoteCommitmentTree,
         orchard_note_commitment_tree: orchard::tree::NoteCommitmentTree,
+        history_tree: HistoryTree,
     ) -> Result<Box<Chain>, ValidateContextError> {
         match self.take_chain_if(|chain| chain.non_finalized_tip_hash() == parent_hash) {
             // An existing chain in the non-finalized state
@@ -376,6 +381,7 @@ impl NonFinalizedState {
                                 parent_hash,
                                 sapling_note_commitment_tree.clone(),
                                 orchard_note_commitment_tree.clone(),
+                                history_tree.clone(),
                             )
                             .transpose()
                     })

zebra-state/src/service/non_finalized_state/chain.rs

@@ -8,14 +8,16 @@ use multiset::HashMultiSet;
 use tracing::instrument;
 
 use zebra_chain::{
-    block, orchard, primitives::Groth16Proof, sapling, sprout, transaction,
+    block, history_tree::HistoryTree, orchard, parameters::Network, primitives::Groth16Proof,
-    transaction::Transaction::*, transparent, work::difficulty::PartialCumulativeWork,
+    sapling, sprout, transaction, transaction::Transaction::*, transparent,
+    work::difficulty::PartialCumulativeWork,
 };
 
 use crate::{service::check, ContextuallyValidBlock, PreparedBlock, ValidateContextError};
 
 #[derive(Debug, Clone)]
 pub struct Chain {
+    network: Network,
     /// The contextually valid blocks which form this non-finalized partial chain, in height order.
     pub(crate) blocks: BTreeMap<block::Height, ContextuallyValidBlock>,
 
@@ -37,6 +39,8 @@ pub struct Chain {
     pub(super) sapling_note_commitment_tree: sapling::tree::NoteCommitmentTree,
     /// The Orchard note commitment tree of the tip of this Chain.
     pub(super) orchard_note_commitment_tree: orchard::tree::NoteCommitmentTree,
+    /// The ZIP-221 history tree of the tip of this Chain.
+    pub(crate) history_tree: HistoryTree,
 
     /// The Sapling anchors created by `blocks`.
     pub(super) sapling_anchors: HashMultiSet<sapling::tree::Root>,
@@ -59,12 +63,15 @@ pub struct Chain {
 }
 
 impl Chain {
-    // Create a new Chain with the given note commitment trees.
+    /// Create a new Chain with the given trees and network.
     pub(crate) fn new(
+        network: Network,
         sapling_note_commitment_tree: sapling::tree::NoteCommitmentTree,
         orchard_note_commitment_tree: orchard::tree::NoteCommitmentTree,
+        history_tree: HistoryTree,
     ) -> Self {
         Self {
+            network,
             blocks: Default::default(),
             height_by_hash: Default::default(),
             tx_by_hash: Default::default(),
@@ -80,6 +87,7 @@ impl Chain {
             sapling_nullifiers: Default::default(),
             orchard_nullifiers: Default::default(),
             partial_cumulative_work: Default::default(),
+            history_tree,
         }
     }
 
@@ -95,6 +103,8 @@ impl Chain {
     /// even if the blocks in the two chains are equal.
     #[cfg(test)]
     pub(crate) fn eq_internal_state(&self, other: &Chain) -> bool {
+        use zebra_chain::history_tree::NonEmptyHistoryTree;
+
         // this method must be updated every time a field is added to Chain
 
         // blocks, heights, hashes
@@ -110,6 +120,9 @@ impl Chain {
             self.sapling_note_commitment_tree.root() == other.sapling_note_commitment_tree.root() &&
             self.orchard_note_commitment_tree.root() == other.orchard_note_commitment_tree.root() &&
 
+            // history tree
+            self.history_tree.as_ref().map(NonEmptyHistoryTree::hash) == other.history_tree.as_ref().map(NonEmptyHistoryTree::hash) &&
+
             // anchors
             self.sapling_anchors == other.sapling_anchors &&
             self.sapling_anchors_by_height == other.sapling_anchors_by_height &&
@@ -169,19 +182,24 @@ impl Chain {
     /// Fork a chain at the block with the given hash, if it is part of this
     /// chain.
     ///
-    /// The note commitment trees must be the trees of the finalized tip.
+    /// The trees must match the trees of the finalized tip and are used
+    /// to rebuild them after the fork.
     pub fn fork(
         &self,
         fork_tip: block::Hash,
         sapling_note_commitment_tree: sapling::tree::NoteCommitmentTree,
         orchard_note_commitment_tree: orchard::tree::NoteCommitmentTree,
+        history_tree: HistoryTree,
     ) -> Result<Option<Self>, ValidateContextError> {
         if !self.height_by_hash.contains_key(&fork_tip) {
             return Ok(None);
         }
 
-        let mut forked =
+        let mut forked = self.with_trees(
-            self.with_trees(sapling_note_commitment_tree, orchard_note_commitment_tree);
+            sapling_note_commitment_tree,
+            orchard_note_commitment_tree,
+            history_tree,
+        );
 
         while forked.non_finalized_tip_hash() != fork_tip {
             forked.pop_tip();
@@ -206,6 +224,24 @@ impl Chain {
                         .expect("must work since it was already appended before the fork");
                 }
             }
+
+            // Note that anchors don't need to be recreated since they are already
+            // handled in revert_chain_state_with.
+
+            let sapling_root = forked
+                .sapling_anchors_by_height
+                .get(&block.height)
+                .expect("Sapling anchors must exist for pre-fork blocks");
+            let orchard_root = forked
+                .orchard_anchors_by_height
+                .get(&block.height)
+                .expect("Orchard anchors must exist for pre-fork blocks");
+            forked.history_tree.push(
+                self.network,
+                block.block.clone(),
+                *sapling_root,
+                *orchard_root,
+            )?;
         }
 
         Ok(Some(forked))
@@ -267,8 +303,10 @@ impl Chain {
         &self,
         sapling_note_commitment_tree: sapling::tree::NoteCommitmentTree,
         orchard_note_commitment_tree: orchard::tree::NoteCommitmentTree,
+        history_tree: HistoryTree,
     ) -> Self {
         Chain {
+            network: self.network,
             blocks: self.blocks.clone(),
             height_by_hash: self.height_by_hash.clone(),
             tx_by_hash: self.tx_by_hash.clone(),
@@ -284,6 +322,7 @@ impl Chain {
             sapling_nullifiers: self.sapling_nullifiers.clone(),
             orchard_nullifiers: self.orchard_nullifiers.clone(),
             partial_cumulative_work: self.partial_cumulative_work,
+            history_tree,
         }
     }
 }
@@ -395,12 +434,19 @@ impl UpdateWith<ContextuallyValidBlock> for Chain {
         // Having updated all the note commitment trees and nullifier sets in
         // this block, the roots of the note commitment trees as of the last
         // transaction are the treestates of this block.
-        let root = self.sapling_note_commitment_tree.root();
+        let sapling_root = self.sapling_note_commitment_tree.root();
-        self.sapling_anchors.insert(root);
+        self.sapling_anchors.insert(sapling_root);
-        self.sapling_anchors_by_height.insert(height, root);
+        self.sapling_anchors_by_height.insert(height, sapling_root);
-        let root = self.orchard_note_commitment_tree.root();
+        let orchard_root = self.orchard_note_commitment_tree.root();
-        self.orchard_anchors.insert(root);
+        self.orchard_anchors.insert(orchard_root);
-        self.orchard_anchors_by_height.insert(height, root);
+        self.orchard_anchors_by_height.insert(height, orchard_root);
+
+        self.history_tree.push(
+            self.network,
+            contextually_valid.block.clone(),
+            sapling_root,
+            orchard_root,
+        )?;
 
         Ok(())
     }
@@ -429,6 +475,11 @@ impl UpdateWith<ContextuallyValidBlock> for Chain {
             .expect("work has already been validated");
         self.partial_cumulative_work -= block_work;
 
+        // Note: the history tree is not modified in this method.
+        // This method is called on two scenarios:
+        // - When popping the root: the history tree does not change.
+        // - When popping the tip: the history tree is rebuilt in fork().
+
         // for each transaction in block
         for (transaction, transaction_hash) in
             block.transactions.iter().zip(transaction_hashes.iter())

zebra-state/src/service/non_finalized_state/tests/prop.rs

@@ -3,9 +3,11 @@ use std::{env, sync::Arc};
 use zebra_test::prelude::*;
 
 use zebra_chain::{
-    block::{self, Block},
+    block::{self, arbitrary::allow_all_transparent_coinbase_spends, Block},
     fmt::DisplayToDebug,
+    history_tree::{HistoryTree, NonEmptyHistoryTree},
     parameters::NetworkUpgrade::*,
+    parameters::{Network, *},
     LedgerState,
 };
 
@@ -33,12 +35,15 @@ fn forked_equals_pushed() -> Result<()> {
                                           .ok()
                                           .and_then(|v| v.parse().ok())
                                           .unwrap_or(DEFAULT_PARTIAL_CHAIN_PROPTEST_CASES)),
-        |((chain, fork_at_count, _network) in PreparedChain::default())| {
+        |((chain, fork_at_count, network, finalized_tree) in PreparedChain::new_heartwood())| {
+            // Skip first block which was used for the history tree; make sure fork_at_count is still valid
+            let fork_at_count = std::cmp::min(fork_at_count, chain.len() - 1);
+            let chain = &chain[1..];
             // use `fork_at_count` as the fork tip
             let fork_tip_hash = chain[fork_at_count - 1].hash;
 
-            let mut full_chain = Chain::new(Default::default(), Default::default());
+            let mut full_chain = Chain::new(network, Default::default(), Default::default(), finalized_tree.clone());
-            let mut partial_chain = Chain::new(Default::default(), Default::default());
+            let mut partial_chain = Chain::new(network, Default::default(), Default::default(), finalized_tree.clone());
 
             for block in chain.iter().take(fork_at_count) {
                 partial_chain = partial_chain.push(block.clone())?;
@@ -65,11 +70,12 @@ fn forked_equals_pushed() -> Result<()> {
                 }
             }
 
-            let forked = full_chain
+            let mut forked = full_chain
                 .fork(
                     fork_tip_hash,
                     Default::default(),
                     Default::default(),
+                    finalized_tree,
                 )
                 .expect("fork works")
                 .expect("hash is present");
@@ -77,6 +83,15 @@ fn forked_equals_pushed() -> Result<()> {
             // the first check is redundant, but it's useful for debugging
             prop_assert_eq!(forked.blocks.len(), partial_chain.blocks.len());
             prop_assert!(forked.eq_internal_state(&partial_chain));
+
+            // Re-add blocks to the fork and check if we arrive at the
+            // same original full chain
+            for block in chain.iter().skip(fork_at_count) {
+                forked = forked.push(block.clone())?;
+            }
+
+            prop_assert_eq!(forked.blocks.len(), full_chain.blocks.len());
+            prop_assert!(forked.eq_internal_state(&full_chain));
         });
 
     Ok(())
@@ -92,17 +107,22 @@ fn finalized_equals_pushed() -> Result<()> {
                                       .ok()
                                       .and_then(|v| v.parse().ok())
                                       .unwrap_or(DEFAULT_PARTIAL_CHAIN_PROPTEST_CASES)),
-    |((chain, end_count, _network) in PreparedChain::default())| {
+    |((chain, end_count, network, finalized_tree) in PreparedChain::new_heartwood())| {
+        // Skip first block which was used for the history tree; make sure end_count is still valid
+        let end_count = std::cmp::min(end_count, chain.len() - 1);
+        let chain = &chain[1..];
         // use `end_count` as the number of non-finalized blocks at the end of the chain
         let finalized_count = chain.len() - end_count;
-        let mut full_chain = Chain::new(Default::default(), Default::default());
+        let mut full_chain = Chain::new(network, Default::default(), Default::default(), finalized_tree);
 
         for block in chain.iter().take(finalized_count) {
             full_chain = full_chain.push(block.clone())?;
         }
         let mut partial_chain = Chain::new(
+            network,
             full_chain.sapling_note_commitment_tree.clone(),
             full_chain.orchard_note_commitment_tree.clone(),
+            full_chain.history_tree.clone(),
         );
         for block in chain.iter().skip(finalized_count) {
             partial_chain = partial_chain.push(block.clone())?;
@@ -133,7 +153,7 @@ fn rejection_restores_internal_state() -> Result<()> {
                                          .and_then(|v| v.parse().ok())
                                          .unwrap_or(DEFAULT_PARTIAL_CHAIN_PROPTEST_CASES)),
               |((chain, valid_count, network, mut bad_block) in (PreparedChain::default(), any::<bool>(), any::<bool>())
-                .prop_flat_map(|((chain, valid_count, network), is_nu5, is_v5)| {
+                .prop_flat_map(|((chain, valid_count, network, _history_tree), is_nu5, is_v5)| {
                     let next_height = chain[valid_count - 1].height;
                     (
                         Just(chain),
@@ -205,7 +225,7 @@ fn different_blocks_different_chains() -> Result<()> {
                                .ok()
                                .and_then(|v| v.parse().ok())
                                .unwrap_or(DEFAULT_PARTIAL_CHAIN_PROPTEST_CASES)),
-    |((block1, block2) in (any::<bool>(), any::<bool>())
+    |((vec1, vec2) in (any::<bool>(), any::<bool>())
       .prop_flat_map(|(is_nu5, is_v5)| {
           // generate a Canopy or NU5 block with v4 or v5 transactions
           LedgerState::coinbase_strategy(
@@ -213,15 +233,28 @@ fn different_blocks_different_chains() -> Result<()> {
               if is_nu5 && is_v5 { 5 } else { 4 },
               true,
           )})
-      .prop_map(Block::arbitrary_with)
+      .prop_map(|ledger_state| Block::partial_chain_strategy(ledger_state, 2, allow_all_transparent_coinbase_spends))
       .prop_flat_map(|block_strategy| (block_strategy.clone(), block_strategy))
-      .prop_map(|(block1, block2)| (DisplayToDebug(block1), DisplayToDebug(block2)))
     )| {
-        let chain1 = Chain::new(Default::default(), Default::default());
+        let prev_block1 = vec1[0].clone();
-        let chain2 = Chain::new(Default::default(), Default::default());
+        let prev_block2 = vec2[0].clone();
+        let height1 = prev_block1.coinbase_height().unwrap();
+        let height2 = prev_block1.coinbase_height().unwrap();
+        let finalized_tree1: HistoryTree = if height1 >= Heartwood.activation_height(Network::Mainnet).unwrap() {
+            NonEmptyHistoryTree::from_block(Network::Mainnet, prev_block1, &Default::default(), &Default::default()).unwrap().into()
+        } else {
+            Default::default()
+        };
+        let finalized_tree2 = if height2 >= NetworkUpgrade::Heartwood.activation_height(Network::Mainnet).unwrap() {
+            NonEmptyHistoryTree::from_block(Network::Mainnet, prev_block2, &Default::default(), &Default::default()).unwrap().into()
+        } else {
+            Default::default()
+        };
+        let chain1 = Chain::new(Network::Mainnet, Default::default(), Default::default(), finalized_tree1);
+        let chain2 = Chain::new(Network::Mainnet, Default::default(), Default::default(), finalized_tree2);
 
-        let block1 = Arc::new(block1.0).prepare();
+        let block1 = vec1[1].clone().prepare();
-        let block2 = Arc::new(block2.0).prepare();
+        let block2 = vec2[1].clone().prepare();
 
         let result1 = chain1.push(block1.clone());
         let result2 = chain2.push(block2.clone());
@@ -256,6 +289,9 @@ fn different_blocks_different_chains() -> Result<()> {
                 chain1.sapling_note_commitment_tree = chain2.sapling_note_commitment_tree.clone();
                 chain1.orchard_note_commitment_tree = chain2.orchard_note_commitment_tree.clone();
 
+                // history tree
+                chain1.history_tree = chain2.history_tree.clone();
+
                 // anchors
                 chain1.sapling_anchors = chain2.sapling_anchors.clone();
                 chain1.sapling_anchors_by_height = chain2.sapling_anchors_by_height.clone();

zebra-state/src/service/non_finalized_state/tests/vectors.rs

@@ -1,6 +1,10 @@
 use std::sync::Arc;
 
-use zebra_chain::{block::Block, parameters::Network, serialization::ZcashDeserializeInto};
+use zebra_chain::{
+    block::Block,
+    parameters::{Network, NetworkUpgrade},
+    serialization::ZcashDeserializeInto,
+};
 use zebra_test::prelude::*;
 
 use crate::{
@@ -18,7 +22,12 @@ use self::assert_eq;
 #[test]
 fn construct_empty() {
     zebra_test::init();
-    let _chain = Chain::new(Default::default(), Default::default());
+    let _chain = Chain::new(
+        Network::Mainnet,
+        Default::default(),
+        Default::default(),
+        Default::default(),
+    );
 }
 
 #[test]
@@ -27,7 +36,12 @@ fn construct_single() -> Result<()> {
     let block: Arc<Block> =
         zebra_test::vectors::BLOCK_MAINNET_434873_BYTES.zcash_deserialize_into()?;
 
-    let mut chain = Chain::new(Default::default(), Default::default());
+    let mut chain = Chain::new(
+        Network::Mainnet,
+        Default::default(),
+        Default::default(),
+        Default::default(),
+    );
     chain = chain.push(block.prepare())?;
 
     assert_eq!(1, chain.blocks.len());
@@ -49,7 +63,12 @@ fn construct_many() -> Result<()> {
         block = next_block;
     }
 
-    let mut chain = Chain::new(Default::default(), Default::default());
+    let mut chain = Chain::new(
+        Network::Mainnet,
+        Default::default(),
+        Default::default(),
+        Default::default(),
+    );
 
     for block in blocks {
         chain = chain.push(block.prepare())?;
@@ -68,10 +87,20 @@ fn ord_matches_work() -> Result<()> {
         .set_work(1);
     let more_block = less_block.clone().set_work(10);
 
-    let mut lesser_chain = Chain::new(Default::default(), Default::default());
+    let mut lesser_chain = Chain::new(
+        Network::Mainnet,
+        Default::default(),
+        Default::default(),
+        Default::default(),
+    );
     lesser_chain = lesser_chain.push(less_block.prepare())?;
 
-    let mut bigger_chain = Chain::new(Default::default(), Default::default());
+    let mut bigger_chain = Chain::new(
+        Network::Mainnet,
+        Default::default(),
+        Default::default(),
+        Default::default(),
+    );
     bigger_chain = bigger_chain.push(more_block.prepare())?;
 
     assert!(bigger_chain > lesser_chain);
@@ -91,11 +120,14 @@ fn best_chain_wins() -> Result<()> {
 
 fn best_chain_wins_for_network(network: Network) -> Result<()> {
     let block1: Arc<Block> = match network {
+        // Since the brand new FinalizedState below will pass a None history tree
+        // to the NonFinalizedState, we must use pre-Heartwood blocks since
+        // they won't trigger the history tree update in the NonFinalizedState.
         Network::Mainnet => {
-            zebra_test::vectors::BLOCK_MAINNET_1180900_BYTES.zcash_deserialize_into()?
+            zebra_test::vectors::BLOCK_MAINNET_653599_BYTES.zcash_deserialize_into()?
         }
         Network::Testnet => {
-            zebra_test::vectors::BLOCK_TESTNET_1326100_BYTES.zcash_deserialize_into()?
+            zebra_test::vectors::BLOCK_TESTNET_583999_BYTES.zcash_deserialize_into()?
         }
     };
 
@@ -128,11 +160,14 @@ fn finalize_pops_from_best_chain() -> Result<()> {
 
 fn finalize_pops_from_best_chain_for_network(network: Network) -> Result<()> {
     let block1: Arc<Block> = match network {
+        // Since the brand new FinalizedState below will pass a None history tree
+        // to the NonFinalizedState, we must use pre-Heartwood blocks since
+        // they won't trigger the history tree update in the NonFinalizedState.
         Network::Mainnet => {
-            zebra_test::vectors::BLOCK_MAINNET_1180900_BYTES.zcash_deserialize_into()?
+            zebra_test::vectors::BLOCK_MAINNET_653599_BYTES.zcash_deserialize_into()?
         }
         Network::Testnet => {
-            zebra_test::vectors::BLOCK_TESTNET_1326100_BYTES.zcash_deserialize_into()?
+            zebra_test::vectors::BLOCK_TESTNET_583999_BYTES.zcash_deserialize_into()?
         }
     };
 
@@ -172,11 +207,14 @@ fn commit_block_extending_best_chain_doesnt_drop_worst_chains_for_network(
     network: Network,
 ) -> Result<()> {
     let block1: Arc<Block> = match network {
+        // Since the brand new FinalizedState below will pass a None history tree
+        // to the NonFinalizedState, we must use pre-Heartwood blocks since
+        // they won't trigger the history tree update in the NonFinalizedState.
         Network::Mainnet => {
-            zebra_test::vectors::BLOCK_MAINNET_1180900_BYTES.zcash_deserialize_into()?
+            zebra_test::vectors::BLOCK_MAINNET_653599_BYTES.zcash_deserialize_into()?
         }
         Network::Testnet => {
-            zebra_test::vectors::BLOCK_TESTNET_1326100_BYTES.zcash_deserialize_into()?
+            zebra_test::vectors::BLOCK_TESTNET_583999_BYTES.zcash_deserialize_into()?
         }
     };
 
@@ -212,11 +250,14 @@ fn shorter_chain_can_be_best_chain() -> Result<()> {
 
 fn shorter_chain_can_be_best_chain_for_network(network: Network) -> Result<()> {
     let block1: Arc<Block> = match network {
+        // Since the brand new FinalizedState below will pass a None history tree
+        // to the NonFinalizedState, we must use pre-Heartwood blocks since
+        // they won't trigger the history tree update in the NonFinalizedState.
         Network::Mainnet => {
-            zebra_test::vectors::BLOCK_MAINNET_1180900_BYTES.zcash_deserialize_into()?
+            zebra_test::vectors::BLOCK_MAINNET_653599_BYTES.zcash_deserialize_into()?
         }
         Network::Testnet => {
-            zebra_test::vectors::BLOCK_TESTNET_1326100_BYTES.zcash_deserialize_into()?
+            zebra_test::vectors::BLOCK_TESTNET_583999_BYTES.zcash_deserialize_into()?
         }
     };
 
@@ -251,11 +292,14 @@ fn longer_chain_with_more_work_wins() -> Result<()> {
 
 fn longer_chain_with_more_work_wins_for_network(network: Network) -> Result<()> {
     let block1: Arc<Block> = match network {
+        // Since the brand new FinalizedState below will pass a None history tree
+        // to the NonFinalizedState, we must use pre-Heartwood blocks since
+        // they won't trigger the history tree update in the NonFinalizedState.
         Network::Mainnet => {
-            zebra_test::vectors::BLOCK_MAINNET_1180900_BYTES.zcash_deserialize_into()?
+            zebra_test::vectors::BLOCK_MAINNET_653599_BYTES.zcash_deserialize_into()?
         }
         Network::Testnet => {
-            zebra_test::vectors::BLOCK_TESTNET_1326100_BYTES.zcash_deserialize_into()?
+            zebra_test::vectors::BLOCK_TESTNET_583999_BYTES.zcash_deserialize_into()?
         }
     };
 
@@ -293,11 +337,14 @@ fn equal_length_goes_to_more_work() -> Result<()> {
 }
 fn equal_length_goes_to_more_work_for_network(network: Network) -> Result<()> {
     let block1: Arc<Block> = match network {
+        // Since the brand new FinalizedState below will pass a None history tree
+        // to the NonFinalizedState, we must use pre-Heartwood blocks since
+        // they won't trigger the history tree update in the NonFinalizedState.
         Network::Mainnet => {
-            zebra_test::vectors::BLOCK_MAINNET_1180900_BYTES.zcash_deserialize_into()?
+            zebra_test::vectors::BLOCK_MAINNET_653599_BYTES.zcash_deserialize_into()?
         }
         Network::Testnet => {
-            zebra_test::vectors::BLOCK_TESTNET_1326100_BYTES.zcash_deserialize_into()?
+            zebra_test::vectors::BLOCK_TESTNET_583999_BYTES.zcash_deserialize_into()?
         }
     };
 
@@ -318,3 +365,73 @@ fn equal_length_goes_to_more_work_for_network(network: Network) -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn history_tree_is_updated() -> Result<()> {
+    history_tree_is_updated_for_network_upgrade(Network::Mainnet, NetworkUpgrade::Heartwood)?;
+    history_tree_is_updated_for_network_upgrade(Network::Testnet, NetworkUpgrade::Heartwood)?;
+    // TODO: we can't test other upgrades until we have a method for creating a FinalizedState
+    // with a HistoryTree.
+    Ok(())
+}
+
+fn history_tree_is_updated_for_network_upgrade(
+    network: Network,
+    network_upgrade: NetworkUpgrade,
+) -> Result<()> {
+    let blocks = match network {
+        Network::Mainnet => &*zebra_test::vectors::MAINNET_BLOCKS,
+        Network::Testnet => &*zebra_test::vectors::TESTNET_BLOCKS,
+    };
+    let height = network_upgrade.activation_height(network).unwrap().0;
+
+    let prev_block = Arc::new(
+        blocks
+            .get(&(height - 1))
+            .expect("test vector exists")
+            .zcash_deserialize_into::<Block>()
+            .expect("block is structurally valid"),
+    );
+    let activation_block = prev_block.make_fake_child();
+    let next_block = activation_block.make_fake_child();
+
+    let mut state = NonFinalizedState::new(network);
+    let finalized_state = FinalizedState::new(&Config::ephemeral(), network);
+
+    state.commit_new_chain(prev_block.prepare(), &finalized_state)?;
+
+    let chain = state.best_chain().unwrap();
+    if network_upgrade == NetworkUpgrade::Heartwood {
+        assert!(
+            chain.history_tree.as_ref().is_none(),
+            "history tree must not exist yet"
+        );
+    } else {
+        assert!(
+            chain.history_tree.as_ref().is_some(),
+            "history tree must already exist"
+        );
+    }
+
+    state.commit_block(activation_block.prepare(), &finalized_state)?;
+
+    let chain = state.best_chain().unwrap();
+    assert!(
+        chain.history_tree.as_ref().is_some(),
+        "history tree must have been (re)created"
+    );
+    assert_eq!(
+        chain.history_tree.as_ref().as_ref().unwrap().size(),
+        1,
+        "history tree must have a single node"
+    );
+
+    state.commit_block(next_block.prepare(), &finalized_state)?;
+
+    assert!(
+        state.best_chain().unwrap().history_tree.as_ref().is_some(),
+        "history tree must still exist"
+    );
+
+    Ok(())
+}