diff --git a/zebra-state/src/service/non_finalized_state/chain.rs b/zebra-state/src/service/non_finalized_state/chain.rs index 0b62abc5a04..9b1fb88570c 100644 --- a/zebra-state/src/service/non_finalized_state/chain.rs +++ b/zebra-state/src/service/non_finalized_state/chain.rs @@ -22,6 +22,12 @@ pub struct Chain { pub created_utxos: HashMap, pub(super) spent_utxos: HashSet, + sprout_note_commitment_tree: sprout::tree::NoteCommitmentTree, + sapling_note_commitment_tree: sapling::tree::NoteCommitmentTree, + orchard_note_commitment_tree: orchard::tree::NoteCommitmentTree, + + pub(crate) history_tree: HistoryTree, + pub(super) sprout_anchors: HashSet, pub(super) sapling_anchors: HashSet, pub(super) orchard_anchors: HashSet, @@ -34,6 +40,30 @@ pub struct Chain { } impl Chain { + /// Is the internal state of `self` the same as `other`? + /// + /// The history tree must contain the history of the previous (finalized) blocks. + pub fn new(history_tree: HistoryTree) -> Self { + Chain { + blocks: Default::default(), + height_by_hash: Default::default(), + tx_by_hash: Default::default(), + created_utxos: Default::default(), + spent_utxos: Default::default(), + sprout_note_commitment_tree: Default::default(), + sapling_note_commitment_tree: Default::default(), + orchard_note_commitment_tree: Default::default(), + sprout_anchors: Default::default(), + sapling_anchors: Default::default(), + orchard_anchors: Default::default(), + sprout_nullifiers: Default::default(), + sapling_nullifiers: Default::default(), + orchard_nullifiers: Default::default(), + partial_cumulative_work: Default::default(), + history_tree, + } + } + /// Is the internal state of `self` the same as `other`? /// /// [`Chain`] has custom [`Eq`] and [`Ord`] implementations based on proof of work, @@ -122,6 +152,20 @@ impl Chain { forked.pop_tip(); } + // Rebuild the history tree starting from the finalized tip tree. + // TODO: change to a more efficient approach by removing nodes + // from the tree of the original chain (in `pop_tip()`). + // See https://github.com/ZcashFoundation/zebra/issues/2378 + forked + .history_tree + .try_extend(forked.blocks.values().map(|prepared_block| { + ( + prepared_block.block.clone(), + &forked.sapling_note_commitment_tree.root(), + Some(&forked.orchard_note_commitment_tree.root()), + ) + }))?; + Ok(Some(forked)) } @@ -161,6 +205,35 @@ impl Chain { pub fn is_empty(&self) -> bool { self.blocks.is_empty() } + + pub fn history_root_hash(&self) -> ChainHistoryMmrRootHash { + self.history_tree.hash() + } + + /// Clone the Chain but not the history tree, using the history tree + /// specified instead. + /// + /// Useful when forking, where the history tree is rebuilt anyway. + fn with_history_tree(&self, history_tree: HistoryTree) -> Self { + Chain { + blocks: self.blocks.clone(), + height_by_hash: self.height_by_hash.clone(), + tx_by_hash: self.tx_by_hash.clone(), + created_utxos: self.created_utxos.clone(), + spent_utxos: self.spent_utxos.clone(), + sprout_note_commitment_tree: self.sprout_note_commitment_tree.clone(), + sapling_note_commitment_tree: self.sapling_note_commitment_tree.clone(), + orchard_note_commitment_tree: self.orchard_note_commitment_tree.clone(), + sprout_anchors: self.sprout_anchors.clone(), + sapling_anchors: self.sapling_anchors.clone(), + orchard_anchors: self.orchard_anchors.clone(), + sprout_nullifiers: self.sprout_nullifiers.clone(), + sapling_nullifiers: self.sapling_nullifiers.clone(), + orchard_nullifiers: self.orchard_nullifiers.clone(), + partial_cumulative_work: self.partial_cumulative_work, + history_tree, + } + } } /// Helper trait to organize inverse operations done on the `Chain` type. Used to @@ -265,6 +338,21 @@ impl UpdateWith for Chain { self.update_chain_state_with(orchard_shielded_data)?; } + // Having updated all the note commitment trees and nullifier sets in + // this block, the roots of the note commitment trees as of the last + // transaction are the treestates of this block. + let sprout_root = self.sprout_note_commitment_tree.root(); + let sapling_root = self.sapling_note_commitment_tree.root(); + let orchard_root = self.orchard_note_commitment_tree.root(); + + self.sprout_anchors.insert(&sprout_root); + self.sapling_anchors.insert(&sapling_root); + self.orchard_anchors.insert(&orchard_root); + + // TODO: pass Sapling and Orchard roots + self.history_tree + .push(prepared.block.clone(), &sapling_root, Some(&orchard_root))?; + Ok(()) } @@ -401,6 +489,10 @@ impl UpdateWith>> for Chain { joinsplit_data: &Option>, ) -> Result<(), ValidateContextError> { if let Some(joinsplit_data) = joinsplit_data { + for cm in joinsplit_data.note_commitments() { + self.sprout_note_commitment_tree.append(cm); + } + check::nullifier::add_to_non_finalized_chain_unique( &mut self.sprout_nullifiers, joinsplit_data.nullifiers(), @@ -437,6 +529,10 @@ where sapling_shielded_data: &Option>, ) -> Result<(), ValidateContextError> { if let Some(sapling_shielded_data) = sapling_shielded_data { + for cm_u in sapling_shielded_data.note_commitments() { + self.sapling_note_commitment_tree.append(cm_u); + } + for nullifier in sapling_shielded_data.nullifiers() { // TODO: check sapling nullifiers are unique (#2231) self.sapling_nullifiers.insert(*nullifier); @@ -467,6 +563,10 @@ impl UpdateWith> for Chain { orchard_shielded_data: &Option, ) -> Result<(), ValidateContextError> { if let Some(orchard_shielded_data) = orchard_shielded_data { + for cm_x in orchard_shielded_data.note_commitments() { + self.orchard_note_commitment_tree.append(cm_x); + } + // TODO: check orchard nullifiers are unique (#2231) for nullifier in orchard_shielded_data.nullifiers() { self.orchard_nullifiers.insert(*nullifier);