zecfoundation
/
zebra
mirror of https://github.com/ZcashFoundation/zebra.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

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>
This commit is contained in:
Marek 2022-09-06 11:32:54 +02:00 committed by GitHub
parent fec012a006
commit b8712d9a1e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
14 changed files with 256 additions and 81 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

72
zebra-state/src/request.rs
View File

@ -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()

4
zebra-state/src/service.rs
View File

@ -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",
);

12
zebra-state/src/service/check/tests/nullifier.rs
View File

@ -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());

10
zebra-state/src/service/check/tests/utxo.rs
View File

@ -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());

89
zebra-state/src/service/finalized_state.rs
View File

@ -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
// 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) {

9
zebra-state/src/service/finalized_state/tests/prop.rs
View File

@ -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());

7
zebra-state/src/service/finalized_state/zebra_db/block.rs
View File

@ -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)?;

12
zebra-state/src/service/finalized_state/zebra_db/chain.rs
View File

@ -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<HistoryTree>,
history_tree: Arc<HistoryTree>,
) -> 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;

5
zebra-state/src/service/finalized_state/zebra_db/shielded.rs
View File

@ -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,

31
zebra-state/src/service/non_finalized_state.rs
View File

@ -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:

73
zebra-state/src/service/non_finalized_state/chain.rs
View File

@ -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<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
(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<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

5
zebra-state/src/service/non_finalized_state/tests/prop.rs
View File

@ -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();

6
zebra-state/src/service/non_finalized_state/tests/vectors.rs
View File

@ -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());

2
zebra-state/src/tests/setup.rs
View File

@ -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());