feat(state): Send treestate from non-finalized state to finalized state (#4721)

* Add history trees for each height in non-fin state

* Refactor formatting

* Pass the treestate to the finalized state

I created a new structure `FinalizedBlockWithTrees` that wraps the
treestate and the finalized block. I did that because the original
`FinalizedBlock` is `Eq`, but `HistoryTree` can't be `Eq`.

This makes Zebra faster because:

1. The finalized state doesn't retrieve the treestate from the disk if
the non-finalized state supplies it.

2.The finalized state doesn't recompute the treestate if the
non-finalized state supplies it.

* Check block commitment before updating hist tree

* Store Sprout commitment trees in non-fin state

* Send trees for the root block to fin-state

When committing a block and sending the treestate from the non-finalized
state to the finalized state, Zebra was sending trees that correspond to
the tip block instead of trees that correspond to the root block of the
best chain. This commit fixes that.

* Refactor doc comments

* Refactor block finalization

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

zebra-state/src/request.rs

@@ -9,8 +9,12 @@ use std::{
 use zebra_chain::{
     amount::NegativeAllowed,
     block::{self, Block},
+    history_tree::HistoryTree,
+    orchard,
+    parallel::tree::NoteCommitmentTrees,
+    sapling,
     serialization::SerializationError,
-    transaction,
+    sprout, transaction,
     transparent::{self, utxos_from_ordered_utxos},
     value_balance::{ValueBalance, ValueBalanceError},
 };
@@ -177,6 +181,72 @@ pub struct FinalizedBlock {
     pub transaction_hashes: Arc<[transaction::Hash]>,
 }
 
+/// Wraps note commitment trees and the history tree together.
+pub struct Treestate {
+    /// Note commitment trees.
+    pub note_commitment_trees: NoteCommitmentTrees,
+    /// History tree.
+    pub history_tree: Arc<HistoryTree>,
+}
+
+impl Treestate {
+    pub fn new(
+        sprout: Arc<sprout::tree::NoteCommitmentTree>,
+        sapling: Arc<sapling::tree::NoteCommitmentTree>,
+        orchard: Arc<orchard::tree::NoteCommitmentTree>,
+        history_tree: Arc<HistoryTree>,
+    ) -> Self {
+        Self {
+            note_commitment_trees: NoteCommitmentTrees {
+                sprout,
+                sapling,
+                orchard,
+            },
+            history_tree,
+        }
+    }
+}
+
+/// Contains a block ready to be committed together with its associated
+/// treestate.
+///
+/// Zebra's non-finalized state passes this `struct` over to the finalized state
+/// when committing a block. The associated treestate is passed so that the
+/// finalized state does not have to retrieve the previous treestate from the
+/// database and recompute the new one.
+pub struct FinalizedWithTrees {
+    /// A block ready to be committed.
+    pub finalized: FinalizedBlock,
+    /// The tresstate associated with the block.
+    pub treestate: Option<Treestate>,
+}
+
+impl FinalizedWithTrees {
+    pub fn new(block: ContextuallyValidBlock, treestate: Treestate) -> Self {
+        let finalized = FinalizedBlock::from(block);
+
+        Self {
+            finalized,
+            treestate: Some(treestate),
+        }
+    }
+}
+
+impl From<Arc<Block>> for FinalizedWithTrees {
+    fn from(block: Arc<Block>) -> Self {
+        Self::from(FinalizedBlock::from(block))
+    }
+}
+
+impl From<FinalizedBlock> for FinalizedWithTrees {
+    fn from(block: FinalizedBlock) -> Self {
+        Self {
+            finalized: block,
+            treestate: None,
+        }
+    }
+}
+
 impl From<&PreparedBlock> for PreparedBlock {
     fn from(prepared: &PreparedBlock) -> Self {
         prepared.clone()

zebra-state/src/service.rs

@@ -296,9 +296,9 @@ impl StateService {
 
         while self.mem.best_chain_len() > crate::constants::MAX_BLOCK_REORG_HEIGHT {
             tracing::trace!("finalizing block past the reorg limit");
-            let finalized = self.mem.finalize();
+            let finalized_with_trees = self.mem.finalize();
             self.disk
-                .commit_finalized_direct(finalized, "best non-finalized chain root")
+                .commit_finalized_direct(finalized_with_trees, "best non-finalized chain root")
                 .expect(
                     "expected that errors would not occur when writing to disk or updating note commitment and history trees",
                 );

zebra-state/src/service/check/tests/nullifier.rs

@@ -82,7 +82,7 @@ proptest! {
         // randomly choose to commit the block to the finalized or non-finalized state
         if use_finalized_state {
             let block1 = FinalizedBlock::from(Arc::new(block1));
-            let commit_result = state.disk.commit_finalized_direct(block1.clone(), "test");
+            let commit_result = state.disk.commit_finalized_direct(block1.clone().into(), "test");
 
             // the block was committed
             prop_assert_eq!(Some((Height(1), block1.hash)), state.best_tip());
@@ -332,7 +332,7 @@ proptest! {
         // randomly choose to commit the next block to the finalized or non-finalized state
         if duplicate_in_finalized_state {
             let block1 = FinalizedBlock::from(Arc::new(block1));
-            let commit_result = state.disk.commit_finalized_direct(block1.clone(), "test");
+            let commit_result = state.disk.commit_finalized_direct(block1.clone().into(), "test");
 
             prop_assert_eq!(Some((Height(1), block1.hash)), state.best_tip());
             prop_assert!(commit_result.is_ok());
@@ -416,7 +416,7 @@ proptest! {
         // randomly choose to commit the block to the finalized or non-finalized state
         if use_finalized_state {
             let block1 = FinalizedBlock::from(Arc::new(block1));
-            let commit_result = state.disk.commit_finalized_direct(block1.clone(), "test");
+            let commit_result = state.disk.commit_finalized_direct(block1.clone().into(), "test");
 
             prop_assert_eq!(Some((Height(1), block1.hash)), state.best_tip());
             prop_assert!(commit_result.is_ok());
@@ -582,7 +582,7 @@ proptest! {
         // randomly choose to commit the next block to the finalized or non-finalized state
         if duplicate_in_finalized_state {
             let block1 = FinalizedBlock::from(Arc::new(block1));
-            let commit_result = state.disk.commit_finalized_direct(block1.clone(), "test");
+            let commit_result = state.disk.commit_finalized_direct(block1.clone().into(), "test");
 
             prop_assert_eq!(Some((Height(1), block1.hash)), state.best_tip());
             prop_assert!(commit_result.is_ok());
@@ -660,7 +660,7 @@ proptest! {
         // randomly choose to commit the block to the finalized or non-finalized state
         if use_finalized_state {
             let block1 = FinalizedBlock::from(Arc::new(block1));
-            let commit_result = state.disk.commit_finalized_direct(block1.clone(), "test");
+            let commit_result = state.disk.commit_finalized_direct(block1.clone().into(), "test");
 
             prop_assert_eq!(Some((Height(1), block1.hash)), state.best_tip());
             prop_assert!(commit_result.is_ok());
@@ -834,7 +834,7 @@ proptest! {
         // randomly choose to commit the next block to the finalized or non-finalized state
         if duplicate_in_finalized_state {
             let block1 = FinalizedBlock::from(Arc::new(block1));
-            let commit_result = state.disk.commit_finalized_direct(block1.clone(), "test");
+            let commit_result = state.disk.commit_finalized_direct(block1.clone().into(), "test");
 
             prop_assert_eq!(Some((Height(1), block1.hash)), state.best_tip());
             prop_assert!(commit_result.is_ok());

zebra-state/src/service/check/tests/utxo.rs

@@ -176,7 +176,7 @@ proptest! {
         // randomly choose to commit the block to the finalized or non-finalized state
         if use_finalized_state {
             let block1 = FinalizedBlock::from(Arc::new(block1));
-            let commit_result = state.disk.commit_finalized_direct(block1.clone(), "test");
+            let commit_result = state.disk.commit_finalized_direct(block1.clone().into(), "test");
 
             // the block was committed
             prop_assert_eq!(Some((Height(1), block1.hash)), state.best_tip());
@@ -262,7 +262,7 @@ proptest! {
 
         if use_finalized_state_spend {
             let block2 = FinalizedBlock::from(Arc::new(block2));
-            let commit_result = state.disk.commit_finalized_direct(block2.clone(), "test");
+            let commit_result = state.disk.commit_finalized_direct(block2.clone().into(), "test");
 
             // the block was committed
             prop_assert_eq!(Some((Height(2), block2.hash)), state.best_tip());
@@ -591,7 +591,7 @@ proptest! {
 
         if use_finalized_state_spend {
             let block2 = FinalizedBlock::from(block2.clone());
-            let commit_result = state.disk.commit_finalized_direct(block2.clone(), "test");
+            let commit_result = state.disk.commit_finalized_direct(block2.clone().into(), "test");
 
             // the block was committed
             prop_assert_eq!(Some((Height(2), block2.hash)), state.best_tip());
@@ -846,7 +846,9 @@ fn new_state_with_mainnet_transparent_data(
 
     if use_finalized_state {
         let block1 = FinalizedBlock::from(block1.clone());
-        let commit_result = state.disk.commit_finalized_direct(block1.clone(), "test");
+        let commit_result = state
+            .disk
+            .commit_finalized_direct(block1.clone().into(), "test");
 
         // the block was committed
         assert_eq!(Some((Height(1), block1.hash)), state.best_tip());

zebra-state/src/service/finalized_state.rs

@@ -19,11 +19,13 @@ use std::{
     collections::HashMap,
     io::{stderr, stdout, Write},
     path::Path,
+    sync::Arc,
 };
 
 use zebra_chain::{block, parameters::Network};
 
 use crate::{
+    request::FinalizedWithTrees,
     service::{check, QueuedFinalized},
     BoxError, Config, FinalizedBlock,
 };
@@ -188,7 +190,8 @@ impl FinalizedState {
     /// public API of the [`FinalizedState`].
     fn commit_finalized(&mut self, queued_block: QueuedFinalized) -> Result<FinalizedBlock, ()> {
         let (finalized, rsp_tx) = queued_block;
-        let result = self.commit_finalized_direct(finalized.clone(), "CommitFinalized request");
+        let result =
+            self.commit_finalized_direct(finalized.clone().into(), "CommitFinalized request");
 
         let block_result = if result.is_ok() {
             metrics::counter!("state.checkpoint.finalized.block.count", 1);
@@ -238,9 +241,10 @@ impl FinalizedState {
     #[allow(clippy::unwrap_in_result)]
     pub fn commit_finalized_direct(
         &mut self,
-        finalized: FinalizedBlock,
+        finalized_with_trees: FinalizedWithTrees,
         source: &str,
     ) -> Result<block::Hash, BoxError> {
+        let finalized = finalized_with_trees.finalized;
         let committed_tip_hash = self.db.finalized_tip_hash();
         let committed_tip_height = self.db.finalized_tip_height();
 
@@ -272,28 +276,73 @@ impl FinalizedState {
             );
         }
 
-        // Check the block commitment. For Nu5-onward, the block hash commits only
+        let (history_tree, note_commitment_trees) = match finalized_with_trees.treestate {
-        // to non-authorizing data (see ZIP-244). This checks the authorizing data
+            // If the treestate associated with the block was supplied, use it
-        // commitment, making sure the entire block contents were committed to.
+            // without recomputing it.
-        // The test is done here (and not during semantic validation) because it needs
+            Some(ref treestate) => (
-        // the history tree root. While it _is_ checked during contextual validation,
+                treestate.history_tree.clone(),
-        // that is not called by the checkpoint verifier, and keeping a history tree there
+                treestate.note_commitment_trees.clone(),
-        // would be harder to implement.
+            ),
+            // If the treestate was not supplied, retrieve a previous treestate
+            // from the database, and update it for the block being committed.
+            None => {
+                let mut history_tree = self.db.history_tree();
+                let mut note_commitment_trees = self.db.note_commitment_trees();
+
+                // Update the note commitment trees.
+                note_commitment_trees.update_trees_parallel(&finalized.block)?;
+
+                // Check the block commitment if the history tree was not
+                // supplied by the non-finalized state. Note that we don't do
+                // this check for history trees supplied by the non-finalized
+                // state because the non-finalized state checks the block
+                // commitment.
                 //
-        // TODO: run this CPU-intensive cryptography in a parallel rayon thread, if it shows up in profiles
+                // For Nu5-onward, the block hash commits only to
-        let history_tree = self.db.history_tree();
+                // non-authorizing data (see ZIP-244). This checks the
+                // authorizing data commitment, making sure the entire block
+                // contents were committed to. The test is done here (and not
+                // during semantic validation) because it needs the history tree
+                // root. While it _is_ checked during contextual validation,
+                // that is not called by the checkpoint verifier, and keeping a
+                // history tree there would be harder to implement.
+                //
+                // TODO: run this CPU-intensive cryptography in a parallel rayon
+                // thread, if it shows up in profiles
                 check::block_commitment_is_valid_for_chain_history(
                     finalized.block.clone(),
                     self.network,
                     &history_tree,
                 )?;
 
+                // Update the history tree.
+                //
+                // TODO: run this CPU-intensive cryptography in a parallel rayon
+                // thread, if it shows up in profiles
+                let history_tree_mut = Arc::make_mut(&mut history_tree);
+                let sapling_root = note_commitment_trees.sapling.root();
+                let orchard_root = note_commitment_trees.orchard.root();
+                history_tree_mut.push(
+                    self.network(),
+                    finalized.block.clone(),
+                    sapling_root,
+                    orchard_root,
+                )?;
+
+                (history_tree, note_commitment_trees)
+            }
+        };
+
         let finalized_height = finalized.height;
         let finalized_hash = finalized.hash;
 
-        let result = self
+        let result = self.db.write_block(
-            .db
+            finalized,
-            .write_block(finalized, history_tree, self.network, source);
+            history_tree,
+            note_commitment_trees,
+            self.network,
+            source,
+        );
 
         // TODO: move the stop height check to the syncer (#3442)
         if result.is_ok() && self.is_at_stop_height(finalized_height) {

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

@@ -28,8 +28,9 @@ fn blocks_with_v5_transactions() -> Result<()> {
             let mut height = Height(0);
             // use `count` to minimize test failures, so they are easier to diagnose
             for block in chain.iter().take(count) {
+                let finalized = FinalizedBlock::from(block.block.clone());
                 let hash = state.commit_finalized_direct(
-                    FinalizedBlock::from(block.block.clone()),
+                    finalized.into(),
                     "blocks_with_v5_transactions test"
                 );
                 prop_assert_eq!(Some(height), state.finalized_tip_height());
@@ -83,16 +84,18 @@ fn all_upgrades_and_wrong_commitments_with_fake_activation_heights() -> Result<(
                         h == nu5_height ||
                         h == nu5_height_plus1 => {
                             let block = block.block.clone().set_block_commitment([0x42; 32]);
+                            let finalized = FinalizedBlock::from(block);
                             state.commit_finalized_direct(
-                                FinalizedBlock::from(block),
+                                finalized.into(),
                                 "all_upgrades test"
                             ).expect_err("Must fail commitment check");
                             failure_count += 1;
                         },
                     _ => {},
                 }
+                let finalized = FinalizedBlock::from(block.block.clone());
                 let hash = state.commit_finalized_direct(
-                    FinalizedBlock::from(block.block.clone()),
+                    finalized.into(),
                     "all_upgrades test"
                 ).unwrap();
                 prop_assert_eq!(Some(height), state.finalized_tip_height());

zebra-state/src/service/finalized_state/zebra_db/block.rs

@@ -242,6 +242,7 @@ impl ZebraDb {
         &mut self,
         finalized: FinalizedBlock,
         history_tree: Arc<HistoryTree>,
+        note_commitment_trees: NoteCommitmentTrees,
         network: Network,
         source: &str,
     ) -> Result<block::Hash, BoxError> {
@@ -329,8 +330,8 @@ impl ZebraDb {
             spent_utxos_by_outpoint,
             spent_utxos_by_out_loc,
             address_balances,
-            self.note_commitment_trees(),
             history_tree,
+            note_commitment_trees,
             self.finalized_value_pool(),
         )?;
 
@@ -382,8 +383,8 @@ impl DiskWriteBatch {
         spent_utxos_by_outpoint: HashMap<transparent::OutPoint, transparent::Utxo>,
         spent_utxos_by_out_loc: BTreeMap<OutputLocation, transparent::Utxo>,
         address_balances: HashMap<transparent::Address, AddressBalanceLocation>,
-        mut note_commitment_trees: NoteCommitmentTrees,
         history_tree: Arc<HistoryTree>,
+        note_commitment_trees: NoteCommitmentTrees,
         value_pool: ValueBalance<NonNegative>,
     ) -> Result<(), BoxError> {
         let FinalizedBlock {
@@ -419,7 +420,7 @@ impl DiskWriteBatch {
             &spent_utxos_by_out_loc,
             address_balances,
         )?;
-        self.prepare_shielded_transaction_batch(db, &finalized, &mut note_commitment_trees)?;
+        self.prepare_shielded_transaction_batch(db, &finalized)?;
 
         self.prepare_note_commitment_batch(db, &finalized, note_commitment_trees, history_tree)?;
 

zebra-state/src/service/finalized_state/zebra_db/chain.rs

@@ -16,7 +16,7 @@ use std::{borrow::Borrow, collections::HashMap, sync::Arc};
 use zebra_chain::{
     amount::NonNegative,
     history_tree::{HistoryTree, NonEmptyHistoryTree},
-    orchard, sapling, transparent,
+    transparent,
     value_balance::ValueBalance,
 };
 
@@ -71,17 +71,11 @@ impl DiskWriteBatch {
         &mut self,
         db: &DiskDb,
         finalized: &FinalizedBlock,
-        sapling_root: sapling::tree::Root,
+        history_tree: Arc<HistoryTree>,
-        orchard_root: orchard::tree::Root,
-        mut history_tree: Arc<HistoryTree>,
     ) -> Result<(), BoxError> {
         let history_tree_cf = db.cf_handle("history_tree").unwrap();
 
-        let FinalizedBlock { block, height, .. } = finalized;
+        let FinalizedBlock { height, .. } = finalized;
-
-        // TODO: run this CPU-intensive cryptography in a parallel rayon thread, if it shows up in profiles
-        let history_tree_mut = Arc::make_mut(&mut history_tree);
-        history_tree_mut.push(self.network(), block.clone(), sapling_root, orchard_root)?;
 
         // Update the tree in state
         let current_tip_height = *height - 1;

zebra-state/src/service/finalized_state/zebra_db/shielded.rs

@@ -180,7 +180,6 @@ impl DiskWriteBatch {
         &mut self,
         db: &DiskDb,
         finalized: &FinalizedBlock,
-        note_commitment_trees: &mut NoteCommitmentTrees,
     ) -> Result<(), BoxError> {
         let FinalizedBlock { block, .. } = finalized;
 
@@ -189,8 +188,6 @@ impl DiskWriteBatch {
             self.prepare_nullifier_batch(db, transaction)?;
         }
 
-        note_commitment_trees.update_trees_parallel(block)?;
-
         Ok(())
     }
 
@@ -290,7 +287,7 @@ impl DiskWriteBatch {
             note_commitment_trees.orchard,
         );
 
-        self.prepare_history_batch(db, finalized, sapling_root, orchard_root, history_tree)
+        self.prepare_history_batch(db, finalized, history_tree)
     }
 
     /// Prepare a database batch containing the initial note commitment trees,

zebra-state/src/service/non_finalized_state.rs

@@ -17,9 +17,9 @@ use zebra_chain::{
 };
 
 use crate::{
-    request::ContextuallyValidBlock,
+    request::{ContextuallyValidBlock, FinalizedWithTrees},
     service::{check, finalized_state::ZebraDb},
-    FinalizedBlock, PreparedBlock, ValidateContextError,
+    PreparedBlock, ValidateContextError,
 };
 
 mod chain;
@@ -80,7 +80,7 @@ impl NonFinalizedState {
 
     /// Finalize the lowest height block in the non-finalized portion of the best
     /// chain and update all side-chains to match.
-    pub fn finalize(&mut self) -> FinalizedBlock {
+    pub fn finalize(&mut self) -> FinalizedWithTrees {
         // Chain::cmp uses the partial cumulative work, and the hash of the tip block.
         // Neither of these fields has interior mutability.
         // (And when the tip block is dropped for a chain, the chain is also dropped.)
@@ -90,14 +90,16 @@ impl NonFinalizedState {
 
         // extract best chain
         let mut best_chain = chains.next_back().expect("there's at least one chain");
+
         // clone if required
-        let write_best_chain = Arc::make_mut(&mut best_chain);
+        let mut_best_chain = Arc::make_mut(&mut best_chain);
 
         // extract the rest into side_chains so they can be mutated
         let side_chains = chains;
 
-        // remove the lowest height block from the best_chain to be finalized
+        // Pop the lowest height block from the best chain to be finalized, and
-        let finalizing = write_best_chain.pop_root();
+        // also obtain its associated treestate.
+        let (best_chain_root, root_treestate) = mut_best_chain.pop_root();
 
         // add best_chain back to `self.chain_set`
         if !best_chain.is_empty() {
@@ -105,11 +107,11 @@ impl NonFinalizedState {
         }
 
         // for each remaining chain in side_chains
-        for mut chain in side_chains {
+        for mut side_chain in side_chains {
-            if chain.non_finalized_root_hash() != finalizing.hash {
+            if side_chain.non_finalized_root_hash() != best_chain_root.hash {
                 // If we popped the root, the chain would be empty or orphaned,
                 // so just drop it now.
-                drop(chain);
+                drop(side_chain);
 
                 continue;
             }
@@ -117,19 +119,20 @@ impl NonFinalizedState {
             // otherwise, the popped root block is the same as the finalizing block
 
             // clone if required
-            let write_chain = Arc::make_mut(&mut chain);
+            let mut_side_chain = Arc::make_mut(&mut side_chain);
 
             // remove the first block from `chain`
-            let chain_start = write_chain.pop_root();
+            let (side_chain_root, _treestate) = mut_side_chain.pop_root();
-            assert_eq!(chain_start.hash, finalizing.hash);
+            assert_eq!(side_chain_root.hash, best_chain_root.hash);
 
             // add the chain back to `self.chain_set`
-            self.chain_set.insert(chain);
+            self.chain_set.insert(side_chain);
         }
 
         self.update_metrics_for_chains();
 
-        finalizing.into()
+        // Add the treestate to the finalized block.
+        FinalizedWithTrees::new(best_chain_root, root_treestate)
     }
 
     /// Commit block to the non-finalized state, on top of:

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

@@ -30,8 +30,8 @@ use zebra_chain::{
 };
 
 use crate::{
-    service::check, ContextuallyValidBlock, HashOrHeight, OutputLocation, TransactionLocation,
+    request::Treestate, service::check, ContextuallyValidBlock, HashOrHeight, OutputLocation,
-    ValidateContextError,
+    TransactionLocation, ValidateContextError,
 };
 
 use self::index::TransparentTransfers;
@@ -71,6 +71,9 @@ pub struct Chain {
     /// This is required for interstitial states.
     pub(crate) sprout_trees_by_anchor:
         HashMap<sprout::tree::Root, Arc<sprout::tree::NoteCommitmentTree>>,
+    /// The Sprout note commitment tree for each height.
+    pub(crate) sprout_trees_by_height:
+        BTreeMap<block::Height, Arc<sprout::tree::NoteCommitmentTree>>,
     /// The Sapling note commitment tree of the tip of this [`Chain`],
     /// including all finalized notes, and the non-finalized notes in this chain.
     pub(super) sapling_note_commitment_tree: Arc<sapling::tree::NoteCommitmentTree>,
@@ -150,6 +153,7 @@ impl Chain {
             sprout_anchors: MultiSet::new(),
             sprout_anchors_by_height: Default::default(),
             sprout_trees_by_anchor: Default::default(),
+            sprout_trees_by_height: Default::default(),
             sapling_anchors: MultiSet::new(),
             sapling_anchors_by_height: Default::default(),
             sapling_trees_by_height: Default::default(),
@@ -191,6 +195,7 @@ impl Chain {
             // note commitment trees
             self.sprout_note_commitment_tree.root() == other.sprout_note_commitment_tree.root() &&
             self.sprout_trees_by_anchor == other.sprout_trees_by_anchor &&
+            self.sprout_trees_by_height == other.sprout_trees_by_height &&
             self.sapling_note_commitment_tree.root() == other.sapling_note_commitment_tree.root() &&
             self.sapling_trees_by_height == other.sapling_trees_by_height &&
             self.orchard_note_commitment_tree.root() == other.orchard_note_commitment_tree.root() &&
@@ -240,22 +245,28 @@ impl Chain {
         Ok(self)
     }
 
-    /// Remove the lowest height block of the non-finalized portion of a chain.
+    /// Pops the lowest height block of the non-finalized portion of a chain,
+    /// and returns it with its associated treestate.
     #[instrument(level = "debug", skip(self))]
-    pub(crate) fn pop_root(&mut self) -> ContextuallyValidBlock {
+    pub(crate) fn pop_root(&mut self) -> (ContextuallyValidBlock, Treestate) {
+        // Obtain the lowest height.
         let block_height = self.non_finalized_root_height();
 
-        // remove the lowest height block from self.blocks
+        // Obtain the treestate associated with the block being finalized.
+        let treestate = self
+            .treestate(block_height.into())
+            .expect("The treestate must be present for the root height.");
+
+        // Remove the lowest height block from `self.blocks`.
         let block = self
             .blocks
             .remove(&block_height)
             .expect("only called while blocks is populated");
 
-        // update cumulative data members
+        // Update cumulative data members.
         self.revert_chain_with(&block, RevertPosition::Root);
 
-        // return the prepared block
+        (block, treestate)
-        block
     }
 
     /// Returns the height of the chain root.
@@ -481,9 +492,22 @@ impl Chain {
         )
     }
 
+    /// Returns the Sprout
+    /// [`NoteCommitmentTree`](sprout::tree::NoteCommitmentTree) specified by a
+    /// [`HashOrHeight`], if it exists in the non-finalized [`Chain`].
+    pub fn sprout_tree(
+        &self,
+        hash_or_height: HashOrHeight,
+    ) -> Option<Arc<sprout::tree::NoteCommitmentTree>> {
+        let height =
+            hash_or_height.height_or_else(|hash| self.height_by_hash.get(&hash).cloned())?;
+
+        self.sprout_trees_by_height.get(&height).cloned()
+    }
+
     /// Returns the Sapling
     /// [`NoteCommitmentTree`](sapling::tree::NoteCommitmentTree) specified by a
-    /// hash or height, if it exists in the non-finalized `chain`.
+    /// [`HashOrHeight`], if it exists in the non-finalized [`Chain`].
     pub fn sapling_tree(
         &self,
         hash_or_height: HashOrHeight,
@@ -496,7 +520,7 @@ impl Chain {
 
     /// Returns the Orchard
     /// [`NoteCommitmentTree`](orchard::tree::NoteCommitmentTree) specified by a
-    /// hash or height, if it exists in the non-finalized `chain`.
+    /// [`HashOrHeight`], if it exists in the non-finalized [`Chain`].
     pub fn orchard_tree(
         &self,
         hash_or_height: HashOrHeight,
@@ -507,6 +531,29 @@ impl Chain {
         self.orchard_trees_by_height.get(&height).cloned()
     }
 
+    /// Returns the [`HistoryTree`] specified by a [`HashOrHeight`], if it
+    /// exists in the non-finalized [`Chain`].
+    pub fn history_tree(&self, hash_or_height: HashOrHeight) -> Option<Arc<HistoryTree>> {
+        let height =
+            hash_or_height.height_or_else(|hash| self.height_by_hash.get(&hash).cloned())?;
+
+        self.history_trees_by_height.get(&height).cloned()
+    }
+
+    fn treestate(&self, hash_or_height: HashOrHeight) -> Option<Treestate> {
+        let sprout_tree = self.sprout_tree(hash_or_height)?;
+        let sapling_tree = self.sapling_tree(hash_or_height)?;
+        let orchard_tree = self.orchard_tree(hash_or_height)?;
+        let history_tree = self.history_tree(hash_or_height)?;
+
+        Some(Treestate::new(
+            sprout_tree,
+            sapling_tree,
+            orchard_tree,
+            history_tree,
+        ))
+    }
+
     /// Returns the block hash of the tip block.
     pub fn non_finalized_tip_hash(&self) -> block::Hash {
         self.blocks
@@ -739,6 +786,7 @@ impl Chain {
             spent_utxos: self.spent_utxos.clone(),
             sprout_note_commitment_tree,
             sprout_trees_by_anchor: self.sprout_trees_by_anchor.clone(),
+            sprout_trees_by_height: self.sprout_trees_by_height.clone(),
             sapling_note_commitment_tree,
             sapling_trees_by_height: self.sapling_trees_by_height.clone(),
             orchard_note_commitment_tree,
@@ -808,6 +856,8 @@ impl Chain {
         // Do the Chain updates with data dependencies on note commitment tree updates
 
         // Update the note commitment trees indexed by height.
+        self.sprout_trees_by_height
+            .insert(height, self.sprout_note_commitment_tree.clone());
         self.sapling_trees_by_height
             .insert(height, self.sapling_note_commitment_tree.clone());
         self.orchard_trees_by_height
@@ -1115,6 +1165,9 @@ impl UpdateWith<ContextuallyValidBlock> for Chain {
         if !self.sprout_anchors.contains(&anchor) {
             self.sprout_trees_by_anchor.remove(&anchor);
         }
+        self.sprout_trees_by_height
+            .remove(&height)
+            .expect("Sprout note commitment tree must be present if block was added to chain");
 
         let anchor = self
             .sapling_anchors_by_height

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

@@ -354,7 +354,7 @@ fn finalized_equals_pushed_genesis() -> Result<()> {
             }
 
         for _ in 0..finalized_count {
-            let _finalized = full_chain.pop_root();
+            full_chain.pop_root();
         }
 
         prop_assert_eq!(full_chain.blocks.len(), partial_chain.blocks.len());
@@ -425,7 +425,7 @@ fn finalized_equals_pushed_history_tree() -> Result<()> {
             }
 
         for _ in 0..finalized_count {
-            let _finalized = full_chain.pop_root();
+            full_chain.pop_root();
         }
 
         prop_assert_eq!(full_chain.blocks.len(), partial_chain.blocks.len());
@@ -608,6 +608,7 @@ fn different_blocks_different_chains() -> Result<()> {
                 // note commitment trees
                 chain1.sprout_note_commitment_tree = chain2.sprout_note_commitment_tree.clone();
                 chain1.sprout_trees_by_anchor = chain2.sprout_trees_by_anchor.clone();
+                chain1.sprout_trees_by_height = chain2.sprout_trees_by_height.clone();
                 chain1.sapling_note_commitment_tree = chain2.sapling_note_commitment_tree.clone();
                 chain1.sapling_trees_by_height = chain2.sapling_trees_by_height.clone();
                 chain1.orchard_note_commitment_tree = chain2.orchard_note_commitment_tree.clone();

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

@@ -198,10 +198,12 @@ fn finalize_pops_from_best_chain_for_network(network: Network) -> Result<()> {
     state.commit_block(block2.clone().prepare(), &finalized_state)?;
     state.commit_block(child.prepare(), &finalized_state)?;
 
-    let finalized = state.finalize();
+    let finalized_with_trees = state.finalize();
+    let finalized = finalized_with_trees.finalized;
     assert_eq!(block1, finalized.block);
 
-    let finalized = state.finalize();
+    let finalized_with_trees = state.finalize();
+    let finalized = finalized_with_trees.finalized;
     assert_eq!(block2, finalized.block);
 
     assert!(state.best_chain().is_none());

zebra-state/src/tests/setup.rs

@@ -93,7 +93,7 @@ pub(crate) fn new_state_with_mainnet_genesis() -> (StateService, FinalizedBlock)
     let genesis = FinalizedBlock::from(genesis);
     state
         .disk
-        .commit_finalized_direct(genesis.clone(), "test")
+        .commit_finalized_direct(genesis.clone().into(), "test")
         .expect("unexpected invalid genesis block test vector");
 
     assert_eq!(Some((Height(0), genesis.hash)), state.best_tip());