change(state): Wrap commitment trees into `Arc` (#4757)

* Wrap Sprout note commitment trees into `Arc`

* Remove a redundant comment

* Rephrase a comment about chain forking

* Remove a redundant comment

The comment is not valid because Zebra uses `bridgetree::Frontier`s from
the `incrementalmerkletree` crate to represent its note commitment
trees. This `struct` does not support popping elements from the tree.

* Wrap Sapling commitment trees into `Arc`

* Remove unnecessary `as_ref`s

* Wrap Orchard commitment trees into `Arc`

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<sprout::tree::NoteCommitmentTree>,
             > = 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");
                 }

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

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<sprout::tree::NoteCommitmentTree>,
+    sapling: Arc<sapling::tree::NoteCommitmentTree>,
+    orchard: Arc<orchard::tree::NoteCommitmentTree>,
 }
 
 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<sprout::tree::NoteCommitmentTree> {
         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<sprout::tree::NoteCommitmentTree> {
-        let sprout_anchors = self.db.cf_handle("sprout_anchors").unwrap();
+    ) -> Option<Arc<sprout::tree::NoteCommitmentTree>> {
+        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<sapling::tree::NoteCommitmentTree> {
         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<sapling::tree::NoteCommitmentTree> {
+    ) -> Option<Arc<sapling::tree::NoteCommitmentTree>> {
         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<orchard::tree::NoteCommitmentTree> {
         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<orchard::tree::NoteCommitmentTree> {
+    ) -> Option<Arc<orchard::tree::NoteCommitmentTree>> {
         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(())

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<sprout::tree::NoteCommitmentTree>,
+        sapling_note_commitment_tree: Arc<sapling::tree::NoteCommitmentTree>,
+        orchard_note_commitment_tree: Arc<orchard::tree::NoteCommitmentTree>,
         history_tree: HistoryTree,
     ) -> Result<Arc<Chain>, ValidateContextError> {
         match self.find_chain(|chain| chain.non_finalized_tip_hash() == parent_hash) {

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<sprout::tree::NoteCommitmentTree>,
     /// The Sprout note commitment tree for each anchor.
     /// This is required for interstitial states.
     pub(crate) sprout_trees_by_anchor:
-        HashMap<sprout::tree::Root, sprout::tree::NoteCommitmentTree>,
+        HashMap<sprout::tree::Root, 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: sapling::tree::NoteCommitmentTree,
+    pub(super) sapling_note_commitment_tree: Arc<sapling::tree::NoteCommitmentTree>,
     /// The Sapling note commitment tree for each height.
-    pub(crate) sapling_trees_by_height: BTreeMap<block::Height, sapling::tree::NoteCommitmentTree>,
+    pub(crate) sapling_trees_by_height:
+        BTreeMap<block::Height, Arc<sapling::tree::NoteCommitmentTree>>,
     /// 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<orchard::tree::NoteCommitmentTree>,
     /// The Orchard note commitment tree for each height.
-    pub(crate) orchard_trees_by_height: BTreeMap<block::Height, orchard::tree::NoteCommitmentTree>,
+    pub(crate) orchard_trees_by_height:
+        BTreeMap<block::Height, Arc<orchard::tree::NoteCommitmentTree>>,
     /// 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<sprout::tree::NoteCommitmentTree>,
+        sapling_note_commitment_tree: Arc<sapling::tree::NoteCommitmentTree>,
+        orchard_note_commitment_tree: Arc<orchard::tree::NoteCommitmentTree>,
         history_tree: HistoryTree,
         finalized_tip_chain_value_pools: ValueBalance<NonNegative>,
     ) -> 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<sprout::tree::NoteCommitmentTree>,
+        sapling_note_commitment_tree: Arc<sapling::tree::NoteCommitmentTree>,
+        orchard_note_commitment_tree: Arc<orchard::tree::NoteCommitmentTree>,
         history_tree: HistoryTree,
     ) -> Result<Option<Self>, 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<Arc<sapling::tree::NoteCommitmentTree>> {
         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<Arc<orchard::tree::NoteCommitmentTree>> {
         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<sprout::tree::NoteCommitmentTree>,
+        sapling_note_commitment_tree: Arc<sapling::tree::NoteCommitmentTree>,
+        orchard_note_commitment_tree: Arc<orchard::tree::NoteCommitmentTree>,
         history_tree: HistoryTree,
     ) -> Self {
         Chain {
@@ -1203,8 +1204,10 @@ impl UpdateWith<Option<transaction::JoinSplitData<Groth16Proof>>> for Chain {
         joinsplit_data: &Option<transaction::JoinSplitData<Groth16Proof>>,
     ) -> 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<sapling::ShieldedData<AnchorV>>,
     ) -> 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<Option<orchard::ShieldedData>> for Chain {
         orchard_shielded_data: &Option<orchard::ShieldedData>,
     ) -> 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(

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))
 }