change(state): Refactor the structure of finalizable blocks (#7035)

* Add and use `FinalizableBlock`

This commit adds `FinalizableBlock`, and uses it instead of
`ContextuallyVerifiedBlockWithTrees` in `commit_finalized_direct()`

* Use `ContextuallyVerifiedBlockWithTrees`

This commit passes `ContextuallyVerifiedBlockWithTrees` instead of
passing separate `finalized`, `history_tree` and `note_commitment_trees`
when storing blocks in the finalized state.

* Apply suggestions from code review

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

* add docs to new methods

* fix existing doc

* rename `ContextuallyVerifiedBlockWithTrees` to `SemanticallyVerifiedBlockWithTrees`

* Refactor docs

* Refactor comments

* Add missing docs, fix typo

* Fix rustfmt

---------

Co-authored-by: teor <teor@riseup.net>
Co-authored-by: Alfredo Garcia <oxarbitrage@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

zebra-state/src/request.rs

@@ -255,46 +255,65 @@ impl Treestate {
 /// 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 a new one.
-pub struct ContextuallyVerifiedBlockWithTrees {
+pub struct SemanticallyVerifiedBlockWithTrees {
     /// A block ready to be committed.
-    pub block: SemanticallyVerifiedBlock,
+    pub verified: SemanticallyVerifiedBlock,
     /// The tresstate associated with the block.
-    pub treestate: Option<Treestate>,
+    pub treestate: Treestate,
 }
 
-impl ContextuallyVerifiedBlockWithTrees {
+/// Contains a block ready to be committed.
+///
+/// Zebra's state service passes this `enum` over to the finalized state
+/// when committing a block.
+pub enum FinalizableBlock {
+    Checkpoint {
+        checkpoint_verified: CheckpointVerifiedBlock,
+    },
+    Contextual {
+        contextually_verified: ContextuallyVerifiedBlock,
+        treestate: Treestate,
+    },
+}
+
+impl FinalizableBlock {
+    /// Create a new [`FinalizableBlock`] given a [`ContextuallyVerifiedBlock`].
     pub fn new(contextually_verified: ContextuallyVerifiedBlock, treestate: Treestate) -> Self {
-        Self {
+        Self::Contextual {
-            block: SemanticallyVerifiedBlock::from(contextually_verified),
+            contextually_verified,
-            treestate: Some(treestate),
+            treestate,
+        }
+    }
+
+    #[cfg(test)]
+    /// Extract a [`Block`] from a [`FinalizableBlock`] variant.
+    pub fn inner_block(&self) -> Arc<Block> {
+        match self {
+            FinalizableBlock::Checkpoint {
+                checkpoint_verified,
+            } => checkpoint_verified.block.clone(),
+            FinalizableBlock::Contextual {
+                contextually_verified,
+                ..
+            } => contextually_verified.block.clone(),
         }
     }
 }
 
-impl From<Arc<Block>> for ContextuallyVerifiedBlockWithTrees {
+impl From<CheckpointVerifiedBlock> for FinalizableBlock {
-    fn from(block: Arc<Block>) -> Self {
-        Self::from(SemanticallyVerifiedBlock::from(block))
-    }
-}
-
-impl From<SemanticallyVerifiedBlock> for ContextuallyVerifiedBlockWithTrees {
-    fn from(semantically_verified: SemanticallyVerifiedBlock) -> Self {
-        Self {
-            block: semantically_verified,
-            treestate: None,
-        }
-    }
-}
-
-impl From<CheckpointVerifiedBlock> for ContextuallyVerifiedBlockWithTrees {
     fn from(checkpoint_verified: CheckpointVerifiedBlock) -> Self {
-        Self {
+        Self::Checkpoint {
-            block: checkpoint_verified.0,
+            checkpoint_verified,
-            treestate: None,
         }
     }
 }
 
+impl From<Arc<Block>> for FinalizableBlock {
+    fn from(block: Arc<Block>) -> Self {
+        Self::from(CheckpointVerifiedBlock::from(block))
+    }
+}
+
 impl From<&SemanticallyVerifiedBlock> for SemanticallyVerifiedBlock {
     fn from(semantically_verified: &SemanticallyVerifiedBlock) -> Self {
         semantically_verified.clone()
@@ -413,6 +432,12 @@ impl From<ContextuallyVerifiedBlock> for SemanticallyVerifiedBlock {
     }
 }
 
+impl From<CheckpointVerifiedBlock> for SemanticallyVerifiedBlock {
+    fn from(checkpoint_verified: CheckpointVerifiedBlock) -> Self {
+        checkpoint_verified.0
+    }
+}
+
 impl Deref for CheckpointVerifiedBlock {
     type Target = SemanticallyVerifiedBlock;
 

zebra-state/src/service/finalized_state.rs

@@ -23,7 +23,7 @@ use std::{
 use zebra_chain::{block, parameters::Network};
 
 use crate::{
-    request::ContextuallyVerifiedBlockWithTrees,
+    request::{FinalizableBlock, SemanticallyVerifiedBlockWithTrees, Treestate},
     service::{check, QueuedCheckpointVerified},
     BoxError, CheckpointVerifiedBlock, CloneError, Config,
 };
@@ -225,53 +225,25 @@ impl FinalizedState {
     #[allow(clippy::unwrap_in_result)]
     pub fn commit_finalized_direct(
         &mut self,
-        contextually_verified_with_trees: ContextuallyVerifiedBlockWithTrees,
+        finalizable_block: FinalizableBlock,
         source: &str,
     ) -> Result<block::Hash, BoxError> {
-        let finalized = contextually_verified_with_trees.block;
+        let (height, hash, finalized) = match finalizable_block {
-        let committed_tip_hash = self.db.finalized_tip_hash();
+            FinalizableBlock::Checkpoint {
-        let committed_tip_height = self.db.finalized_tip_height();
+                checkpoint_verified,
+            } => {
+                // Checkpoint-verified blocks don't have an associated treestate, so we retrieve the
+                // treestate of the finalized tip from the database and update it for the block
+                // being committed, assuming the retrieved treestate is the parent block's
+                // treestate. Later on, this function proves this assumption by asserting that the
+                // finalized tip is the parent block of the block being committed.
 
-        // Assert that callers (including unit tests) get the chain order correct
+                let block = checkpoint_verified.block.clone();
-        if self.db.is_empty() {
-            assert_eq!(
-                committed_tip_hash, finalized.block.header.previous_block_hash,
-                "the first block added to an empty state must be a genesis block, source: {source}",
-            );
-            assert_eq!(
-                block::Height(0),
-                finalized.height,
-                "cannot commit genesis: invalid height, source: {source}",
-            );
-        } else {
-            assert_eq!(
-                committed_tip_height.expect("state must have a genesis block committed") + 1,
-                Some(finalized.height),
-                "committed block height must be 1 more than the finalized tip height, source: {source}",
-            );
-
-            assert_eq!(
-                committed_tip_hash, finalized.block.header.previous_block_hash,
-                "committed block must be a child of the finalized tip, source: {source}",
-            );
-        }
-
-        let (history_tree, note_commitment_trees) = match contextually_verified_with_trees.treestate
-        {
-            // If the treestate associated with the block was supplied, use it
-            // without recomputing it.
-            Some(ref treestate) => (
-                treestate.history_tree.clone(),
-                treestate.note_commitment_trees.clone(),
-            ),
-            // 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)?;
+                note_commitment_trees.update_trees_parallel(&block)?;
 
                 // Check the block commitment if the history tree was not
                 // supplied by the non-finalized state. Note that we don't do
@@ -291,7 +263,7 @@ impl FinalizedState {
                 // 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(),
+                    block.clone(),
                     self.network,
                     &history_tree,
                 )?;
@@ -303,30 +275,64 @@ impl FinalizedState {
                 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(
+                history_tree_mut.push(self.network(), block.clone(), sapling_root, orchard_root)?;
-                    self.network(),
-                    finalized.block.clone(),
-                    sapling_root,
-                    orchard_root,
-                )?;
 
-                (history_tree, note_commitment_trees)
+                (
+                    checkpoint_verified.height,
+                    checkpoint_verified.hash,
+                    SemanticallyVerifiedBlockWithTrees {
+                        verified: checkpoint_verified.0,
+                        treestate: Treestate {
+                            note_commitment_trees,
+                            history_tree,
+                        },
+                    },
+                )
             }
+            FinalizableBlock::Contextual {
+                contextually_verified,
+                treestate,
+            } => (
+                contextually_verified.height,
+                contextually_verified.hash,
+                SemanticallyVerifiedBlockWithTrees {
+                    verified: contextually_verified.into(),
+                    treestate,
+                },
+            ),
         };
 
-        let finalized_height = finalized.height;
+        let committed_tip_hash = self.db.finalized_tip_hash();
-        let finalized_hash = finalized.hash;
+        let committed_tip_height = self.db.finalized_tip_height();
+
+        // Assert that callers (including unit tests) get the chain order correct
+        if self.db.is_empty() {
+            assert_eq!(
+                committed_tip_hash, finalized.verified.block.header.previous_block_hash,
+                "the first block added to an empty state must be a genesis block, source: {source}",
+            );
+            assert_eq!(
+                block::Height(0),
+                height,
+                "cannot commit genesis: invalid height, source: {source}",
+            );
+        } else {
+            assert_eq!(
+                committed_tip_height.expect("state must have a genesis block committed") + 1,
+                Some(height),
+                "committed block height must be 1 more than the finalized tip height, source: {source}",
+            );
+
+            assert_eq!(
+                committed_tip_hash, finalized.verified.block.header.previous_block_hash,
+                "committed block must be a child of the finalized tip, source: {source}",
+            );
+        }
 
         #[cfg(feature = "elasticsearch")]
-        let finalized_block = finalized.block.clone();
+        let finalized_block = finalized.verified.block.clone();
 
-        let result = self.db.write_block(
+        let result = self.db.write_block(finalized, self.network, source);
-            finalized,
-            history_tree,
-            note_commitment_trees,
-            self.network,
-            source,
-        );
 
         if result.is_ok() {
             // Save blocks to elasticsearch if the feature is enabled.
@@ -334,10 +340,10 @@ impl FinalizedState {
             self.elasticsearch(&finalized_block);
 
             // TODO: move the stop height check to the syncer (#3442)
-            if self.is_at_stop_height(finalized_height) {
+            if self.is_at_stop_height(height) {
                 tracing::info!(
-                    height = ?finalized_height,
+                    ?height,
-                    hash = ?finalized_hash,
+                    ?hash,
                     block_source = ?source,
                     "stopping at configured height, flushing database to disk"
                 );

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

@@ -19,9 +19,7 @@ use itertools::Itertools;
 use zebra_chain::{
     amount::NonNegative,
     block::{self, Block, Height},
-    history_tree::HistoryTree,
     orchard,
-    parallel::tree::NoteCommitmentTrees,
     parameters::{Network, GENESIS_PREVIOUS_BLOCK_HASH},
     sapling,
     serialization::TrustedPreallocate,
@@ -31,6 +29,7 @@ use zebra_chain::{
 };
 
 use crate::{
+    request::SemanticallyVerifiedBlockWithTrees,
     service::finalized_state::{
         disk_db::{DiskDb, DiskWriteBatch, ReadDisk, WriteDisk},
         disk_format::{
@@ -281,15 +280,12 @@ impl ZebraDb {
     /// - Propagates any errors from updating history and note commitment trees
     pub(in super::super) fn write_block(
         &mut self,
-        finalized: SemanticallyVerifiedBlock,
+        finalized: SemanticallyVerifiedBlockWithTrees,
-        history_tree: Arc<HistoryTree>,
-        note_commitment_trees: NoteCommitmentTrees,
         network: Network,
         source: &str,
     ) -> Result<block::Hash, BoxError> {
-        let finalized_hash = finalized.hash;
-
         let tx_hash_indexes: HashMap<transaction::Hash, usize> = finalized
+            .verified
             .transaction_hashes
             .iter()
             .enumerate()
@@ -302,11 +298,12 @@ impl ZebraDb {
         //       simplify the spent_utxos location lookup code,
         //       and remove the extra new_outputs_by_out_loc argument
         let new_outputs_by_out_loc: BTreeMap<OutputLocation, transparent::Utxo> = finalized
+            .verified
             .new_outputs
             .iter()
             .map(|(outpoint, ordered_utxo)| {
                 (
-                    lookup_out_loc(finalized.height, outpoint, &tx_hash_indexes),
+                    lookup_out_loc(finalized.verified.height, outpoint, &tx_hash_indexes),
                     ordered_utxo.utxo.clone(),
                 )
             })
@@ -315,6 +312,7 @@ impl ZebraDb {
         // Get a list of the spent UTXOs, before we delete any from the database
         let spent_utxos: Vec<(transparent::OutPoint, OutputLocation, transparent::Utxo)> =
             finalized
+                .verified
                 .block
                 .transactions
                 .iter()
@@ -326,12 +324,13 @@ impl ZebraDb {
                         // Some utxos are spent in the same block, so they will be in
                         // `tx_hash_indexes` and `new_outputs`
                         self.output_location(&outpoint).unwrap_or_else(|| {
-                            lookup_out_loc(finalized.height, &outpoint, &tx_hash_indexes)
+                            lookup_out_loc(finalized.verified.height, &outpoint, &tx_hash_indexes)
                         }),
                         self.utxo(&outpoint)
                             .map(|ordered_utxo| ordered_utxo.utxo)
                             .or_else(|| {
                                 finalized
+                                    .verified
                                     .new_outputs
                                     .get(&outpoint)
                                     .map(|ordered_utxo| ordered_utxo.utxo.clone())
@@ -356,6 +355,7 @@ impl ZebraDb {
             .values()
             .chain(
                 finalized
+                    .verified
                     .new_outputs
                     .values()
                     .map(|ordered_utxo| &ordered_utxo.utxo),
@@ -376,13 +376,11 @@ impl ZebraDb {
         // In case of errors, propagate and do not write the batch.
         batch.prepare_block_batch(
             &self.db,
-            finalized,
+            &finalized,
             new_outputs_by_out_loc,
             spent_utxos_by_outpoint,
             spent_utxos_by_out_loc,
             address_balances,
-            history_tree,
-            note_commitment_trees,
             self.finalized_value_pool(),
         )?;
 
@@ -390,7 +388,7 @@ impl ZebraDb {
 
         tracing::trace!(?source, "committed block from");
 
-        Ok(finalized_hash)
+        Ok(finalized.verified.hash)
     }
 }
 
@@ -429,25 +427,16 @@ impl DiskWriteBatch {
     pub fn prepare_block_batch(
         &mut self,
         db: &DiskDb,
-        finalized: SemanticallyVerifiedBlock,
+        finalized: &SemanticallyVerifiedBlockWithTrees,
         new_outputs_by_out_loc: BTreeMap<OutputLocation, transparent::Utxo>,
         spent_utxos_by_outpoint: HashMap<transparent::OutPoint, transparent::Utxo>,
         spent_utxos_by_out_loc: BTreeMap<OutputLocation, transparent::Utxo>,
         address_balances: HashMap<transparent::Address, AddressBalanceLocation>,
-        history_tree: Arc<HistoryTree>,
-        note_commitment_trees: NoteCommitmentTrees,
         value_pool: ValueBalance<NonNegative>,
     ) -> Result<(), BoxError> {
-        let SemanticallyVerifiedBlock {
-            block,
-            hash,
-            height,
-            ..
-        } = &finalized;
-
         // Commit block and transaction data.
         // (Transaction indexes, note commitments, and UTXOs are committed later.)
-        self.prepare_block_header_and_transaction_data_batch(db, &finalized)?;
+        self.prepare_block_header_and_transaction_data_batch(db, &finalized.verified)?;
 
         // # Consensus
         //
@@ -458,28 +447,37 @@ impl DiskWriteBatch {
         //
         // By returning early, Zebra commits the genesis block and transaction data,
         // but it ignores the genesis UTXO and value pool updates.
-        if self.prepare_genesis_batch(db, &finalized) {
+        if self.prepare_genesis_batch(db, &finalized.verified) {
             return Ok(());
         }
 
         // Commit transaction indexes
         self.prepare_transparent_transaction_batch(
             db,
-            &finalized,
+            &finalized.verified,
             &new_outputs_by_out_loc,
             &spent_utxos_by_outpoint,
             &spent_utxos_by_out_loc,
             address_balances,
         )?;
-        self.prepare_shielded_transaction_batch(db, &finalized)?;
+        self.prepare_shielded_transaction_batch(db, &finalized.verified)?;
 
-        self.prepare_note_commitment_batch(db, &finalized, note_commitment_trees, history_tree)?;
+        self.prepare_note_commitment_batch(db, finalized)?;
 
         // Commit UTXOs and value pools
-        self.prepare_chain_value_pools_batch(db, &finalized, spent_utxos_by_outpoint, value_pool)?;
+        self.prepare_chain_value_pools_batch(
+            db,
+            &finalized.verified,
+            spent_utxos_by_outpoint,
+            value_pool,
+        )?;
 
         // The block has passed contextual validation, so update the metrics
-        block_precommit_metrics(block, *hash, *height);
+        block_precommit_metrics(
+            &finalized.verified.block,
+            finalized.verified.hash,
+            finalized.verified.height,
+        );
 
         Ok(())
     }

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

@@ -21,6 +21,7 @@ use zebra_chain::{
 };
 
 use crate::{
+    request::SemanticallyVerifiedBlockWithTrees,
     service::finalized_state::{
         disk_db::{DiskDb, DiskWriteBatch, ReadDisk, WriteDisk},
         zebra_db::ZebraDb,
@@ -69,15 +70,14 @@ impl DiskWriteBatch {
     pub fn prepare_history_batch(
         &mut self,
         db: &DiskDb,
-        finalized: &SemanticallyVerifiedBlock,
+        finalized: &SemanticallyVerifiedBlockWithTrees,
-        history_tree: Arc<HistoryTree>,
     ) -> Result<(), BoxError> {
         let history_tree_cf = db.cf_handle("history_tree").unwrap();
 
-        let SemanticallyVerifiedBlock { height, .. } = finalized;
+        let height = finalized.verified.height;
 
         // Update the tree in state
-        let current_tip_height = *height - 1;
+        let current_tip_height = height - 1;
         if let Some(h) = current_tip_height {
             self.zs_delete(&history_tree_cf, h);
         }
@@ -87,7 +87,7 @@ impl DiskWriteBatch {
         // Otherwise, the ReadStateService could access a height
         // that was just deleted by a concurrent StateService write.
         // This requires a database version update.
-        if let Some(history_tree) = history_tree.as_ref().as_ref() {
+        if let Some(history_tree) = finalized.treestate.history_tree.as_ref().as_ref() {
             self.zs_insert(&history_tree_cf, height, history_tree);
         }
 

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

@@ -15,11 +15,12 @@
 use std::sync::Arc;
 
 use zebra_chain::{
-    block::Height, history_tree::HistoryTree, orchard, parallel::tree::NoteCommitmentTrees,
+    block::Height, orchard, parallel::tree::NoteCommitmentTrees, sapling, sprout,
-    sapling, sprout, transaction::Transaction,
+    transaction::Transaction,
 };
 
 use crate::{
+    request::SemanticallyVerifiedBlockWithTrees,
     service::finalized_state::{
         disk_db::{DiskDb, DiskWriteBatch, ReadDisk, WriteDisk},
         zebra_db::ZebraDb,
@@ -264,9 +265,7 @@ impl DiskWriteBatch {
     pub fn prepare_note_commitment_batch(
         &mut self,
         db: &DiskDb,
-        finalized: &SemanticallyVerifiedBlock,
+        finalized: &SemanticallyVerifiedBlockWithTrees,
-        note_commitment_trees: NoteCommitmentTrees,
-        history_tree: Arc<HistoryTree>,
     ) -> Result<(), BoxError> {
         let sprout_anchors = db.cf_handle("sprout_anchors").unwrap();
         let sapling_anchors = db.cf_handle("sapling_anchors").unwrap();
@@ -276,7 +275,8 @@ impl DiskWriteBatch {
         let sapling_note_commitment_tree_cf = db.cf_handle("sapling_note_commitment_tree").unwrap();
         let orchard_note_commitment_tree_cf = db.cf_handle("orchard_note_commitment_tree").unwrap();
 
-        let SemanticallyVerifiedBlock { height, .. } = finalized;
+        let height = finalized.verified.height;
+        let note_commitment_trees = finalized.treestate.note_commitment_trees.clone();
 
         // Use the cached values that were previously calculated in parallel.
         let sprout_root = note_commitment_trees.sprout.root();
@@ -290,7 +290,7 @@ impl DiskWriteBatch {
         self.zs_insert(&orchard_anchors, orchard_root, ());
 
         // Delete the previously stored Sprout note commitment tree.
-        let current_tip_height = *height - 1;
+        let current_tip_height = height - 1;
         if let Some(h) = current_tip_height {
             self.zs_delete(&sprout_note_commitment_tree_cf, h);
         }
@@ -317,7 +317,7 @@ impl DiskWriteBatch {
             note_commitment_trees.orchard,
         );
 
-        self.prepare_history_batch(db, finalized, history_tree)
+        self.prepare_history_batch(db, finalized)
     }
 
     /// Prepare a database batch containing the initial note commitment trees,

zebra-state/src/service/non_finalized_state.rs

@@ -16,7 +16,7 @@ use zebra_chain::{
 
 use crate::{
     constants::MAX_NON_FINALIZED_CHAIN_FORKS,
-    request::{ContextuallyVerifiedBlock, ContextuallyVerifiedBlockWithTrees},
+    request::{ContextuallyVerifiedBlock, FinalizableBlock},
     service::{check, finalized_state::ZebraDb},
     SemanticallyVerifiedBlock, ValidateContextError,
 };
@@ -174,7 +174,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) -> ContextuallyVerifiedBlockWithTrees {
+    pub fn finalize(&mut self) -> FinalizableBlock {
         // 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.)
@@ -226,7 +226,7 @@ impl NonFinalizedState {
         self.update_metrics_for_chains();
 
         // Add the treestate to the finalized block.
-        ContextuallyVerifiedBlockWithTrees::new(best_chain_root, root_treestate)
+        FinalizableBlock::new(best_chain_root, root_treestate)
     }
 
     /// Commit block to the non-finalized state, on top of:

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

@@ -213,13 +213,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_with_trees = state.finalize();
+    let finalized = state.finalize().inner_block();
-    let finalized = finalized_with_trees.block;
-    assert_eq!(block1, finalized.block);
 
-    let finalized_with_trees = state.finalize();
+    assert_eq!(block1, finalized);
-    let finalized = finalized_with_trees.block;
+
-    assert_eq!(block2, finalized.block);
+    let finalized = state.finalize().inner_block();
+    assert_eq!(block2, finalized);
 
     assert!(state.best_chain().is_none());