feat(state): Store history trees by height in the non-finalized state (#4928)

* Add history trees for each height in non-fin state

* Refactor formatting

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

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 {}

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<HistoryTree>,
+    pub(crate) history_trees_by_height: BTreeMap<block::Height, Arc<HistoryTree>>,
 
     /// The Sprout anchors created by `blocks`.
     pub(crate) sprout_anchors: MultiSet<sprout::tree::Root>,
@@ -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<ContextuallyValidBlock> 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

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();