diff --git a/zebra-state/src/request.rs b/zebra-state/src/request.rs index 98bb5c9d..e40db4f5 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 888e9b99..e20e1a6e 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 e4958586..89a4f4fa 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 4929dffb..6ab272bc 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 cc2207b9..2e86d194 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 a96f9e4c..f5b03dc4 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 9c1fbb20..72ce15f8 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 25c316b7..a9a63cb3 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 1963d2d2..42803585 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 5a23bd07..c1540b24 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 ab833c56..214b2d6a 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 d9398d74..a8aa175c 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 2a864942..ba02fda4 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 44d86687..11e2f2a2 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());