diff --git a/zebra-chain/src/history_tree.rs b/zebra-chain/src/history_tree.rs index fcd19fb478a..223d1b372ee 100644 --- a/zebra-chain/src/history_tree.rs +++ b/zebra-chain/src/history_tree.rs @@ -507,3 +507,11 @@ impl Deref for HistoryTree { &self.0 } } + +impl PartialEq for HistoryTree { + fn eq(&self, other: &Self) -> bool { + self.as_ref().map(|tree| tree.hash()) == other.as_ref().map(|other_tree| other_tree.hash()) + } +} + +impl Eq for HistoryTree {} diff --git a/zebra-state/src/service/non_finalized_state/chain.rs b/zebra-state/src/service/non_finalized_state/chain.rs index aabfac839a5..ab833c568d9 100644 --- a/zebra-state/src/service/non_finalized_state/chain.rs +++ b/zebra-state/src/service/non_finalized_state/chain.rs @@ -86,6 +86,7 @@ pub struct Chain { /// 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: Arc, + pub(crate) history_trees_by_height: BTreeMap>, /// The Sprout anchors created by `blocks`. pub(crate) sprout_anchors: MultiSet, @@ -161,6 +162,7 @@ impl Chain { partial_transparent_transfers: Default::default(), partial_cumulative_work: Default::default(), history_tree, + history_trees_by_height: Default::default(), chain_value_pools: finalized_tip_chain_value_pools, } } @@ -190,12 +192,13 @@ impl Chain { self.sprout_note_commitment_tree.root() == other.sprout_note_commitment_tree.root() && self.sprout_trees_by_anchor == other.sprout_trees_by_anchor && self.sapling_note_commitment_tree.root() == other.sapling_note_commitment_tree.root() && - self.sapling_trees_by_height== other.sapling_trees_by_height && + self.sapling_trees_by_height == other.sapling_trees_by_height && self.orchard_note_commitment_tree.root() == other.orchard_note_commitment_tree.root() && - self.orchard_trees_by_height== other.orchard_trees_by_height && + self.orchard_trees_by_height == other.orchard_trees_by_height && - // history tree - self.history_tree.as_ref().as_ref().map(|tree| tree.hash()) == other.history_tree.as_ref().as_ref().map(|other_tree| other_tree.hash()) && + // history trees + self.history_tree == other.history_tree && + self.history_trees_by_height == other.history_trees_by_height && // anchors self.sprout_anchors == other.sprout_anchors && @@ -752,6 +755,7 @@ impl Chain { partial_transparent_transfers: self.partial_transparent_transfers.clone(), partial_cumulative_work: self.partial_cumulative_work, history_tree, + history_trees_by_height: self.history_trees_by_height.clone(), chain_value_pools: self.chain_value_pools, } } @@ -836,6 +840,9 @@ impl Chain { orchard_root, )?; + self.history_trees_by_height + .insert(height, self.history_tree.clone()); + Ok(()) } @@ -1035,6 +1042,11 @@ impl UpdateWith for Chain { // This method is called on two scenarios: // - When popping the root: the history tree does not change. // - When popping the tip: the history tree is rebuilt in fork(). + // + // However, `history_trees_by_height` is reverted. + self.history_trees_by_height + .remove(&height) + .expect("History tree must be present if block was added to chain"); // for each transaction in block for (transaction, transaction_hash) in 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 fdf968acaa6..d9398d7411f 100644 --- a/zebra-state/src/service/non_finalized_state/tests/prop.rs +++ b/zebra-state/src/service/non_finalized_state/tests/prop.rs @@ -613,8 +613,9 @@ fn different_blocks_different_chains() -> Result<()> { chain1.orchard_note_commitment_tree = chain2.orchard_note_commitment_tree.clone(); chain1.orchard_trees_by_height = chain2.orchard_trees_by_height.clone(); - // history tree + // history trees chain1.history_tree = chain2.history_tree.clone(); + chain1.history_trees_by_height = chain2.history_trees_by_height.clone(); // anchors chain1.sprout_anchors = chain2.sprout_anchors.clone();