From 1d45938e0f41d0fd82eb5efe10bee8e1a043abdc Mon Sep 17 00:00:00 2001 From: Alfredo Garcia Date: Sun, 8 Oct 2023 23:02:04 -0300 Subject: [PATCH] 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 * 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 --------- Co-authored-by: Arya Co-authored-by: teor --- zebra-state/src/service/finalized_state.rs | 4 +- .../disk_format/upgrade/add_subtrees.rs | 18 ++--- .../zebra_db/block/tests/snapshot.rs | 6 +- .../finalized_state/zebra_db/shielded.rs | 72 +++++++++++++++---- .../src/service/non_finalized_state.rs | 6 +- .../src/service/non_finalized_state/chain.rs | 40 ++++++++--- .../service/non_finalized_state/tests/prop.rs | 12 ++-- .../non_finalized_state/tests/vectors.rs | 8 +-- 8 files changed, 116 insertions(+), 50 deletions(-) diff --git a/zebra-state/src/service/finalized_state.rs b/zebra-state/src/service/finalized_state.rs index 3f4e0d4d..acafda7b 100644 --- a/zebra-state/src/service/finalized_state.rs +++ b/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 = - prev_note_commitment_trees.unwrap_or_else(|| self.db.note_commitment_trees()); + let prev_note_commitment_trees = prev_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(); diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs index 35209b46..72dc4d55 100644 --- a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs +++ b/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, ) -> Result, 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, ) -> Result, 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; } 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 2754cd69..6fc96f8d 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 @@ -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 sapling_tree_at_tip = state.sapling_tree(); - let orchard_tree_at_tip = state.orchard_tree(); + let sprout_tree_at_tip = state.sprout_tree_for_tip(); + let sapling_tree_at_tip = state.sapling_tree_for_tip(); + let orchard_tree_at_tip = state.orchard_tree_for_tip(); // Test the history tree. // 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 fc0cca9d..75b5db8d 100644 --- a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs +++ b/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 { + pub fn sprout_tree_for_tip(&self) -> Arc { if self.is_empty() { return Arc::::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 { + pub fn sapling_tree_for_tip(&self) -> Arc { 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> { + let sapling_subtrees = self + .db + .cf_handle("sapling_note_commitment_subtree") + .unwrap(); + + let (index, subtree_data): ( + NoteCommitmentSubtreeIndex, + NoteCommitmentSubtreeData, + ) = 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 { + pub fn orchard_tree_for_tip(&self) -> Arc { 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> { + let orchard_subtrees = self + .db + .cf_handle("orchard_note_commitment_subtree") + .unwrap(); + + let (index, subtree_data): ( + NoteCommitmentSubtreeIndex, + NoteCommitmentSubtreeData, + ) = 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(), - sapling: self.sapling_tree(), - sapling_subtree: None, - orchard: self.orchard_tree(), - orchard_subtree: None, + sprout: self.sprout_tree_for_tip(), + sapling: self.sapling_tree_for_tip(), + sapling_subtree: self.sapling_subtree_for_tip(), + orchard: self.orchard_tree_for_tip(), + 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 - .as_ref() - .map_or_else(|| zebra_db.sapling_tree(), |trees| trees.sapling.clone()) - != trees.sapling + || prev_note_commitment_trees.as_ref().map_or_else( + || zebra_db.sapling_tree_for_tip(), + |trees| trees.sapling.clone(), + ) != 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); diff --git a/zebra-state/src/service/non_finalized_state.rs b/zebra-state/src/service/non_finalized_state.rs index 6b303360..bc342e5b 100644 --- a/zebra-state/src/service/non_finalized_state.rs +++ b/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.sapling_tree(), - finalized_state.orchard_tree(), + finalized_state.sprout_tree_for_tip(), + finalized_state.sapling_tree_for_tip(), + finalized_state.orchard_tree_for_tip(), finalized_state.history_tree(), finalized_state.finalized_value_pool(), ); diff --git a/zebra-state/src/service/non_finalized_state/chain.rs b/zebra-state/src/service/non_finalized_state/chain.rs index 670901ba..0297c93f 100644 --- a/zebra-state/src/service/non_finalized_state/chain.rs +++ b/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 { + pub fn sprout_note_commitment_tree_for_tip(&self) -> Arc { 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 { + pub fn sapling_note_commitment_tree_for_tip(&self) -> Arc { 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> { + 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 { + pub fn orchard_note_commitment_tree_for_tip(&self) -> Arc { 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> { + 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(), - sapling: self.sapling_note_commitment_tree(), - sapling_subtree: None, - orchard: self.orchard_note_commitment_tree(), - orchard_subtree: None, + sprout: self.sprout_note_commitment_tree_for_tip(), + sapling: self.sapling_note_commitment_tree_for_tip(), + sapling_subtree: self.sapling_subtree_for_tip(), + orchard: self.orchard_note_commitment_tree_for_tip(), + 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 orchard_root = self.orchard_note_commitment_tree().root(); + let sapling_root = self.sapling_note_commitment_tree_for_tip().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(); diff --git a/zebra-state/src/service/non_finalized_state/tests/prop.rs b/zebra-state/src/service/non_finalized_state/tests/prop.rs index 76ebd377..b7a3e5b0 100644 --- a/zebra-state/src/service/non_finalized_state/tests/prop.rs +++ b/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.sapling_note_commitment_tree(), - full_chain.orchard_note_commitment_tree(), + full_chain.sprout_note_commitment_tree_for_tip(), + full_chain.sapling_note_commitment_tree_for_tip(), + 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.sapling_note_commitment_tree(), - full_chain.orchard_note_commitment_tree(), + full_chain.sprout_note_commitment_tree_for_tip(), + full_chain.sapling_note_commitment_tree_for_tip(), + full_chain.orchard_note_commitment_tree_for_tip(), full_chain.history_block_commitment_tree(), full_chain.chain_value_pools, ); diff --git a/zebra-state/src/service/non_finalized_state/tests/vectors.rs b/zebra-state/src/service/non_finalized_state/tests/vectors.rs index 34242be7..78c5de9d 100644 --- a/zebra-state/src/service/non_finalized_state/tests/vectors.rs +++ b/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.orchard_note_commitment_tree().root(), + &chain.sapling_note_commitment_tree_for_tip().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.orchard_note_commitment_tree().root(), + &chain.sapling_note_commitment_tree_for_tip().root(), + &chain.orchard_note_commitment_tree_for_tip().root(), ) .unwrap();