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>

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

zebra-chain/src/history_tree.rs

@@ -32,6 +32,7 @@ pub enum HistoryTreeError {
 
 /// History tree (Merkle mountain range) structure that contains information about
 // the block history, as specified in [ZIP-221][https://zips.z.cash/zip-0221].
+#[derive(Debug)]
 pub struct HistoryTree {
     network: Network,
     network_upgrade: NetworkUpgrade,
@@ -244,3 +245,11 @@ impl Clone for HistoryTree {
         }
     }
 }
+
+impl PartialEq for HistoryTree {
+    fn eq(&self, other: &Self) -> bool {
+        self.hash() == other.hash()
+    }
+}
+
+impl Eq for HistoryTree {}

zebra-chain/src/primitives/zcash_history.rs

@@ -45,7 +45,7 @@ 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)]
+#[derive(Clone, Debug)]
 pub struct Entry {
     inner: [u8; zcash_history::MAX_ENTRY_SIZE],
 }
@@ -231,6 +231,15 @@ impl Tree {
     }
 }
 
+impl std::fmt::Debug for Tree {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("Tree")
+            .field("network", &self.network)
+            .field("network_upgrade", &self.network_upgrade)
+            .finish()
+    }
+}
+
 /// 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.

zebra-state/src/error.rs

@@ -3,7 +3,11 @@ use std::sync::Arc;
 use chrono::{DateTime, Utc};
 use thiserror::Error;
 
-use zebra_chain::{block, work::difficulty::CompactDifficulty};
+use zebra_chain::{
+    block::{self, ChainHistoryMmrRootHash},
+    history_tree::HistoryTreeError,
+    work::difficulty::CompactDifficulty,
+};
 
 /// A wrapper for type erased errors that is itself clonable and implements the
 /// Error trait
@@ -74,4 +78,17 @@ pub enum ValidateContextError {
         difficulty_threshold: CompactDifficulty,
         expected_difficulty: CompactDifficulty,
     },
