From b8712d9a1e358cdabe4a65a5fdb9872afe7ce35d Mon Sep 17 00:00:00 2001 From: Marek Date: Tue, 6 Sep 2022 11:32:54 +0200 Subject: [PATCH] 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 | 72 ++++++++++++++- zebra-state/src/service.rs | 4 +- .../src/service/check/tests/nullifier.rs | 12 +-- zebra-state/src/service/check/tests/utxo.rs | 10 ++- zebra-state/src/service/finalized_state.rs | 89 ++++++++++++++----- .../src/service/finalized_state/tests/prop.rs | 9 +- .../service/finalized_state/zebra_db/block.rs | 7 +- .../service/finalized_state/zebra_db/chain.rs | 12 +-- .../finalized_state/zebra_db/shielded.rs | 5 +- .../src/service/non_finalized_state.rs | 31 ++++--- .../src/service/non_finalized_state/chain.rs | 73 ++++++++++++--- .../service/non_finalized_state/tests/prop.rs | 5 +- .../non_finalized_state/tests/vectors.rs | 6 +- zebra-state/src/tests/setup.rs | 2 +- 14 files changed, 256 insertions(+), 81 deletions(-) diff --git a/zebra-state/src/request.rs b/zebra-state/src/request.rs index 98bb5c9d6..e40db4f53 100644 --- a/zebra-state/src/request.rs +++ b/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, +} + +impl Treestate { + pub fn new( + sprout: Arc, + sapling: Arc, + orchard: Arc, + history_tree: Arc, + ) -> 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, +} + +impl FinalizedWithTrees { + pub fn new(block: ContextuallyValidBlock, treestate: Treestate) -> Self { + let finalized = FinalizedBlock::from(block); + + Self { + finalized, + treestate: Some(treestate), + } + } +} + +impl From> for FinalizedWithTrees { + fn from(block: Arc) -> Self { + Self::from(FinalizedBlock::from(block)) + } +} + +impl From for FinalizedWithTrees { + fn from(block: FinalizedBlock) -> Self { + Self { + finalized: block, + treestate: None, + } + } +} + impl From<&PreparedBlock> for PreparedBlock { fn from(prepared: &PreparedBlock) -> Self { prepared.clone() diff --git a/zebra-state/src/service.rs b/zebra-state/src/service.rs index 888e9b99a..e20e1a6e9 100644 --- a/zebra-state/src/service.rs +++ b/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", ); diff --git a/zebra-state/src/service/check/tests/nullifier.rs b/zebra-state/src/service/check/tests/nullifier.rs index e49585863..89a4f4fa4 100644 --- a/zebra-state/src/service/check/tests/nullifier.rs +++ b/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()); diff --git a/zebra-state/src/service/check/tests/utxo.rs b/zebra-state/src/service/check/tests/utxo.rs index 4929dffb4..6ab272bc3 100644 --- a/zebra-state/src/service/check/tests/utxo.rs +++ b/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()); diff --git a/zebra-state/src/service/finalized_state.rs b/zebra-state/src/service/finalized_state.rs index cc2207b9e..2e86d1943 100644 --- a/zebra-state/src/service/finalized_state.rs +++ b/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 { 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 { + 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 - // to 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 - let history_tree = self.db.history_tree(); - check::block_commitment_is_valid_for_chain_history( - finalized.block.clone(), - self.network, - &history_tree, - )?; + let (history_tree, note_commitment_trees) = match finalized_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)?; + + // 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. + // + // For Nu5-onward, the block hash commits only to + // 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 - .db - .write_block(finalized, history_tree, self.network, source); + let result = self.db.write_block( + finalized, + 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) { diff --git a/zebra-state/src/service/finalized_state/tests/prop.rs b/zebra-state/src/service/finalized_state/tests/prop.rs index a96f9e4c0..f5b03dc4b 100644 --- a/zebra-state/src/service/finalized_state/tests/prop.rs +++ b/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()); diff --git a/zebra-state/src/service/finalized_state/zebra_db/block.rs b/zebra-state/src/service/finalized_state/zebra_db/block.rs index 9c1fbb208..72ce15f89 100644 --- a/zebra-state/src/service/finalized_state/zebra_db/block.rs +++ b/zebra-state/src/service/finalized_state/zebra_db/block.rs @@ -242,6 +242,7 @@ impl ZebraDb { &mut self, finalized: FinalizedBlock, history_tree: Arc, + note_commitment_trees: NoteCommitmentTrees, network: Network, source: &str, ) -> Result { @@ -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, spent_utxos_by_out_loc: BTreeMap, address_balances: HashMap, - mut note_commitment_trees: NoteCommitmentTrees, history_tree: Arc, + note_commitment_trees: NoteCommitmentTrees, value_pool: ValueBalance, ) -> 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)?; diff --git a/zebra-state/src/service/finalized_state/zebra_db/chain.rs b/zebra-state/src/service/finalized_state/zebra_db/chain.rs index 25c316b72..a9a63cb3f 100644 --- a/zebra-state/src/service/finalized_state/zebra_db/chain.rs +++ b/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, - orchard_root: orchard::tree::Root, - mut history_tree: Arc, + history_tree: Arc, ) -> Result<(), BoxError> { let history_tree_cf = db.cf_handle("history_tree").unwrap(); - let FinalizedBlock { block, 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)?; + let FinalizedBlock { height, .. } = finalized; // Update the tree in state let current_tip_height = *height - 1; diff --git a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs index 1963d2d2f..42803585e 100644 --- a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs +++ b/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, diff --git a/zebra-state/src/service/non_finalized_state.rs b/zebra-state/src/service/non_finalized_state.rs index 5a23bd078..c1540b247 100644 --- a/zebra-state/src/service/non_finalized_state.rs +++ b/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 - let finalizing = write_best_chain.pop_root(); + // Pop the lowest height block from the best chain to be finalized, and + // 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 { - if chain.non_finalized_root_hash() != finalizing.hash { + for mut side_chain in side_chains { + 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(); - assert_eq!(chain_start.hash, finalizing.hash); + let (side_chain_root, _treestate) = mut_side_chain.pop_root(); + 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: diff --git a/zebra-state/src/service/non_finalized_state/chain.rs b/zebra-state/src/service/non_finalized_state/chain.rs index ab833c568..214b2d6aa 100644 --- a/zebra-state/src/service/non_finalized_state/chain.rs +++ b/zebra-state/src/service/non_finalized_state/chain.rs @@ -30,8 +30,8 @@ use zebra_chain::{ }; use crate::{ - service::check, ContextuallyValidBlock, HashOrHeight, OutputLocation, TransactionLocation, - ValidateContextError, + request::Treestate, service::check, ContextuallyValidBlock, HashOrHeight, OutputLocation, + 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>, + /// The Sprout note commitment tree for each height. + pub(crate) sprout_trees_by_height: + BTreeMap>, /// 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, @@ -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 + (block, treestate) } /// 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> { + 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> { + 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 { + 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 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 diff --git a/zebra-state/src/service/non_finalized_state/tests/prop.rs b/zebra-state/src/service/non_finalized_state/tests/prop.rs index d9398d741..a8aa175c8 100644 --- a/zebra-state/src/service/non_finalized_state/tests/prop.rs +++ b/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(); diff --git a/zebra-state/src/service/non_finalized_state/tests/vectors.rs b/zebra-state/src/service/non_finalized_state/tests/vectors.rs index 2a864942e..ba02fda4d 100644 --- a/zebra-state/src/service/non_finalized_state/tests/vectors.rs +++ b/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()); diff --git a/zebra-state/src/tests/setup.rs b/zebra-state/src/tests/setup.rs index 44d86687d..11e2f2a26 100644 --- a/zebra-state/src/tests/setup.rs +++ b/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());