diff --git a/zebra-state/src/service/check/anchors.rs b/zebra-state/src/service/check/anchors.rs index 4413fc76..ab62f03a 100644 --- a/zebra-state/src/service/check/anchors.rs +++ b/zebra-state/src/service/check/anchors.rs @@ -1,7 +1,7 @@ //! Checks for whether cited anchors are previously-computed note commitment //! tree roots. -use std::collections::HashMap; +use std::{collections::HashMap, sync::Arc}; use zebra_chain::sprout; @@ -29,7 +29,7 @@ pub(crate) fn anchors_refer_to_earlier_treestates( if transaction.has_sprout_joinsplit_data() { let mut interstitial_trees: HashMap< sprout::tree::Root, - sprout::tree::NoteCommitmentTree, + Arc, > = HashMap::new(); for joinsplit in transaction.sprout_groth16_joinsplits() { @@ -90,11 +90,13 @@ pub(crate) fn anchors_refer_to_earlier_treestates( } }; + let input_tree_inner = Arc::make_mut(&mut input_tree); + tracing::debug!(?joinsplit.anchor, "validated sprout anchor"); // Add new anchors to the interstitial note commitment tree. for cm in joinsplit.commitments { - input_tree + input_tree_inner .append(cm) .expect("note commitment should be appendable to the tree"); } diff --git a/zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshot.rs b/zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshot.rs index e40c2c28..e5be8ec5 100644 --- a/zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshot.rs +++ b/zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshot.rs @@ -219,8 +219,8 @@ fn snapshot_block_and_transaction_data(state: &FinalizedState) { .orchard_note_commitment_tree_by_height(&block::Height::MIN) .expect("the genesis block in the database has an Orchard tree"); - assert_eq!(sapling_tree, sapling::tree::NoteCommitmentTree::default()); - assert_eq!(orchard_tree, orchard::tree::NoteCommitmentTree::default()); + assert_eq!(*sapling_tree, sapling::tree::NoteCommitmentTree::default()); + assert_eq!(*orchard_tree, orchard::tree::NoteCommitmentTree::default()); // Blocks let mut stored_block_hashes = Vec::new(); 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 c69a481b..95b0b63f 100644 --- a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs +++ b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs @@ -12,6 +12,8 @@ //! The [`crate::constants::DATABASE_FORMAT_VERSION`] constant must //! be incremented each time the database format (column, serialization, etc) changes. +use std::sync::Arc; + use zebra_chain::{ block::Height, history_tree::HistoryTree, orchard, sapling, sprout, transaction::Transaction, }; @@ -28,9 +30,9 @@ use crate::{ /// An argument wrapper struct for note commitment trees. #[derive(Clone, Debug)] pub struct NoteCommitmentTrees { - sprout: sprout::tree::NoteCommitmentTree, - sapling: sapling::tree::NoteCommitmentTree, - orchard: orchard::tree::NoteCommitmentTree, + sprout: Arc, + sapling: Arc, + orchard: Arc, } impl ZebraDb { @@ -75,16 +77,17 @@ impl ZebraDb { /// Returns the Sprout note commitment tree of the finalized tip /// or the empty tree if the state is empty. - pub fn sprout_note_commitment_tree(&self) -> sprout::tree::NoteCommitmentTree { + pub fn sprout_note_commitment_tree(&self) -> Arc { let height = match self.finalized_tip_height() { Some(h) => h, None => return Default::default(), }; - let sprout_note_commitment_tree = self.db.cf_handle("sprout_note_commitment_tree").unwrap(); + let sprout_nct_handle = self.db.cf_handle("sprout_note_commitment_tree").unwrap(); self.db - .zs_get(&sprout_note_commitment_tree, &height) + .zs_get(&sprout_nct_handle, &height) + .map(Arc::new) .expect("Sprout note commitment tree must exist if there is a finalized tip") } @@ -95,25 +98,27 @@ impl ZebraDb { pub fn sprout_note_commitment_tree_by_anchor( &self, sprout_anchor: &sprout::tree::Root, - ) -> Option { - let sprout_anchors = self.db.cf_handle("sprout_anchors").unwrap(); + ) -> Option> { + let sprout_anchors_handle = self.db.cf_handle("sprout_anchors").unwrap(); - self.db.zs_get(&sprout_anchors, sprout_anchor) + self.db + .zs_get(&sprout_anchors_handle, sprout_anchor) + .map(Arc::new) } /// Returns the Sapling note commitment tree of the finalized tip /// or the empty tree if the state is empty. - pub fn sapling_note_commitment_tree(&self) -> sapling::tree::NoteCommitmentTree { + pub fn sapling_note_commitment_tree(&self) -> Arc { let height = match self.finalized_tip_height() { Some(h) => h, None => return Default::default(), }; - let sapling_note_commitment_tree = - self.db.cf_handle("sapling_note_commitment_tree").unwrap(); + let sapling_nct_handle = self.db.cf_handle("sapling_note_commitment_tree").unwrap(); self.db - .zs_get(&sapling_note_commitment_tree, &height) + .zs_get(&sapling_nct_handle, &height) + .map(Arc::new) .expect("Sapling note commitment tree must exist if there is a finalized tip") } @@ -123,25 +128,25 @@ impl ZebraDb { pub fn sapling_note_commitment_tree_by_height( &self, height: &Height, - ) -> Option { + ) -> Option> { let sapling_trees = self.db.cf_handle("sapling_note_commitment_tree").unwrap(); - self.db.zs_get(&sapling_trees, height) + self.db.zs_get(&sapling_trees, height).map(Arc::new) } /// Returns the Orchard note commitment tree of the finalized tip /// or the empty tree if the state is empty. - pub fn orchard_note_commitment_tree(&self) -> orchard::tree::NoteCommitmentTree { + pub fn orchard_note_commitment_tree(&self) -> Arc { let height = match self.finalized_tip_height() { Some(h) => h, None => return Default::default(), }; - let orchard_note_commitment_tree = - self.db.cf_handle("orchard_note_commitment_tree").unwrap(); + let orchard_nct_handle = self.db.cf_handle("orchard_note_commitment_tree").unwrap(); self.db - .zs_get(&orchard_note_commitment_tree, &height) + .zs_get(&orchard_nct_handle, &height) + .map(Arc::new) .expect("Orchard note commitment tree must exist if there is a finalized tip") } @@ -151,10 +156,10 @@ impl ZebraDb { pub fn orchard_note_commitment_tree_by_height( &self, height: &Height, - ) -> Option { + ) -> Option> { let orchard_trees = self.db.cf_handle("orchard_note_commitment_tree").unwrap(); - self.db.zs_get(&orchard_trees, height) + self.db.zs_get(&orchard_trees, height).map(Arc::new) } /// Returns the shielded note commitment trees of the finalized tip @@ -238,21 +243,19 @@ impl DiskWriteBatch { transaction: &Transaction, note_commitment_trees: &mut NoteCommitmentTrees, ) -> Result<(), BoxError> { - // Update the note commitment trees + let sprout_nct = Arc::make_mut(&mut note_commitment_trees.sprout); for sprout_note_commitment in transaction.sprout_note_commitments() { - note_commitment_trees - .sprout - .append(*sprout_note_commitment)?; + sprout_nct.append(*sprout_note_commitment)?; } + + let sapling_nct = Arc::make_mut(&mut note_commitment_trees.sapling); for sapling_note_commitment in transaction.sapling_note_commitments() { - note_commitment_trees - .sapling - .append(*sapling_note_commitment)?; + sapling_nct.append(*sapling_note_commitment)?; } + + let orchard_nct = Arc::make_mut(&mut note_commitment_trees.orchard); for orchard_note_commitment in transaction.orchard_note_commitments() { - note_commitment_trees - .orchard - .append(*orchard_note_commitment)?; + orchard_nct.append(*orchard_note_commitment)?; } Ok(()) diff --git a/zebra-state/src/service/non_finalized_state.rs b/zebra-state/src/service/non_finalized_state.rs index 875a26f4..251582a8 100644 --- a/zebra-state/src/service/non_finalized_state.rs +++ b/zebra-state/src/service/non_finalized_state.rs @@ -385,9 +385,9 @@ impl NonFinalizedState { fn parent_chain( &mut self, parent_hash: block::Hash, - sprout_note_commitment_tree: sprout::tree::NoteCommitmentTree, - sapling_note_commitment_tree: sapling::tree::NoteCommitmentTree, - orchard_note_commitment_tree: orchard::tree::NoteCommitmentTree, + sprout_note_commitment_tree: Arc, + sapling_note_commitment_tree: Arc, + orchard_note_commitment_tree: Arc, history_tree: HistoryTree, ) -> Result, ValidateContextError> { match self.find_chain(|chain| chain.non_finalized_tip_hash() == parent_hash) { diff --git a/zebra-state/src/service/non_finalized_state/chain.rs b/zebra-state/src/service/non_finalized_state/chain.rs index 0ad9794d..43379110 100644 --- a/zebra-state/src/service/non_finalized_state/chain.rs +++ b/zebra-state/src/service/non_finalized_state/chain.rs @@ -63,21 +63,23 @@ pub struct Chain { /// The Sprout note commitment tree of the tip of this [`Chain`], /// including all finalized notes, and the non-finalized notes in this chain. - pub(super) sprout_note_commitment_tree: sprout::tree::NoteCommitmentTree, + pub(super) sprout_note_commitment_tree: Arc, /// The Sprout note commitment tree for each anchor. /// This is required for interstitial states. pub(crate) sprout_trees_by_anchor: - HashMap, + HashMap>, /// 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: sapling::tree::NoteCommitmentTree, + pub(super) sapling_note_commitment_tree: Arc, /// The Sapling note commitment tree for each height. - pub(crate) sapling_trees_by_height: BTreeMap, + pub(crate) sapling_trees_by_height: + BTreeMap>, /// The Orchard note commitment tree of the tip of this [`Chain`], /// including all finalized notes, and the non-finalized notes in this chain. - pub(super) orchard_note_commitment_tree: orchard::tree::NoteCommitmentTree, + pub(super) orchard_note_commitment_tree: Arc, /// The Orchard note commitment tree for each height. - pub(crate) orchard_trees_by_height: BTreeMap, + pub(crate) orchard_trees_by_height: + BTreeMap>, /// The ZIP-221 history tree of the tip of this [`Chain`], /// including all finalized blocks, and the non-finalized `blocks` in this chain. pub(crate) history_tree: HistoryTree, @@ -125,9 +127,9 @@ impl Chain { /// Create a new Chain with the given trees and network. pub(crate) fn new( network: Network, - sprout_note_commitment_tree: sprout::tree::NoteCommitmentTree, - sapling_note_commitment_tree: sapling::tree::NoteCommitmentTree, - orchard_note_commitment_tree: orchard::tree::NoteCommitmentTree, + sprout_note_commitment_tree: Arc, + sapling_note_commitment_tree: Arc, + orchard_note_commitment_tree: Arc, history_tree: HistoryTree, finalized_tip_chain_value_pools: ValueBalance, ) -> Self { @@ -264,15 +266,16 @@ impl Chain { /// Fork a chain at the block with the given hash, if it is part of this /// chain. /// - /// The trees must match the trees of the finalized tip and are used - /// to rebuild them after the fork. + /// The passed trees must match the trees of the finalized tip. They are + /// extended by the commitments from the newly forked chain up to the passed + /// `fork_tip`. #[allow(clippy::unwrap_in_result)] pub fn fork( &self, fork_tip: block::Hash, - sprout_note_commitment_tree: sprout::tree::NoteCommitmentTree, - sapling_note_commitment_tree: sapling::tree::NoteCommitmentTree, - orchard_note_commitment_tree: orchard::tree::NoteCommitmentTree, + sprout_note_commitment_tree: Arc, + sapling_note_commitment_tree: Arc, + orchard_note_commitment_tree: Arc, history_tree: HistoryTree, ) -> Result, ValidateContextError> { if !self.height_by_hash.contains_key(&fork_tip) { @@ -290,29 +293,27 @@ impl Chain { forked.pop_tip(); } + let sprout_nct = Arc::make_mut(&mut forked.sprout_note_commitment_tree); + let sapling_nct = Arc::make_mut(&mut forked.sapling_note_commitment_tree); + let orchard_nct = Arc::make_mut(&mut forked.orchard_note_commitment_tree); + // Rebuild the note commitment trees, starting from the finalized tip tree. - // TODO: change to a more efficient approach by removing nodes - // from the tree of the original chain (in [`Self::pop_tip`]). - // See https://github.com/ZcashFoundation/zebra/issues/2378 for block in forked.blocks.values() { for transaction in block.block.transactions.iter() { for sprout_note_commitment in transaction.sprout_note_commitments() { - forked - .sprout_note_commitment_tree + sprout_nct .append(*sprout_note_commitment) .expect("must work since it was already appended before the fork"); } for sapling_note_commitment in transaction.sapling_note_commitments() { - forked - .sapling_note_commitment_tree + sapling_nct .append(*sapling_note_commitment) .expect("must work since it was already appended before the fork"); } for orchard_note_commitment in transaction.orchard_note_commitments() { - forked - .orchard_note_commitment_tree + orchard_nct .append(*orchard_note_commitment) .expect("must work since it was already appended before the fork"); } @@ -401,11 +402,11 @@ impl Chain { pub fn sapling_tree( &self, hash_or_height: HashOrHeight, - ) -> Option<&sapling::tree::NoteCommitmentTree> { + ) -> Option> { let height = hash_or_height.height_or_else(|hash| self.height_by_hash.get(&hash).cloned())?; - self.sapling_trees_by_height.get(&height) + self.sapling_trees_by_height.get(&height).cloned() } /// Returns the Orchard @@ -414,11 +415,11 @@ impl Chain { pub fn orchard_tree( &self, hash_or_height: HashOrHeight, - ) -> Option<&orchard::tree::NoteCommitmentTree> { + ) -> Option> { let height = hash_or_height.height_or_else(|hash| self.height_by_hash.get(&hash).cloned())?; - self.orchard_trees_by_height.get(&height) + self.orchard_trees_by_height.get(&height).cloned() } /// Returns the block hash of the tip block. @@ -639,9 +640,9 @@ impl Chain { /// Useful when forking, where the trees are rebuilt anyway. fn with_trees( &self, - sprout_note_commitment_tree: sprout::tree::NoteCommitmentTree, - sapling_note_commitment_tree: sapling::tree::NoteCommitmentTree, - orchard_note_commitment_tree: orchard::tree::NoteCommitmentTree, + sprout_note_commitment_tree: Arc, + sapling_note_commitment_tree: Arc, + orchard_note_commitment_tree: Arc, history_tree: HistoryTree, ) -> Self { Chain { @@ -1203,8 +1204,10 @@ impl UpdateWith>> for Chain { joinsplit_data: &Option>, ) -> Result<(), ValidateContextError> { if let Some(joinsplit_data) = joinsplit_data { + let sprout_ncm = Arc::make_mut(&mut self.sprout_note_commitment_tree); + for cm in joinsplit_data.note_commitments() { - self.sprout_note_commitment_tree.append(*cm)?; + sprout_ncm.append(*cm)?; } check::nullifier::add_to_non_finalized_chain_unique( @@ -1245,11 +1248,13 @@ where sapling_shielded_data: &Option>, ) -> Result<(), ValidateContextError> { if let Some(sapling_shielded_data) = sapling_shielded_data { + let sapling_nct = Arc::make_mut(&mut self.sapling_note_commitment_tree); + // The `_u` here indicates that the Sapling note commitment is // specified only by the `u`-coordinate of the Jubjub curve // point `(u, v)`. for cm_u in sapling_shielded_data.note_commitments() { - self.sapling_note_commitment_tree.append(*cm_u)?; + sapling_nct.append(*cm_u)?; } check::nullifier::add_to_non_finalized_chain_unique( @@ -1291,8 +1296,10 @@ impl UpdateWith> for Chain { orchard_shielded_data: &Option, ) -> Result<(), ValidateContextError> { if let Some(orchard_shielded_data) = orchard_shielded_data { + let orchard_nct = Arc::make_mut(&mut self.orchard_note_commitment_tree); + for cm_x in orchard_shielded_data.note_commitments() { - self.orchard_note_commitment_tree.append(*cm_x)?; + orchard_nct.append(*cm_x)?; } check::nullifier::add_to_non_finalized_chain_unique( diff --git a/zebra-state/src/service/read.rs b/zebra-state/src/service/read.rs index 06d182d0..e296bb3a 100644 --- a/zebra-state/src/service/read.rs +++ b/zebra-state/src/service/read.rs @@ -97,7 +97,6 @@ where // we check the most efficient alternative first. (`chain` is always in // memory, but `db` stores transactions on disk, with a memory cache.) chain - .as_ref() .and_then(|chain| { chain .as_ref() @@ -128,9 +127,7 @@ where // state, we check the most efficient alternative first. (`chain` is always // in memory, but `db` stores blocks on disk, with a memory cache.) chain - .as_ref() - .and_then(|chain| chain.as_ref().sapling_tree(hash_or_height).cloned()) - .map(Arc::new) + .and_then(|chain| chain.as_ref().sapling_tree(hash_or_height)) .or_else(|| db.sapling_tree(hash_or_height)) } @@ -155,9 +152,7 @@ where // state, we check the most efficient alternative first. (`chain` is always // in memory, but `db` stores blocks on disk, with a memory cache.) chain - .as_ref() - .and_then(|chain| chain.as_ref().orchard_tree(hash_or_height).cloned()) - .map(Arc::new) + .and_then(|chain| chain.as_ref().orchard_tree(hash_or_height)) .or_else(|| db.orchard_tree(hash_or_height)) }