Track anchors and note commitment trees in Chain

Append note commitments to their trees when doing update_chain_state_with,
then use the resulting Sapling and Orchard roots to pass to history_tree, and add
new roots to the anchor sets.

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

@@ -22,6 +22,12 @@ pub struct Chain {
     pub created_utxos: HashMap<transparent::OutPoint, transparent::Utxo>,
     pub(super) spent_utxos: HashSet<transparent::OutPoint>,
 
+    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<sprout::tree::Root>,
     pub(super) sapling_anchors: HashSet<sapling::tree::Root>,
     pub(super) orchard_anchors: HashSet<orchard::tree::Root>,
@@ -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<PreparedBlock> 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<Option<transaction::JoinSplitData<Groth16Proof>>> for Chain {
         joinsplit_data: &Option<transaction::JoinSplitData<Groth16Proof>>,
     ) -> 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<sapling::ShieldedData<AnchorV>>,
     ) -> 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<Option<orchard::ShieldedData>> for Chain {
         orchard_shielded_data: &Option<orchard::ShieldedData>,
     ) -> 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);