fix(note-commitment-trees): Populate subtrees (#7636)

* add `sapling_subtree_for_tip` and `orchard_subtree_for_tip` methods to `ZebraDb`

* add methods for non finalized state, move functions

* call `zs_last_key_value` the right way

* fix and simplify `*_subtree_for_tip` methods

Co-authored-by: Arya <aryasolhi@gmail.com>

* apply filter

* rename all tree and subtree methods that use tip

* rename tip tree and subtree methods in non finalized chain

* apply simplify suggestions

Co-authored-by: teor <teor@riseup.net>

---------

Co-authored-by: Arya <aryasolhi@gmail.com>
Co-authored-by: teor <teor@riseup.net>

zebra-state/src/service/finalized_state.rs

@@ -242,8 +242,8 @@ impl FinalizedState {
 
                 let block = checkpoint_verified.block.clone();
                 let mut history_tree = self.db.history_tree();
-                let prev_note_commitment_trees =
+                let prev_note_commitment_trees = prev_note_commitment_trees
-                    prev_note_commitment_trees.unwrap_or_else(|| self.db.note_commitment_trees());
+                    .unwrap_or_else(|| self.db.note_commitment_trees_for_tip());
 
                 // Update the note commitment trees.
                 let mut note_commitment_trees = prev_note_commitment_trees.clone();

zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs

@@ -207,12 +207,13 @@ fn quick_check_sapling_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
         return Ok(());
     }
 
-    let Some(NoteCommitmentSubtreeIndex(tip_subtree_index)) = db.sapling_tree().subtree_index()
+    let Some(NoteCommitmentSubtreeIndex(tip_subtree_index)) =
+        db.sapling_tree_for_tip().subtree_index()
     else {
         return Ok(());
     };
 
-    if tip_subtree_index == 0 && !db.sapling_tree().is_complete_subtree() {
+    if tip_subtree_index == 0 && !db.sapling_tree_for_tip().is_complete_subtree() {
         return Ok(());
     }
 
@@ -260,12 +261,13 @@ fn quick_check_orchard_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
         return Ok(());
     }
 
-    let Some(NoteCommitmentSubtreeIndex(tip_subtree_index)) = db.orchard_tree().subtree_index()
+    let Some(NoteCommitmentSubtreeIndex(tip_subtree_index)) =
+        db.orchard_tree_for_tip().subtree_index()
     else {
         return Ok(());
     };
 
-    if tip_subtree_index == 0 && !db.orchard_tree().is_complete_subtree() {
+    if tip_subtree_index == 0 && !db.orchard_tree_for_tip().is_complete_subtree() {
         return Ok(());
     }
 
@@ -333,13 +335,13 @@ fn check_sapling_subtrees(
     cancel_receiver: &mpsc::Receiver<CancelFormatChange>,
 ) -> Result<Result<(), &'static str>, CancelFormatChange> {
     let Some(NoteCommitmentSubtreeIndex(mut first_incomplete_subtree_index)) =
-        db.sapling_tree().subtree_index()
+        db.sapling_tree_for_tip().subtree_index()
     else {
         return Ok(Ok(()));
     };
 
     // If there are no incomplete subtrees in the tree, also expect a subtree for the final index.
-    if db.sapling_tree().is_complete_subtree() {
+    if db.sapling_tree_for_tip().is_complete_subtree() {
         first_incomplete_subtree_index += 1;
     }
 
@@ -463,13 +465,13 @@ fn check_orchard_subtrees(
     cancel_receiver: &mpsc::Receiver<CancelFormatChange>,
 ) -> Result<Result<(), &'static str>, CancelFormatChange> {
     let Some(NoteCommitmentSubtreeIndex(mut first_incomplete_subtree_index)) =
-        db.orchard_tree().subtree_index()
+        db.orchard_tree_for_tip().subtree_index()
     else {
         return Ok(Ok(()));
     };
 
     // If there are no incomplete subtrees in the tree, also expect a subtree for the final index.
-    if db.orchard_tree().is_complete_subtree() {
+    if db.orchard_tree_for_tip().is_complete_subtree() {
         first_incomplete_subtree_index += 1;
     }
 

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

@@ -247,9 +247,9 @@ fn snapshot_block_and_transaction_data(state: &FinalizedState) {
         let mut stored_sapling_trees = Vec::new();
         let mut stored_orchard_trees = Vec::new();
 
-        let sprout_tree_at_tip = state.sprout_tree();
+        let sprout_tree_at_tip = state.sprout_tree_for_tip();
-        let sapling_tree_at_tip = state.sapling_tree();
+        let sapling_tree_at_tip = state.sapling_tree_for_tip();
-        let orchard_tree_at_tip = state.orchard_tree();
+        let orchard_tree_at_tip = state.orchard_tree_for_tip();
 
         // Test the history tree.
         //

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

@@ -83,7 +83,7 @@ impl ZebraDb {
 
     /// Returns the Sprout note commitment tree of the finalized tip
     /// or the empty tree if the state is empty.
-    pub fn sprout_tree(&self) -> Arc<sprout::tree::NoteCommitmentTree> {
+    pub fn sprout_tree_for_tip(&self) -> Arc<sprout::tree::NoteCommitmentTree> {
         if self.is_empty() {
             return Arc::<sprout::tree::NoteCommitmentTree>::default();
         }
@@ -161,7 +161,7 @@ impl ZebraDb {
 
     /// Returns the Sapling note commitment tree of the finalized tip or the empty tree if the state
     /// is empty.
-    pub fn sapling_tree(&self) -> Arc<sapling::tree::NoteCommitmentTree> {
+    pub fn sapling_tree_for_tip(&self) -> Arc<sapling::tree::NoteCommitmentTree> {
         let height = match self.finalized_tip_height() {
             Some(h) => h,
             None => return Default::default(),
@@ -303,11 +303,32 @@ impl ZebraDb {
         }
     }
 
+    /// Get the sapling note commitment subtress for the finalized tip.
+    #[allow(clippy::unwrap_in_result)]
+    fn sapling_subtree_for_tip(&self) -> Option<NoteCommitmentSubtree<sapling::tree::Node>> {
+        let sapling_subtrees = self
+            .db
+            .cf_handle("sapling_note_commitment_subtree")
+            .unwrap();
+
+        let (index, subtree_data): (
+            NoteCommitmentSubtreeIndex,
+            NoteCommitmentSubtreeData<sapling::tree::Node>,
+        ) = self.db.zs_last_key_value(&sapling_subtrees)?;
+
+        let tip_height = self.finalized_tip_height()?;
+        if subtree_data.end != tip_height {
+            return None;
+        }
+
+        Some(subtree_data.with_index(index))
+    }
+
     // Orchard trees
 
     /// Returns the Orchard note commitment tree of the finalized tip or the empty tree if the state
     /// is empty.
-    pub fn orchard_tree(&self) -> Arc<orchard::tree::NoteCommitmentTree> {
+    pub fn orchard_tree_for_tip(&self) -> Arc<orchard::tree::NoteCommitmentTree> {
         let height = match self.finalized_tip_height() {
             Some(h) => h,
             None => return Default::default(),
@@ -449,15 +470,38 @@ impl ZebraDb {
         }
     }
 
+    /// Get the orchard note commitment subtress for the finalized tip.
+    #[allow(clippy::unwrap_in_result)]
+    fn orchard_subtree_for_tip(&self) -> Option<NoteCommitmentSubtree<orchard::tree::Node>> {
+        let orchard_subtrees = self
+            .db
+            .cf_handle("orchard_note_commitment_subtree")
+            .unwrap();
+
+        let (index, subtree_data): (
+            NoteCommitmentSubtreeIndex,
+            NoteCommitmentSubtreeData<orchard::tree::Node>,
+        ) = self.db.zs_last_key_value(&orchard_subtrees)?;
+
+        let tip_height = self.finalized_tip_height()?;
+        if subtree_data.end != tip_height {
+            return None;
+        }
+
+        Some(subtree_data.with_index(index))
+    }
+
     /// Returns the shielded note commitment trees of the finalized tip
     /// or the empty trees if the state is empty.
-    pub fn note_commitment_trees(&self) -> NoteCommitmentTrees {
+    /// Additionally, returns the sapling and orchard subtrees for the finalized tip if
+    /// the current subtree is finalizing in the tip, None otherwise.
+    pub fn note_commitment_trees_for_tip(&self) -> NoteCommitmentTrees {
         NoteCommitmentTrees {
-            sprout: self.sprout_tree(),
+            sprout: self.sprout_tree_for_tip(),
-            sapling: self.sapling_tree(),
+            sapling: self.sapling_tree_for_tip(),
-            sapling_subtree: None,
+            sapling_subtree: self.sapling_subtree_for_tip(),
-            orchard: self.orchard_tree(),
+            orchard: self.orchard_tree_for_tip(),
-            orchard_subtree: None,
+            orchard_subtree: self.orchard_subtree_for_tip(),
         }
     }
 }
@@ -571,10 +615,10 @@ impl DiskWriteBatch {
 
         // Store the Sapling tree only if it is not already present at the previous height.
         if height.is_min()
-            || prev_note_commitment_trees
+            || prev_note_commitment_trees.as_ref().map_or_else(
-                .as_ref()
+                || zebra_db.sapling_tree_for_tip(),
-                .map_or_else(|| zebra_db.sapling_tree(), |trees| trees.sapling.clone())
+                |trees| trees.sapling.clone(),
-                != trees.sapling
+            ) != trees.sapling
         {
             self.zs_insert(&sapling_tree_cf, height, trees.sapling);
         }
@@ -582,7 +626,7 @@ impl DiskWriteBatch {
         // Store the Orchard tree only if it is not already present at the previous height.
         if height.is_min()
             || prev_note_commitment_trees
-                .map_or_else(|| zebra_db.orchard_tree(), |trees| trees.orchard)
+                .map_or_else(|| zebra_db.orchard_tree_for_tip(), |trees| trees.orchard)
                 != trees.orchard
         {
             self.zs_insert(&orchard_tree_cf, height, trees.orchard);

zebra-state/src/service/non_finalized_state.rs

@@ -284,9 +284,9 @@ impl NonFinalizedState {
         let chain = Chain::new(
             self.network,
             finalized_tip_height,
-            finalized_state.sprout_tree(),
+            finalized_state.sprout_tree_for_tip(),
-            finalized_state.sapling_tree(),
+            finalized_state.sapling_tree_for_tip(),
-            finalized_state.orchard_tree(),
+            finalized_state.orchard_tree_for_tip(),
             finalized_state.history_tree(),
             finalized_state.finalized_value_pool(),
         );

zebra-state/src/service/non_finalized_state/chain.rs

@@ -504,7 +504,7 @@ impl Chain {
     /// # Panics
     ///
     /// If this chain has no sprout trees. (This should be impossible.)
-    pub fn sprout_note_commitment_tree(&self) -> Arc<sprout::tree::NoteCommitmentTree> {
+    pub fn sprout_note_commitment_tree_for_tip(&self) -> Arc<sprout::tree::NoteCommitmentTree> {
         self.sprout_trees_by_height
             .last_key_value()
             .expect("only called while sprout_trees_by_height is populated")
@@ -668,7 +668,7 @@ impl Chain {
     /// # Panics
     ///
     /// If this chain has no sapling trees. (This should be impossible.)
-    pub fn sapling_note_commitment_tree(&self) -> Arc<sapling::tree::NoteCommitmentTree> {
+    pub fn sapling_note_commitment_tree_for_tip(&self) -> Arc<sapling::tree::NoteCommitmentTree> {
         self.sapling_trees_by_height
             .last_key_value()
             .expect("only called while sapling_trees_by_height is populated")
@@ -737,6 +737,16 @@ impl Chain {
             .collect()
     }
 
+    /// Returns the Sapling [`NoteCommitmentSubtree`] if it was completed at the tip height.
+    pub fn sapling_subtree_for_tip(&self) -> Option<NoteCommitmentSubtree<sapling::tree::Node>> {
+        if !self.is_empty() {
+            let tip = self.non_finalized_tip_height();
+            self.sapling_subtree(tip.into())
+        } else {
+            None
+        }
+    }
+
     /// Adds the Sapling `tree` to the tree and anchor indexes at `height`.
     ///
     /// `height` can be either:
@@ -869,7 +879,7 @@ impl Chain {
     /// # Panics
     ///
     /// If this chain has no orchard trees. (This should be impossible.)
-    pub fn orchard_note_commitment_tree(&self) -> Arc<orchard::tree::NoteCommitmentTree> {
+    pub fn orchard_note_commitment_tree_for_tip(&self) -> Arc<orchard::tree::NoteCommitmentTree> {
         self.orchard_trees_by_height
             .last_key_value()
             .expect("only called while orchard_trees_by_height is populated")
@@ -939,6 +949,16 @@ impl Chain {
             .collect()
     }
 
+    /// Returns the Orchard [`NoteCommitmentSubtree`] if it was completed at the tip height.
+    pub fn orchard_subtree_for_tip(&self) -> Option<NoteCommitmentSubtree<orchard::tree::Node>> {
+        if !self.is_empty() {
+            let tip = self.non_finalized_tip_height();
+            self.orchard_subtree(tip.into())
+        } else {
+            None
+        }
+    }
+
     /// Adds the Orchard `tree` to the tree and anchor indexes at `height`.
     ///
     /// `height` can be either:
@@ -1387,11 +1407,11 @@ impl Chain {
 
         // Prepare data for parallel execution
         let mut nct = NoteCommitmentTrees {
-            sprout: self.sprout_note_commitment_tree(),
+            sprout: self.sprout_note_commitment_tree_for_tip(),
-            sapling: self.sapling_note_commitment_tree(),
+            sapling: self.sapling_note_commitment_tree_for_tip(),
-            sapling_subtree: None,
+            sapling_subtree: self.sapling_subtree_for_tip(),
-            orchard: self.orchard_note_commitment_tree(),
+            orchard: self.orchard_note_commitment_tree_for_tip(),
-            orchard_subtree: None,
+            orchard_subtree: self.orchard_subtree_for_tip(),
         };
 
         let mut tree_result = None;
@@ -1427,8 +1447,8 @@ impl Chain {
                 .insert(subtree.index, subtree.into_data());
         }
 
-        let sapling_root = self.sapling_note_commitment_tree().root();
+        let sapling_root = self.sapling_note_commitment_tree_for_tip().root();
-        let orchard_root = self.orchard_note_commitment_tree().root();
+        let orchard_root = self.orchard_note_commitment_tree_for_tip().root();
 
         // TODO: update the history trees in a rayon thread, if they show up in CPU profiles
         let mut history_tree = self.history_block_commitment_tree();

zebra-state/src/service/non_finalized_state/tests/prop.rs

@@ -328,9 +328,9 @@ fn finalized_equals_pushed_genesis() -> Result<()> {
         let mut partial_chain = Chain::new(
             network,
             full_chain.non_finalized_tip_height(),
-            full_chain.sprout_note_commitment_tree(),
+            full_chain.sprout_note_commitment_tree_for_tip(),
-            full_chain.sapling_note_commitment_tree(),
+            full_chain.sapling_note_commitment_tree_for_tip(),
-            full_chain.orchard_note_commitment_tree(),
+            full_chain.orchard_note_commitment_tree_for_tip(),
             full_chain.history_block_commitment_tree(),
             full_chain.chain_value_pools,
         );
@@ -406,9 +406,9 @@ fn finalized_equals_pushed_history_tree() -> Result<()> {
         let mut partial_chain = Chain::new(
             network,
             Height(finalized_count.try_into().unwrap()),
-            full_chain.sprout_note_commitment_tree(),
+            full_chain.sprout_note_commitment_tree_for_tip(),
-            full_chain.sapling_note_commitment_tree(),
+            full_chain.sapling_note_commitment_tree_for_tip(),
-            full_chain.orchard_note_commitment_tree(),
+            full_chain.orchard_note_commitment_tree_for_tip(),
             full_chain.history_block_commitment_tree(),
             full_chain.chain_value_pools,
         );

zebra-state/src/service/non_finalized_state/tests/vectors.rs

@@ -513,8 +513,8 @@ fn history_tree_is_updated_for_network_upgrade(
     let tree = NonEmptyHistoryTree::from_block(
         Network::Mainnet,
         activation_block.clone(),
-        &chain.sapling_note_commitment_tree().root(),
+        &chain.sapling_note_commitment_tree_for_tip().root(),
-        &chain.orchard_note_commitment_tree().root(),
+        &chain.orchard_note_commitment_tree_for_tip().root(),
     )
     .unwrap();
 
@@ -598,8 +598,8 @@ fn commitment_is_validated_for_network_upgrade(network: Network, network_upgrade
     let tree = NonEmptyHistoryTree::from_block(
         Network::Mainnet,
         activation_block.clone(),
-        &chain.sapling_note_commitment_tree().root(),
+        &chain.sapling_note_commitment_tree_for_tip().root(),
-        &chain.orchard_note_commitment_tree().root(),
+        &chain.orchard_note_commitment_tree_for_tip().root(),
     )
     .unwrap();