Wrap Orchard commitment trees into Arc

zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshot.rs

@@ -220,7 +220,7 @@ fn snapshot_block_and_transaction_data(state: &FinalizedState) {
             .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!(*orchard_tree, orchard::tree::NoteCommitmentTree::default());
 
         // Blocks
         let mut stored_block_hashes = Vec::new();

zebra-state/src/service/finalized_state/zebra_db/shielded.rs

@@ -32,7 +32,7 @@ use crate::{
 pub struct NoteCommitmentTrees {
     sprout: Arc<sprout::tree::NoteCommitmentTree>,
     sapling: Arc<sapling::tree::NoteCommitmentTree>,
-    orchard: orchard::tree::NoteCommitmentTree,
+    orchard: Arc<orchard::tree::NoteCommitmentTree>,
 }
 
 impl ZebraDb {
@@ -136,17 +136,17 @@ impl ZebraDb {
 
     /// 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")
     }
 
@@ -156,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
@@ -253,10 +253,9 @@ impl DiskWriteBatch {
             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

@@ -387,7 +387,7 @@ impl NonFinalizedState {
         parent_hash: block::Hash,
         sprout_note_commitment_tree: Arc<sprout::tree::NoteCommitmentTree>,
         sapling_note_commitment_tree: Arc<sapling::tree::NoteCommitmentTree>,
-        orchard_note_commitment_tree: orchard::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

@@ -76,9 +76,10 @@ pub struct Chain {
         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,
@@ -128,7 +129,7 @@ impl Chain {
         network: Network,
         sprout_note_commitment_tree: Arc<sprout::tree::NoteCommitmentTree>,
         sapling_note_commitment_tree: Arc<sapling::tree::NoteCommitmentTree>,
-        orchard_note_commitment_tree: orchard::tree::NoteCommitmentTree,
+        orchard_note_commitment_tree: Arc<orchard::tree::NoteCommitmentTree>,
         history_tree: HistoryTree,
         finalized_tip_chain_value_pools: ValueBalance<NonNegative>,
     ) -> Self {
@@ -274,7 +275,7 @@ impl Chain {
         fork_tip: block::Hash,
         sprout_note_commitment_tree: Arc<sprout::tree::NoteCommitmentTree>,
         sapling_note_commitment_tree: Arc<sapling::tree::NoteCommitmentTree>,
-        orchard_note_commitment_tree: orchard::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) {
@@ -294,6 +295,7 @@ impl Chain {
 
         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.
         for block in forked.blocks.values() {
@@ -311,8 +313,7 @@ impl Chain {
                 }
 
                 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");
                 }
@@ -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.
@@ -641,7 +642,7 @@ impl Chain {
         &self,
         sprout_note_commitment_tree: Arc<sprout::tree::NoteCommitmentTree>,
         sapling_note_commitment_tree: Arc<sapling::tree::NoteCommitmentTree>,
-        orchard_note_commitment_tree: orchard::tree::NoteCommitmentTree,
+        orchard_note_commitment_tree: Arc<orchard::tree::NoteCommitmentTree>,
         history_tree: HistoryTree,
     ) -> Self {
         Chain {
@@ -1295,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

@@ -152,8 +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
-        .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))
 }