+
+    #[error("block contains an invalid commitment")]
+    InvalidBlockCommitment(#[from] block::CommitmentError),
+
+    #[error("block history commitment {candidate_commitment:?} is different to the expected commitment {expected_commitment:?}")]
+    #[non_exhaustive]
+    InvalidHistoryCommitment {
+        candidate_commitment: ChainHistoryMmrRootHash,
+        expected_commitment: ChainHistoryMmrRootHash,
+    },
+
+    #[error("error building the history tree")]
+    HistoryTreeError(#[from] HistoryTreeError),
 }

zebra-state/src/service.rs

@@ -178,9 +178,10 @@ impl StateService {
         let parent_hash = prepared.block.header.previous_block_hash;
 
         if self.disk.finalized_tip_hash() == parent_hash {
-            self.mem.commit_new_chain(prepared)?;
+            self.mem
+                .commit_new_chain(prepared, self.disk.history_tree().clone())?;
         } else {
-            self.mem.commit_block(prepared)?;
+            self.mem.commit_block(prepared, self.disk.history_tree())?;
         }
 
         Ok(())
@@ -222,7 +223,7 @@ impl StateService {
         assert!(relevant_chain.len() >= POW_AVERAGING_WINDOW + POW_MEDIAN_BLOCK_SPAN,
                 "contextual validation requires at least 28 (POW_AVERAGING_WINDOW + POW_MEDIAN_BLOCK_SPAN) blocks");
 
-        check::block_is_contextually_valid(
+        check::block_is_valid_for_recent_chain(
             prepared,
             self.network,
             self.disk.finalized_tip_height(),

zebra-state/src/service/arbitrary.rs

@@ -50,6 +50,29 @@ impl ValueTree for PreparedChainTree {
 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>>>)>>,
+    // the height from which to start the chain. If None, starts at the genesis block
+    start_height: Option<Height>,
+}
+
+impl PreparedChain {
+    /// Create a PreparedChain strategy with Heartwood-onward blocks.
+    pub(super) 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) + 1).expect("must be valid");
+
+        PreparedChain {
+            start_height: Some(height),
+            ..Default::default()
+        }
+    }
 }
 
 impl Strategy for PreparedChain {
@@ -60,7 +83,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 ledger_strategy = match self.start_height {
+                Some(start_height) => {
+                    LedgerState::height_strategy(start_height, NetworkUpgrade::Nu5, None, false)
+                }
+                None => LedgerState::genesis_strategy(NetworkUpgrade::Nu5, None, false),
+            };
 
             let (network, blocks) = ledger_strategy
                 .prop_flat_map(|ledger| {
@@ -83,7 +111,9 @@ impl Strategy for PreparedChain {
         }
 
         let chain = chain.clone().expect("should be generated");
-        let count = (1..chain.1.len()).new_tree(runner)?;
+        // `count` must be 1 less since the first block is used to build the
+        // history tree.
+        let count = (1..chain.1.len() - 1).new_tree(runner)?;
         Ok(PreparedChainTree {
             chain: chain.1,
             count,

zebra-state/src/service/check.rs

@@ -4,7 +4,7 @@ use std::borrow::Borrow;
 
 use chrono::Duration;
 use zebra_chain::{
-    block::{self, Block},
+    block::{self, Block, ChainHistoryMmrRootHash},
     parameters::POW_AVERAGING_WINDOW,
     parameters::{Network, NetworkUpgrade},
     work::difficulty::CompactDifficulty,
@@ -18,8 +18,11 @@ use difficulty::{AdjustedDifficulty, POW_MEDIAN_BLOCK_SPAN};
 
 pub(crate) mod difficulty;
 
-/// Check that `block` is contextually valid for `network`, based on the
+/// Check that the `prepared` block is contextually valid for `network`, based
-/// `finalized_tip_height` and `relevant_chain`.
+/// on the `finalized_tip_height` and `relevant_chain`.
+///
+/// This function performs checks that require a small number of recent blocks,
+/// including previous hash, previous height, and block difficulty.
 ///
 /// The relevant chain is an iterator over the ancestors of `block`, starting
 /// with its parent block.
@@ -28,12 +31,8 @@ pub(crate) mod difficulty;
 ///
 /// If the state contains less than 28
 /// (`POW_AVERAGING_WINDOW + POW_MEDIAN_BLOCK_SPAN`) blocks.
-#[tracing::instrument(
+#[tracing::instrument(skip(prepared, finalized_tip_height, relevant_chain))]
-    name = "contextual_validation",
+pub(crate) fn block_is_valid_for_recent_chain<C>(
-    fields(?network),
-    skip(prepared, network, finalized_tip_height, relevant_chain)
-)]
-pub(crate) fn block_is_contextually_valid<C>(
     prepared: &PreparedBlock,
     network: Network,
     finalized_tip_height: Option<block::Height>,
@@ -86,6 +85,40 @@ where
     Ok(())
 }
 
+/// Check that the `prepared` block is contextually valid for `network`, based
+/// on the `history_root_hash` of the history tree up to and including the
+/// previous block.
+#[tracing::instrument(skip(prepared))]
+pub(crate) fn block_commitment_is_valid_for_chain_history(
+    prepared: &PreparedBlock,
+    network: Network,
+    history_root_hash: &ChainHistoryMmrRootHash,
+) -> Result<(), ValidateContextError> {
+    match prepared.block.commitment(network)? {
+        block::Commitment::PreSaplingReserved(_)
+        | block::Commitment::FinalSaplingRoot(_)
+        | block::Commitment::ChainHistoryActivationReserved => {
+            // No contextual checks needed for those.
+            Ok(())
+        }
+        block::Commitment::ChainHistoryRoot(block_history_root_hash) => {
+            if block_history_root_hash == *history_root_hash {
+                Ok(())
+            } else {
+                Err(ValidateContextError::InvalidHistoryCommitment {
+                    candidate_commitment: block_history_root_hash,
+                    expected_commitment: *history_root_hash,
+                })
+            }
+        }
+        block::Commitment::ChainHistoryBlockTxAuthCommitment(_) => {
+            // TODO: Get auth_hash from block (ZIP-244), e.g.
+            // let auth_hash = prepared.block.auth_hash();
+            todo!("hash mmr_hash and auth_hash per ZIP-244 and compare")
+        }
+    }
+}
+
 /// Returns `ValidateContextError::OrphanedBlock` if the height of the given
 /// block is less than or equal to the finalized tip height.
 fn block_is_not_orphaned(

zebra-state/src/service/finalized_state.rs

@@ -7,6 +7,7 @@ mod tests;
 
 use std::{collections::HashMap, convert::TryInto, path::Path, sync::Arc};
 
+use zebra_chain::history_tree::HistoryTree;
 use zebra_chain::transparent;
 use zebra_chain::{
     block::{self, Block},
@@ -378,6 +379,11 @@ impl FinalizedState {
             })
     }
 
+    /// Returns the history tree for the finalized state.
+    pub fn history_tree(&self) -> &HistoryTree {
+        todo!("add history tree to finalized state");
+    }
+
     /// If the database is `ephemeral`, delete it.
     fn delete_ephemeral(&self) {
         if self.ephemeral {

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,
     parameters::Network,
     transaction::{self, Transaction},
     transparent,
@@ -23,6 +24,8 @@ use crate::{FinalizedBlock, HashOrHeight, PreparedBlock, Utxo, ValidateContextEr
 
 use self::chain::Chain;
 
+use super::check;
+
 /// The state of the chains in memory, incuding queued blocks.
 #[derive(Default)]
 pub struct NonFinalizedState {
@@ -74,7 +77,14 @@ impl NonFinalizedState {
     }
 
     /// Commit block to the non-finalized state.
-    pub fn commit_block(&mut self, prepared: PreparedBlock) -> Result<(), ValidateContextError> {
+    ///
+    /// `finalized_tip_history_tree`: the history tree of the finalized tip used to recompute
+    /// the history tree, if needed.
+    pub fn commit_block(
+        &mut self,
+        prepared: PreparedBlock,
+        finalized_tip_history_tree: &HistoryTree,
+    ) -> Result<(), ValidateContextError> {
         let parent_hash = prepared.block.header.previous_block_hash;
         let (height, hash) = (prepared.height, prepared.hash);
 
@@ -85,8 +95,12 @@ impl NonFinalizedState {
             );
         }
 
-        let mut parent_chain = self.parent_chain(parent_hash)?;
+        let mut parent_chain = self.parent_chain(parent_hash, finalized_tip_history_tree)?;
-
+        check::block_commitment_is_valid_for_chain_history(
+            &prepared,
+            self.network,
+            &parent_chain.history_root_hash(),
+        )?;
         parent_chain.push(prepared)?;
         self.chain_set.insert(parent_chain);
         self.update_metrics_for_committed_block(height, hash);
@@ -95,12 +109,20 @@ impl NonFinalizedState {
 
     /// Commit block to the non-finalized state as a new chain where its parent
     /// is the finalized tip.
+    ///
+    /// `history_tree` must contain the history of the finalized tip.
     pub fn commit_new_chain(
         &mut self,
         prepared: PreparedBlock,
+        finalized_tip_history_tree: HistoryTree,
     ) -> Result<(), ValidateContextError> {
-        let mut chain = Chain::default();
+        let mut chain = Chain::new(finalized_tip_history_tree);
         let (height, hash) = (prepared.height, prepared.hash);
+        check::block_commitment_is_valid_for_chain_history(
+            &prepared,
+            self.network,
+            &chain.history_root_hash(),
+        )?;
         chain.push(prepared)?;
         self.chain_set.insert(Box::new(chain));
         self.update_metrics_for_committed_block(height, hash);
@@ -246,9 +268,13 @@ impl NonFinalizedState {
     ///
     /// The chain can be an existing chain in the non-finalized state or a freshly
     /// created fork, if needed.
+    ///
+    /// `finalized_tip_history_tree`: the history tree of the finalized tip used to recompute
+    /// the history tree, if needed.
     fn parent_chain(
         &mut self,
         parent_hash: block::Hash,
+        finalized_tip_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
@@ -257,7 +283,11 @@ impl NonFinalizedState {
             None => Ok(Box::new(
                 self.chain_set
                     .iter()
-                    .find_map(|chain| chain.fork(parent_hash).transpose())
+                    .find_map(|chain| {
+                        chain
+                            .fork(parent_hash, finalized_tip_history_tree)
+                            .transpose()
+                    })
                     .expect(
                         "commit_block is only called with blocks that are ready to be commited",
                     )?,

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

@@ -6,13 +6,19 @@ use std::{
 
 use tracing::{debug_span, instrument, trace};
 use zebra_chain::{
-    block, orchard, primitives::Groth16Proof, sapling, sprout, transaction,
+    block::{self, ChainHistoryMmrRootHash},
-    transaction::Transaction::*, transparent, work::difficulty::PartialCumulativeWork,
+    history_tree::HistoryTree,
+    orchard,
+    primitives::Groth16Proof,
+    sapling, sprout, transaction,
+    transaction::Transaction::*,
+    transparent,
+    work::difficulty::PartialCumulativeWork,
 };
 
 use crate::{PreparedBlock, Utxo, ValidateContextError};
 
-#[derive(Default, Clone)]
+#[derive(PartialEq, Eq, Debug)]
 pub struct Chain {
     pub blocks: BTreeMap<block::Height, PreparedBlock>,
     pub height_by_hash: HashMap<block::Hash, block::Height>,
@@ -27,9 +33,30 @@ pub struct Chain {
     sapling_nullifiers: HashSet<sapling::Nullifier>,
     orchard_nullifiers: HashSet<orchard::Nullifier>,
     partial_cumulative_work: PartialCumulativeWork,
+    pub(crate) history_tree: HistoryTree,
 }
 
 impl Chain {
+    /// Create a new empty non-finalized chain with the given history tree.
+    ///
+    /// The history tree must contain the history of the previous (finalized) blocks.
+    pub fn new(history_tree: HistoryTree) -> Self {
+        Chain {
+            blocks: Default::default(),
+            height_by_hash: Default::default(),
+            tx_by_hash: Default::default(),
+            created_utxos: Default::default(),
+            spent_utxos: Default::default(),
+            sprout_anchors: Default::default(),
+            sapling_anchors: Default::default(),
+            sprout_nullifiers: Default::default(),
+            sapling_nullifiers: Default::default(),
+            orchard_nullifiers: Default::default(),
+            partial_cumulative_work: Default::default(),
+            history_tree,
+        }
+    }
+
     /// Push a contextually valid non-finalized block into a chain as the new tip.
     #[instrument(level = "debug", skip(self, block), fields(block = %block.block))]
     pub fn push(&mut self, block: PreparedBlock) -> Result<(), ValidateContextError> {
@@ -68,17 +95,39 @@ impl Chain {
 
     /// Fork a chain at the block with the given hash, if it is part of this
     /// chain.
-    pub fn fork(&self, fork_tip: block::Hash) -> Result<Option<Self>, ValidateContextError> {
+    ///
+    /// `finalized_tip_history_tree`: the history tree for the finalized tip
+    /// from which the tree of the fork will be computed.
+    pub fn fork(
+        &self,
+        fork_tip: block::Hash,
+        finalized_tip_history_tree: &HistoryTree,
+    ) -> Result<Option<Self>, ValidateContextError> {
         if !self.height_by_hash.contains_key(&fork_tip) {
             return Ok(None);
         }
 
-        let mut forked = self.clone();
+        let mut forked = self.with_history_tree(finalized_tip_history_tree.clone());
 
         while forked.non_finalized_tip_hash() != fork_tip {
             forked.pop_tip();
         }
 
+        // Rebuild the history tree starting from the finalized tip tree.
+        // TODO: change to a more efficient approach by removing nodes
+        // from the tree of the original chain (in `pop_tip()`).
+        // See https://github.com/ZcashFoundation/zebra/issues/2378
+        forked
+            .history_tree
+            .try_extend(forked.blocks.values().map(|prepared_block| {
+                (
+                    prepared_block.block.clone(),
+                    // TODO: pass Sapling and Orchard roots
+                    &sapling::tree::Root([0; 32]),
+                    None,
+                )
+            }))?;
+
         Ok(Some(forked))
     }
 
@@ -118,6 +167,31 @@ impl Chain {
     pub fn is_empty(&self) -> bool {
         self.blocks.is_empty()
     }
+
+    pub fn history_root_hash(&self) -> ChainHistoryMmrRootHash {
+        self.history_tree.hash()
+    }
+
+    /// Clone the Chain but not the history tree, using the history tree
+    /// specified instead.
+    ///
+    /// Useful when forking, where the history tree is rebuilt anyway.
+    fn with_history_tree(&self, history_tree: HistoryTree) -> Self {
+        Chain {
+            blocks: self.blocks.clone(),
+            height_by_hash: self.height_by_hash.clone(),
+            tx_by_hash: self.tx_by_hash.clone(),
+            created_utxos: self.created_utxos.clone(),
+            spent_utxos: self.spent_utxos.clone(),
+            sprout_anchors: self.sprout_anchors.clone(),
+            sapling_anchors: self.sapling_anchors.clone(),
+            sprout_nullifiers: self.sprout_nullifiers.clone(),
+            sapling_nullifiers: self.sapling_nullifiers.clone(),
+            orchard_nullifiers: self.orchard_nullifiers.clone(),
+            partial_cumulative_work: self.partial_cumulative_work,
+            history_tree,
+        }
+    }
 }
 
 /// Helper trait to organize inverse operations done on the `Chain` type. Used to
@@ -164,6 +238,10 @@ impl UpdateWith<PreparedBlock> for Chain {
             .expect("work has already been validated");
         self.partial_cumulative_work += block_work;
 
+        // TODO: pass Sapling and Orchard roots
+        self.history_tree
+            .push(prepared.block.clone(), &sapling::tree::Root([0; 32]), None)?;
+
         // for each transaction in block
         for (transaction_index, (transaction, transaction_hash)) in block
             .transactions
@@ -247,6 +325,11 @@ impl UpdateWith<PreparedBlock> 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())
@@ -445,14 +528,6 @@ impl UpdateWith<Option<orchard::ShieldedData>> for Chain {
     }
 }
 
-impl PartialEq for Chain {
-    fn eq(&self, other: &Self) -> bool {
-        self.partial_cmp(other) == Some(Ordering::Equal)
-    }
-}
-
-impl Eq for Chain {}
-
 impl PartialOrd for Chain {
     fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
         Some(self.cmp(other))

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

@@ -1,5 +1,6 @@
 use std::env;
 
+use zebra_chain::{history_tree::HistoryTree, sapling};
 use zebra_test::prelude::*;
 
 use crate::service::{arbitrary::PreparedChain, non_finalized_state::Chain};
@@ -14,10 +15,14 @@ fn forked_equals_pushed() -> Result<()> {
                                           .ok()
                                           .and_then(|v| v.parse().ok())
                                           .unwrap_or(DEFAULT_PARTIAL_CHAIN_PROPTEST_CASES)),
-        |((chain, count, _network) in PreparedChain::default())| {
+        |((chain, count, network) in PreparedChain::new_heartwood())| {
+            // Build a history tree with the first block to simulate the tree of
+            // the finalized state.
+            let finalized_tree = HistoryTree::from_block(network, chain[0].block.clone(), &sapling::tree::Root::default(), None).unwrap();
+            let chain = &chain[1..];
             let fork_tip_hash = chain[count - 1].hash;
-            let mut full_chain = Chain::default();
+            let mut full_chain = Chain::new(finalized_tree.clone());
-            let mut partial_chain = Chain::default();
+            let mut partial_chain = Chain::new(finalized_tree.clone());
 
             for block in chain.iter().take(count) {
                 partial_chain.push(block.clone())?;
@@ -26,9 +31,17 @@ fn forked_equals_pushed() -> Result<()> {
                 full_chain.push(block.clone())?;
             }
 
-            let forked = full_chain.fork(fork_tip_hash).expect("fork works").expect("hash is present");
+            let mut forked = full_chain.fork(fork_tip_hash, &finalized_tree).expect("fork works").expect("hash is present");
 
             prop_assert_eq!(forked.blocks.len(), partial_chain.blocks.len());
+            prop_assert_eq!(&forked, &partial_chain);
+
+            for block in chain.iter().skip(count) {
+                forked.push(block.clone())?;
+            }
+
+            prop_assert_eq!(forked.blocks.len(), full_chain.blocks.len());
+            prop_assert_eq!(&forked, &full_chain);
         });
 
     Ok(())
@@ -42,23 +55,32 @@ 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) in PreparedChain::new_heartwood())| {
+        // Build a history tree with the first block to simulate the tree of
+        // the finalized state.
+        let finalized_tree = HistoryTree::from_block(network, chain[0].block.clone(), &sapling::tree::Root::default(), None).unwrap();
+        let chain = &chain[1..];
         let finalized_count = chain.len() - end_count;
-        let mut full_chain = Chain::default();
+        let mut full_chain = Chain::new(finalized_tree);
-        let mut partial_chain = Chain::default();
+
+        for block in chain.iter().take(finalized_count) {
+            full_chain.push(block.clone())?;
+        }
+        let mut partial_chain = Chain::new(full_chain.history_tree.clone());
+        for block in chain.iter().skip(finalized_count) {
+            full_chain.push(block.clone())?;
+        }
 
         for block in chain.iter().skip(finalized_count) {
             partial_chain.push(block.clone())?;
         }
-        for block in chain.iter() {
-            full_chain.push(block.clone())?;
-        }
 
         for _ in 0..finalized_count {
             let _finalized = full_chain.pop_root();
         }
 
         prop_assert_eq!(full_chain.blocks.len(), partial_chain.blocks.len());
+        prop_assert_eq!(&full_chain, &partial_chain);
     });
 
     Ok(())

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

@@ -1,6 +1,11 @@
 use std::sync::Arc;
 
-use zebra_chain::{block::Block, parameters::Network, serialization::ZcashDeserializeInto};
+use zebra_chain::{
+    block::Block,
+    history_tree::{HistoryTree, HistoryTreeError},
+    parameters::Network,
+    serialization::ZcashDeserializeInto,
+};
 use zebra_test::prelude::*;
 
 use crate::{
@@ -10,20 +15,32 @@ use crate::{
 
 use self::assert_eq;
 
-#[test]
+/// Make a history tree for the given block givens the history tree of its parent.
-fn construct_empty() {
+fn make_tree(
-    zebra_test::init();
+    block: Arc<Block>,
-    let _chain = Chain::default();
+    parent_tree: &HistoryTree,
+) -> Result<HistoryTree, HistoryTreeError> {
+    let mut tree = parent_tree.clone();
+    tree.push(block, &Default::default(), None)?;
+    Ok(tree)
 }
 
 #[test]
 fn construct_single() -> Result<()> {
     zebra_test::init();
-    let block: Arc<Block> =
+    let block0: Arc<Block> =
         zebra_test::vectors::BLOCK_MAINNET_434873_BYTES.zcash_deserialize_into()?;
 
-    let mut chain = Chain::default();
+    let finalized_tree =
-    chain.push(block.prepare())?;
+        HistoryTree::from_block(Network::Mainnet, block0.clone(), &Default::default(), None)
+            .unwrap();
+
+    let block1 = block0
+        .make_fake_child()
+        .set_block_commitment(finalized_tree.hash().into());
+
+    let mut chain = Chain::new(finalized_tree);
+    chain.push(block1.prepare())?;
 
     assert_eq!(1, chain.blocks.len());
 
@@ -36,15 +53,22 @@ fn construct_many() -> Result<()> {
 
     let mut block: Arc<Block> =
         zebra_test::vectors::BLOCK_MAINNET_434873_BYTES.zcash_deserialize_into()?;
+    let finalized_tree =
+        HistoryTree::from_block(Network::Mainnet, block.clone(), &Default::default(), None)
+            .unwrap();
     let mut blocks = vec![];
 
+    let mut tree = finalized_tree.clone();
     while blocks.len() < 100 {
-        let next_block = block.make_fake_child();
+        let next_block = block
-        blocks.push(block);
+            .make_fake_child()
+            .set_block_commitment(tree.hash().into());
+        blocks.push(next_block.clone());
         block = next_block;
+        tree = make_tree(block.clone(), &tree)?;
     }
 
-    let mut chain = Chain::default();
+    let mut chain = Chain::new(finalized_tree);
 
     for block in blocks {
         chain.push(block.prepare())?;
@@ -58,15 +82,25 @@ fn construct_many() -> Result<()> {
 #[test]
 fn ord_matches_work() -> Result<()> {
     zebra_test::init();
-    let less_block = zebra_test::vectors::BLOCK_MAINNET_434873_BYTES
+    let block =
-        .zcash_deserialize_into::<Arc<Block>>()?
+        zebra_test::vectors::BLOCK_MAINNET_434873_BYTES.zcash_deserialize_into::<Arc<Block>>()?;
-        .set_work(1);
+    let finalized_tree =
-    let more_block = less_block.clone().set_work(10);
+        HistoryTree::from_block(Network::Mainnet, block.clone(), &Default::default(), None)
+            .unwrap();
 
-    let mut lesser_chain = Chain::default();
+    let less_block = block
+        .make_fake_child()
+        .set_work(1)
+        .set_block_commitment(finalized_tree.hash().into());
+    let more_block = block
+        .make_fake_child()
+        .set_work(10)
+        .set_block_commitment(finalized_tree.hash().into());
+
+    let mut lesser_chain = Chain::new(finalized_tree.clone());
     lesser_chain.push(less_block.prepare())?;
 
-    let mut bigger_chain = Chain::default();
+    let mut bigger_chain = Chain::new(finalized_tree);
     bigger_chain.push(more_block.prepare())?;
 
     assert!(bigger_chain > lesser_chain);
@@ -93,15 +127,23 @@ fn best_chain_wins_for_network(network: Network) -> Result<()> {
             zebra_test::vectors::BLOCK_TESTNET_1326100_BYTES.zcash_deserialize_into()?
         }
     };
+    let finalized_tree =
+        HistoryTree::from_block(network, block1.clone(), &Default::default(), None).unwrap();
 
-    let block2 = block1.make_fake_child().set_work(10);
+    let block2 = block1
-    let child = block1.make_fake_child().set_work(1);
+        .make_fake_child()
+        .set_work(10)
+        .set_block_commitment(finalized_tree.hash().into());
+    let child = block1
+        .make_fake_child()
+        .set_work(1)
+        .set_block_commitment(finalized_tree.hash().into());
 
     let expected_hash = block2.hash();
 
     let mut state = NonFinalizedState::default();
-    state.commit_new_chain(block2.prepare())?;
+    state.commit_new_chain(block2.prepare(), finalized_tree.clone())?;
-    state.commit_new_chain(child.prepare())?;
+    state.commit_new_chain(child.prepare(), finalized_tree)?;
 
     let best_chain = state.best_chain().unwrap();
     assert!(best_chain.height_by_hash.contains_key(&expected_hash));
@@ -120,7 +162,7 @@ fn finalize_pops_from_best_chain() -> Result<()> {
 }
 
 fn finalize_pops_from_best_chain_for_network(network: Network) -> Result<()> {
-    let block1: Arc<Block> = match network {
+    let block0: Arc<Block> = match network {
         Network::Mainnet => {
             zebra_test::vectors::BLOCK_MAINNET_1180900_BYTES.zcash_deserialize_into()?
         }
@@ -128,14 +170,27 @@ fn finalize_pops_from_best_chain_for_network(network: Network) -> Result<()> {
             zebra_test::vectors::BLOCK_TESTNET_1326100_BYTES.zcash_deserialize_into()?
         }
     };
+    let finalized_tree =
+        HistoryTree::from_block(network, block0.clone(), &Default::default(), None).unwrap();
 
-    let block2 = block1.make_fake_child().set_work(10);
+    let block1 = block0
-    let child = block1.make_fake_child().set_work(1);
+        .make_fake_child()
+        .set_block_commitment(finalized_tree.hash().into());
+    let block1_tree = make_tree(block1.clone(), &finalized_tree)?;
+
+    let block2 = block1
+        .make_fake_child()
+        .set_work(10)
+        .set_block_commitment(block1_tree.hash().into());
+    let child = block1
+        .make_fake_child()
+        .set_work(1)
+        .set_block_commitment(block1_tree.hash().into());
 
     let mut state = NonFinalizedState::default();
-    state.commit_new_chain(block1.clone().prepare())?;
+    state.commit_new_chain(block1.clone().prepare(), finalized_tree.clone())?;
-    state.commit_block(block2.clone().prepare())?;
+    state.commit_block(block2.clone().prepare(), &finalized_tree)?;
-    state.commit_block(child.prepare())?;
+    state.commit_block(child.prepare(), &finalized_tree)?;
 
     let finalized = state.finalize();
     assert_eq!(block1, finalized.block);
@@ -162,7 +217,7 @@ fn commit_block_extending_best_chain_doesnt_drop_worst_chains() -> Result<()> {
 fn commit_block_extending_best_chain_doesnt_drop_worst_chains_for_network(
     network: Network,
 ) -> Result<()> {
-    let block1: Arc<Block> = match network {
+    let block0: Arc<Block> = match network {
         Network::Mainnet => {
             zebra_test::vectors::BLOCK_MAINNET_1180900_BYTES.zcash_deserialize_into()?
         }
@@ -170,20 +225,37 @@ fn commit_block_extending_best_chain_doesnt_drop_worst_chains_for_network(
             zebra_test::vectors::BLOCK_TESTNET_1326100_BYTES.zcash_deserialize_into()?
         }
     };
+    let finalized_tree =
+        HistoryTree::from_block(network, block0.clone(), &Default::default(), None).unwrap();
 
-    let block2 = block1.make_fake_child().set_work(10);
+    let block1 = block0
-    let child1 = block1.make_fake_child().set_work(1);
+        .make_fake_child()
-    let child2 = block2.make_fake_child().set_work(1);
+        .set_block_commitment(finalized_tree.hash().into());
+    let block1_tree = make_tree(block1.clone(), &finalized_tree)?;
+
+    let block2 = block1
+        .make_fake_child()
+        .set_work(10)
+        .set_block_commitment(block1_tree.hash().into());
+    let block2_tree = make_tree(block2.clone(), &block1_tree)?;
+    let child1 = block1
+        .make_fake_child()
+        .set_work(1)
+        .set_block_commitment(block1_tree.hash().into());
+    let child2 = block2
+        .make_fake_child()
+        .set_work(1)
+        .set_block_commitment(block2_tree.hash().into());
 
     let mut state = NonFinalizedState::default();
     assert_eq!(0, state.chain_set.len());
-    state.commit_new_chain(block1.prepare())?;
+    state.commit_new_chain(block1.prepare(), finalized_tree.clone())?;
     assert_eq!(1, state.chain_set.len());
-    state.commit_block(block2.prepare())?;
+    state.commit_block(block2.prepare(), &finalized_tree)?;
     assert_eq!(1, state.chain_set.len());
-    state.commit_block(child1.prepare())?;
+    state.commit_block(child1.prepare(), &finalized_tree)?;
     assert_eq!(2, state.chain_set.len());
-    state.commit_block(child2.prepare())?;
+    state.commit_block(child2.prepare(), &finalized_tree)?;
     assert_eq!(2, state.chain_set.len());
 
     Ok(())
@@ -200,7 +272,7 @@ 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 {
+    let block0: Arc<Block> = match network {
         Network::Mainnet => {
             zebra_test::vectors::BLOCK_MAINNET_1180900_BYTES.zcash_deserialize_into()?
         }
@@ -208,17 +280,34 @@ fn shorter_chain_can_be_best_chain_for_network(network: Network) -> Result<()> {
             zebra_test::vectors::BLOCK_TESTNET_1326100_BYTES.zcash_deserialize_into()?
         }
     };
+    let finalized_tree =
+        HistoryTree::from_block(network, block0.clone(), &Default::default(), None).unwrap();
 
-    let long_chain_block1 = block1.make_fake_child().set_work(1);
+    let block1 = block0
-    let long_chain_block2 = long_chain_block1.make_fake_child().set_work(1);
+        .make_fake_child()
+        .set_block_commitment(finalized_tree.hash().into());
+    let block1_tree = make_tree(block1.clone(), &finalized_tree)?;
 
-    let short_chain_block = block1.make_fake_child().set_work(3);
+    let long_chain_block1 = block1
+        .make_fake_child()
+        .set_work(1)
+        .set_block_commitment(block1_tree.hash().into());
+    let long_chain_block1_tree = make_tree(long_chain_block1.clone(), &block1_tree)?;
+    let long_chain_block2 = long_chain_block1
+        .make_fake_child()
+        .set_work(1)
+        .set_block_commitment(long_chain_block1_tree.hash().into());
+
+    let short_chain_block = block1
+        .make_fake_child()
+        .set_work(3)
+        .set_block_commitment(block1_tree.hash().into());
 
     let mut state = NonFinalizedState::default();
-    state.commit_new_chain(block1.prepare())?;
+    state.commit_new_chain(block1.prepare(), finalized_tree.clone())?;
-    state.commit_block(long_chain_block1.prepare())?;
+    state.commit_block(long_chain_block1.prepare(), &finalized_tree)?;
-    state.commit_block(long_chain_block2.prepare())?;
+    state.commit_block(long_chain_block2.prepare(), &finalized_tree)?;
-    state.commit_block(short_chain_block.prepare())?;
+    state.commit_block(short_chain_block.prepare(), &finalized_tree)?;
     assert_eq!(2, state.chain_set.len());
 
     assert_eq!(2, state.best_chain_len());
@@ -237,7 +326,7 @@ 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 {
+    let block0: Arc<Block> = match network {
         Network::Mainnet => {
             zebra_test::vectors::BLOCK_MAINNET_1180900_BYTES.zcash_deserialize_into()?
         }
@@ -245,21 +334,46 @@ fn longer_chain_with_more_work_wins_for_network(network: Network) -> Result<()>
             zebra_test::vectors::BLOCK_TESTNET_1326100_BYTES.zcash_deserialize_into()?
         }
     };
+    let finalized_tree =
+        HistoryTree::from_block(network, block0.clone(), &Default::default(), None).unwrap();
 
-    let long_chain_block1 = block1.make_fake_child().set_work(1);
+    let block1 = block0
-    let long_chain_block2 = long_chain_block1.make_fake_child().set_work(1);
+        .make_fake_child()
-    let long_chain_block3 = long_chain_block2.make_fake_child().set_work(1);
+        .set_block_commitment(finalized_tree.hash().into());
-    let long_chain_block4 = long_chain_block3.make_fake_child().set_work(1);
+    let block1_tree = make_tree(block1.clone(), &finalized_tree)?;
 
-    let short_chain_block = block1.make_fake_child().set_work(3);
+    let long_chain_block1 = block1
+        .make_fake_child()
+        .set_work(1)
+        .set_block_commitment(block1_tree.hash().into());
+    let long_chain_block1_tree = make_tree(long_chain_block1.clone(), &block1_tree)?;
+    let long_chain_block2 = long_chain_block1
+        .make_fake_child()
+        .set_work(1)
+        .set_block_commitment(long_chain_block1_tree.hash().into());
+    let long_chain_block2_tree = make_tree(long_chain_block2.clone(), &long_chain_block1_tree)?;
+    let long_chain_block3 = long_chain_block2
+        .make_fake_child()
+        .set_work(1)
+        .set_block_commitment(long_chain_block2_tree.hash().into());
+    let long_chain_block3_tree = make_tree(long_chain_block3.clone(), &long_chain_block2_tree)?;
+    let long_chain_block4 = long_chain_block3
+        .make_fake_child()
+        .set_work(1)
+        .set_block_commitment(long_chain_block3_tree.hash().into());
+
+    let short_chain_block = block1
+        .make_fake_child()
+        .set_work(3)
+        .set_block_commitment(block1_tree.hash().into());
 
     let mut state = NonFinalizedState::default();
-    state.commit_new_chain(block1.prepare())?;
+    state.commit_new_chain(block1.prepare(), finalized_tree.clone())?;
-    state.commit_block(long_chain_block1.prepare())?;
+    state.commit_block(long_chain_block1.prepare(), &finalized_tree)?;
-    state.commit_block(long_chain_block2.prepare())?;
+    state.commit_block(long_chain_block2.prepare(), &finalized_tree)?;
-    state.commit_block(long_chain_block3.prepare())?;
+    state.commit_block(long_chain_block3.prepare(), &finalized_tree)?;
-    state.commit_block(long_chain_block4.prepare())?;
+    state.commit_block(long_chain_block4.prepare(), &finalized_tree)?;
-    state.commit_block(short_chain_block.prepare())?;
+    state.commit_block(short_chain_block.prepare(), &finalized_tree)?;
     assert_eq!(2, state.chain_set.len());
 
     assert_eq!(5, state.best_chain_len());
@@ -277,7 +391,7 @@ fn equal_length_goes_to_more_work() -> Result<()> {
     Ok(())
 }
 fn equal_length_goes_to_more_work_for_network(network: Network) -> Result<()> {
-    let block1: Arc<Block> = match network {
+    let block0: Arc<Block> = match network {
         Network::Mainnet => {
             zebra_test::vectors::BLOCK_MAINNET_1180900_BYTES.zcash_deserialize_into()?
         }
@@ -285,15 +399,28 @@ fn equal_length_goes_to_more_work_for_network(network: Network) -> Result<()> {
             zebra_test::vectors::BLOCK_TESTNET_1326100_BYTES.zcash_deserialize_into()?
         }
     };
+    let finalized_tree =
+        HistoryTree::from_block(network, block0.clone(), &Default::default(), None).unwrap();
 
-    let less_work_child = block1.make_fake_child().set_work(1);
+    let block1 = block0
-    let more_work_child = block1.make_fake_child().set_work(3);
+        .make_fake_child()
+        .set_block_commitment(finalized_tree.hash().into());
+    let block1_tree = make_tree(block1.clone(), &finalized_tree)?;
+
+    let less_work_child = block1
+        .make_fake_child()
+        .set_work(1)
+        .set_block_commitment(block1_tree.hash().into());
+    let more_work_child = block1
+        .make_fake_child()
+        .set_work(3)
+        .set_block_commitment(block1_tree.hash().into());
     let expected_hash = more_work_child.hash();
 
     let mut state = NonFinalizedState::default();
-    state.commit_new_chain(block1.prepare())?;
+    state.commit_new_chain(block1.prepare(), finalized_tree.clone())?;
-    state.commit_block(less_work_child.prepare())?;
+    state.commit_block(less_work_child.prepare(), &finalized_tree)?;
-    state.commit_block(more_work_child.prepare())?;
+    state.commit_block(more_work_child.prepare(), &finalized_tree)?;
     assert_eq!(2, state.chain_set.len());
 
     let tip_hash = state.best_tip().unwrap().1;

zebra-state/src/tests.rs

@@ -38,6 +38,8 @@ pub trait FakeChainHelper {
     fn make_fake_child(&self) -> Arc<Block>;
 
     fn set_work(self, work: u128) -> Arc<Block>;
+
+    fn set_block_commitment(self, commitment: [u8; 32]) -> Arc<Block>;
 }
 
 impl FakeChainHelper for Arc<Block> {
@@ -74,6 +76,12 @@ impl FakeChainHelper for Arc<Block> {
         block.header.difficulty_threshold = expanded.into();
         self
     }
+
+    fn set_block_commitment(mut self, block_commitment: [u8; 32]) -> Arc<Block> {
+        let block = Arc::make_mut(&mut self);
+        block.header.commitment_bytes = block_commitment;
+        self
+    }
 }
 
 fn work_to_expanded(work: U256) -> ExpandedDifficulty {