Add Sprout anchors to `zebra-state` (#3100)

* Add Sprout anchors to the state

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

Co-authored-by: Deirdre Connolly <deirdre@zfnd.org>

* Return new types of note commitments from Sprout transactions

* Refactor the tests

* Refactor some comments

Co-authored-by: teor <teor@riseup.net>

* Increment `DATABASE_FORMAT_VERSION`

* Update `test.yml` with the new image name

* Refactor the `version = 5` transaction description

Co-authored-by: Deirdre Connolly <deirdre@zfnd.org>

* Update comment

Co-authored-by: Deirdre Connolly <deirdre@zfnd.org>
Co-authored-by: teor <teor@riseup.net>

.github/workflows/test.yml

@@ -54,7 +54,7 @@ jobs:
           --container-image rust:buster \
           --container-mount-disk mount-path='/mainnet',name="zebrad-cache-$SHORT_SHA-mainnet-canopy" \
           --container-restart-policy never \
-          --create-disk name="zebrad-cache-$SHORT_SHA-mainnet-canopy",image=zebrad-cache-0fafd6af-mainnet-canopy \
+          --create-disk name="zebrad-cache-$SHORT_SHA-mainnet-canopy",image=zebrad-cache-13c6a826-mainnet-canopy \
           --machine-type n2-standard-8 \
           --service-account cos-vm@zealous-zebra.iam.gserviceaccount.com \
           --scopes cloud-platform \

zebra-chain/src/transaction.rs

@@ -109,7 +109,7 @@ pub enum Transaction {
         /// The sapling shielded data for this transaction, if any.
         sapling_shielded_data: Option<sapling::ShieldedData<sapling::PerSpendAnchor>>,
     },
-    /// A `version = 5` transaction, which supports `Sapling` and `Orchard`.
+    /// A `version = 5` transaction , which supports Orchard, Sapling, and transparent, but not Sprout.
     V5 {
         /// The Network Upgrade for this transaction.
         ///
@@ -692,18 +692,36 @@ impl Transaction {
         &self,
     ) -> Box<dyn Iterator<Item = &sprout::commitment::NoteCommitment> + '_> {
         match self {
+            // Return [`NoteCommitment`]s with [`Bctv14Proof`]s.
             Transaction::V2 {
                 joinsplit_data: Some(joinsplit_data),
                 ..
+            }
+            | Transaction::V3 {
+                joinsplit_data: Some(joinsplit_data),
+                ..
             } => Box::new(joinsplit_data.note_commitments()),
 
-            Transaction::V1 { .. }
+            // Return [`NoteCommitment`]s with [`Groth16Proof`]s.
-            | Transaction::V2 {
+            Transaction::V4 {
+                joinsplit_data: Some(joinsplit_data),
+                ..
+            } => Box::new(joinsplit_data.note_commitments()),
+
+            // Return an empty iterator.
+            Transaction::V2 {
                 joinsplit_data: None,
                 ..
             }
-            | Transaction::V3 { .. }
+            | Transaction::V3 {
-            | Transaction::V4 { .. }
+                joinsplit_data: None,
+                ..
+            }
+            | Transaction::V4 {
+                joinsplit_data: None,
+                ..
+            }
+            | Transaction::V1 { .. }
             | Transaction::V5 { .. } => Box::new(std::iter::empty()),
         }
     }

zebra-state/src/constants.rs

@@ -18,7 +18,7 @@ pub use zebra_chain::transparent::MIN_TRANSPARENT_COINBASE_MATURITY;
 pub const MAX_BLOCK_REORG_HEIGHT: u32 = MIN_TRANSPARENT_COINBASE_MATURITY - 1;
 
 /// The database format version, incremented each time the database format changes.
-pub const DATABASE_FORMAT_VERSION: u32 = 10;
+pub const DATABASE_FORMAT_VERSION: u32 = 11;
 
 /// The maximum number of blocks to check for NU5 transactions,
 /// before we assume we are on a pre-NU5 legacy chain.

zebra-state/src/error.rs

@@ -214,6 +214,9 @@ pub enum ValidateContextError {
         height: Option<block::Height>,
     },
 
+    #[error("error in Sprout note commitment tree")]
+    SproutNoteCommitmentTreeError(#[from] zebra_chain::sprout::tree::NoteCommitmentTreeError),
+
     #[error("error in Sapling note commitment tree")]
     SaplingNoteCommitmentTreeError(#[from] zebra_chain::sapling::tree::NoteCommitmentTreeError),
 

zebra-state/src/service/finalized_state.rs

@@ -53,6 +53,8 @@ pub struct FinalizedState {
 impl FinalizedState {
     pub fn new(config: &Config, network: Network) -> Self {
         let (path, db_options) = config.db_config(network);
+        // Note: The [`crate::constants::DATABASE_FORMAT_VERSION`] constant must
+        // be incremented each time the database format (column, serialization, etc) changes.
         let column_families = vec![
             rocksdb::ColumnFamilyDescriptor::new("hash_by_height", db_options.clone()),
             rocksdb::ColumnFamilyDescriptor::new("height_by_hash", db_options.clone()),
@@ -62,8 +64,10 @@ impl FinalizedState {
             rocksdb::ColumnFamilyDescriptor::new("sprout_nullifiers", db_options.clone()),
             rocksdb::ColumnFamilyDescriptor::new("sapling_nullifiers", db_options.clone()),
             rocksdb::ColumnFamilyDescriptor::new("orchard_nullifiers", db_options.clone()),
+            rocksdb::ColumnFamilyDescriptor::new("sprout_anchors", db_options.clone()),
             rocksdb::ColumnFamilyDescriptor::new("sapling_anchors", db_options.clone()),
             rocksdb::ColumnFamilyDescriptor::new("orchard_anchors", db_options.clone()),
+            rocksdb::ColumnFamilyDescriptor::new("sprout_note_commitment_tree", db_options.clone()),
             rocksdb::ColumnFamilyDescriptor::new(
                 "sapling_note_commitment_tree",
                 db_options.clone(),
@@ -256,9 +260,12 @@ impl FinalizedState {
         let sapling_nullifiers = self.db.cf_handle("sapling_nullifiers").unwrap();
         let orchard_nullifiers = self.db.cf_handle("orchard_nullifiers").unwrap();
 
+        let sprout_anchors = self.db.cf_handle("sprout_anchors").unwrap();
         let sapling_anchors = self.db.cf_handle("sapling_anchors").unwrap();
         let orchard_anchors = self.db.cf_handle("orchard_anchors").unwrap();
 
+        let sprout_note_commitment_tree_cf =
+            self.db.cf_handle("sprout_note_commitment_tree").unwrap();
         let sapling_note_commitment_tree_cf =
             self.db.cf_handle("sapling_note_commitment_tree").unwrap();
         let orchard_note_commitment_tree_cf =
@@ -298,6 +305,7 @@ impl FinalizedState {
 
         // Read the current note commitment trees. If there are no blocks in the
         // state, these will contain the empty trees.
+        let mut sprout_note_commitment_tree = self.sprout_note_commitment_tree();
         let mut sapling_note_commitment_tree = self.sapling_note_commitment_tree();
         let mut orchard_note_commitment_tree = self.orchard_note_commitment_tree();
         let mut history_tree = self.history_tree();
@@ -344,6 +352,11 @@ impl FinalizedState {
                 // used too early (e.g. the Orchard tree before Nu5 activates)
                 // since the block validation will make sure only appropriate
                 // transactions are allowed in a block.
+                batch.zs_insert(
+                    sprout_note_commitment_tree_cf,
+                    height,
+                    sprout_note_commitment_tree,
+                );
                 batch.zs_insert(
                     sapling_note_commitment_tree_cf,
                     height,
@@ -406,6 +419,9 @@ impl FinalizedState {
                     batch.zs_insert(orchard_nullifiers, orchard_nullifier, ());
                 }
 
+                for sprout_note_commitment in transaction.sprout_note_commitments() {
+                    sprout_note_commitment_tree.append(*sprout_note_commitment)?;
+                }
                 for sapling_note_commitment in transaction.sapling_note_commitments() {
                     sapling_note_commitment_tree.append(*sapling_note_commitment)?;
                 }
@@ -414,6 +430,7 @@ impl FinalizedState {
                 }
             }
 
+            let sprout_root = sprout_note_commitment_tree.root();
             let sapling_root = sapling_note_commitment_tree.root();
             let orchard_root = orchard_note_commitment_tree.root();
 
@@ -421,25 +438,36 @@ impl FinalizedState {
 
             // Compute the new anchors and index them
             // Note: if the root hasn't changed, we write the same value again.
+            batch.zs_insert(sprout_anchors, sprout_root, ());
             batch.zs_insert(sapling_anchors, sapling_root, ());
             batch.zs_insert(orchard_anchors, orchard_root, ());
 
             // Update the trees in state
             if let Some(h) = finalized_tip_height {
+                batch.zs_delete(sprout_note_commitment_tree_cf, h);
                 batch.zs_delete(sapling_note_commitment_tree_cf, h);
                 batch.zs_delete(orchard_note_commitment_tree_cf, h);
                 batch.zs_delete(history_tree_cf, h);
             }
+
+            batch.zs_insert(
+                sprout_note_commitment_tree_cf,
+                height,
+                sprout_note_commitment_tree,
+            );
+
             batch.zs_insert(
                 sapling_note_commitment_tree_cf,
                 height,
                 sapling_note_commitment_tree,
             );
+
             batch.zs_insert(
                 orchard_note_commitment_tree_cf,
                 height,
                 orchard_note_commitment_tree,
             );
+
             if let Some(history_tree) = history_tree.as_ref() {
                 batch.zs_insert(history_tree_cf, height, history_tree);
             }
@@ -624,6 +652,21 @@ impl FinalizedState {
             })
     }
 
+    /// Returns the Sprout note commitment tree of the finalized tip
+    /// or the empty tree if the state is empty.
+    pub fn sprout_note_commitment_tree(&self) -> sprout::tree::NoteCommitmentTree {
+        let height = match self.finalized_tip_height() {
+            Some(h) => h,
+            None => return Default::default(),
+        };
+
+        let sprout_note_commitment_tree = self.db.cf_handle("sprout_note_commitment_tree").unwrap();
+
+        self.db
+            .zs_get(sprout_note_commitment_tree, &height)
+            .expect("Sprout note commitment tree must exist if there is a finalized tip")
+    }
+
     /// Returns the Sapling note commitment tree of the finalized tip
     /// or the empty tree if the state is empty.
     pub fn sapling_note_commitment_tree(&self) -> sapling::tree::NoteCommitmentTree {
@@ -631,11 +674,13 @@ impl FinalizedState {
             Some(h) => h,
             None => return Default::default(),
         };
+
         let sapling_note_commitment_tree =
             self.db.cf_handle("sapling_note_commitment_tree").unwrap();
+
         self.db
             .zs_get(sapling_note_commitment_tree, &height)
-            .expect("note commitment tree must exist if there is a finalized tip")
+            .expect("Sapling note commitment tree must exist if there is a finalized tip")
     }
 
     /// Returns the Orchard note commitment tree of the finalized tip
@@ -645,11 +690,13 @@ impl FinalizedState {
             Some(h) => h,
             None => return Default::default(),
         };
+
         let orchard_note_commitment_tree =
             self.db.cf_handle("orchard_note_commitment_tree").unwrap();
+
         self.db
             .zs_get(orchard_note_commitment_tree, &height)
-            .expect("note commitment tree must exist if there is a finalized tip")
+            .expect("Orchard note commitment tree must exist if there is a finalized tip")
     }
 
     /// Returns the ZIP-221 history tree of the finalized tip or `None`

zebra-state/src/service/finalized_state/disk_format.rs

@@ -239,6 +239,14 @@ impl IntoDisk for transparent::OutPoint {
     }
 }
 
+impl IntoDisk for sprout::tree::Root {
+    type Bytes = [u8; 32];
+
+    fn as_bytes(&self) -> Self::Bytes {
+        self.into()
+    }
+}
+
 impl IntoDisk for sapling::tree::Root {
     type Bytes = [u8; 32];
 
@@ -277,6 +285,23 @@ impl FromDisk for ValueBalance<NonNegative> {
 // in particular to disallow trailing bytes; see
 // https://docs.rs/bincode/1.3.3/bincode/config/index.html#options-struct-vs-bincode-functions
 
+impl IntoDisk for sprout::tree::NoteCommitmentTree {
+    type Bytes = Vec<u8>;
+
+    fn as_bytes(&self) -> Self::Bytes {
+        bincode::DefaultOptions::new()
+            .serialize(self)
+            .expect("serialization to vec doesn't fail")
+    }
+}
+
+impl FromDisk for sprout::tree::NoteCommitmentTree {
+    fn from_bytes(bytes: impl AsRef<[u8]>) -> Self {
+        bincode::DefaultOptions::new()
+            .deserialize(bytes.as_ref())
+            .expect("deserialization format should match the serialization format used by IntoDisk")
+    }
+}
 impl IntoDisk for sapling::tree::NoteCommitmentTree {
     type Bytes = Vec<u8>;
 

zebra-state/src/service/non_finalized_state.rs

@@ -17,14 +17,11 @@ use zebra_chain::{
     history_tree::HistoryTree,
     orchard,
     parameters::Network,
-    sapling,
+    sapling, sprout,
     transaction::{self, Transaction},
     transparent,
 };
 
-#[cfg(test)]
-use zebra_chain::sprout;
-
 use crate::{
     request::ContextuallyValidBlock, FinalizedBlock, HashOrHeight, PreparedBlock,
     ValidateContextError,
@@ -135,6 +132,7 @@ impl NonFinalizedState {
 
         let parent_chain = self.parent_chain(
             parent_hash,
+            finalized_state.sprout_note_commitment_tree(),
             finalized_state.sapling_note_commitment_tree(),
             finalized_state.orchard_note_commitment_tree(),
             finalized_state.history_tree(),
@@ -171,6 +169,7 @@ impl NonFinalizedState {
     ) -> Result<(), ValidateContextError> {
         let chain = Chain::new(
             self.network,
+            finalized_state.sprout_note_commitment_tree(),
             finalized_state.sapling_note_commitment_tree(),
             finalized_state.orchard_note_commitment_tree(),
             finalized_state.history_tree(),
@@ -397,6 +396,7 @@ impl NonFinalizedState {
     fn parent_chain(
         &mut self,
         parent_hash: block::Hash,
+        sprout_note_commitment_tree: sprout::tree::NoteCommitmentTree,
         sapling_note_commitment_tree: sapling::tree::NoteCommitmentTree,
         orchard_note_commitment_tree: orchard::tree::NoteCommitmentTree,
         history_tree: HistoryTree,
@@ -412,6 +412,7 @@ impl NonFinalizedState {
                         chain
                             .fork(
                                 parent_hash,
+                                sprout_note_commitment_tree.clone(),
                                 sapling_note_commitment_tree.clone(),
                                 orchard_note_commitment_tree.clone(),
                                 history_tree.clone(),

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

@@ -1,3 +1,5 @@
+//! Chain that is a part of the non-finalized state.
+
 use std::{
     cmp::Ordering,
     collections::{BTreeMap, HashMap, HashSet},
@@ -25,6 +27,7 @@ use crate::{service::check, ContextuallyValidBlock, ValidateContextError};
 
 #[derive(Debug, Clone)]
 pub struct Chain {
+    // The function `eq_internal_state` must be updated every time a field is added to `Chain`.
     network: Network,
     /// The contextually valid blocks which form this non-finalized partial chain, in height order.
     pub(crate) blocks: BTreeMap<block::Height, ContextuallyValidBlock>,
@@ -43,6 +46,9 @@ pub struct Chain {
     /// including those created by earlier transactions or blocks in the chain.
     pub(crate) spent_utxos: HashSet<transparent::OutPoint>,
 
+    /// The Sprout note commitment tree of the tip of this `Chain`,
+    /// including all finalized notes, and the non-finalized notes in this chain.
+    pub(super) sprout_note_commitment_tree: sprout::tree::NoteCommitmentTree,
     /// The Sapling note commitment tree of the tip of this `Chain`,
     /// including all finalized notes, and the non-finalized notes in this chain.
     pub(super) sapling_note_commitment_tree: sapling::tree::NoteCommitmentTree,
@@ -53,6 +59,10 @@ pub struct Chain {
     /// including all finalized blocks, and the non-finalized `blocks` in this chain.
     pub(crate) history_tree: HistoryTree,
 
+    /// The Sprout anchors created by `blocks`.
+    pub(super) sprout_anchors: HashMultiSet<sprout::tree::Root>,
+    /// The Sprout anchors created by each block in `blocks`.
+    pub(super) sprout_anchors_by_height: BTreeMap<block::Height, sprout::tree::Root>,
     /// The Sapling anchors created by `blocks`.
     pub(super) sapling_anchors: HashMultiSet<sapling::tree::Root>,
     /// The Sapling anchors created by each block in `blocks`.
@@ -89,6 +99,7 @@ impl Chain {
     /// Create a new Chain with the given trees and network.
     pub(crate) fn new(
         network: Network,
+        sprout_note_commitment_tree: sprout::tree::NoteCommitmentTree,
         sapling_note_commitment_tree: sapling::tree::NoteCommitmentTree,
         orchard_note_commitment_tree: orchard::tree::NoteCommitmentTree,
         history_tree: HistoryTree,
@@ -100,9 +111,12 @@ impl Chain {
             height_by_hash: Default::default(),
             tx_by_hash: Default::default(),
             created_utxos: Default::default(),
+            sprout_note_commitment_tree,
             sapling_note_commitment_tree,
             orchard_note_commitment_tree,
             spent_utxos: Default::default(),
+            sprout_anchors: HashMultiSet::new(),
+            sprout_anchors_by_height: Default::default(),
             sapling_anchors: HashMultiSet::new(),
             sapling_anchors_by_height: Default::default(),
             orchard_anchors: HashMultiSet::new(),
@@ -130,8 +144,6 @@ impl Chain {
     pub(crate) fn eq_internal_state(&self, other: &Chain) -> bool {
         use zebra_chain::history_tree::NonEmptyHistoryTree;
 
-        // this method must be updated every time a field is added to Chain
-
         // blocks, heights, hashes
         self.blocks == other.blocks &&
             self.height_by_hash == other.height_by_hash &&
@@ -142,6 +154,7 @@ impl Chain {
             self.spent_utxos == other.spent_utxos &&
 
             // note commitment trees
+            self.sprout_note_commitment_tree.root() == other.sprout_note_commitment_tree.root() &&
             self.sapling_note_commitment_tree.root() == other.sapling_note_commitment_tree.root() &&
             self.orchard_note_commitment_tree.root() == other.orchard_note_commitment_tree.root() &&
 
@@ -149,6 +162,8 @@ impl Chain {
             self.history_tree.as_ref().map(NonEmptyHistoryTree::hash) == other.history_tree.as_ref().map(NonEmptyHistoryTree::hash) &&
 
             // anchors
+            self.sprout_anchors == other.sprout_anchors &&
+            self.sprout_anchors_by_height == other.sprout_anchors_by_height &&
             self.sapling_anchors == other.sapling_anchors &&
             self.sapling_anchors_by_height == other.sapling_anchors_by_height &&
             self.orchard_anchors == other.orchard_anchors &&
@@ -216,6 +231,7 @@ impl Chain {
     pub fn fork(
         &self,
         fork_tip: block::Hash,
+        sprout_note_commitment_tree: sprout::tree::NoteCommitmentTree,
         sapling_note_commitment_tree: sapling::tree::NoteCommitmentTree,
         orchard_note_commitment_tree: orchard::tree::NoteCommitmentTree,
         history_tree: HistoryTree,
@@ -225,6 +241,7 @@ impl Chain {
         }
 
         let mut forked = self.with_trees(
+            sprout_note_commitment_tree,
             sapling_note_commitment_tree,
             orchard_note_commitment_tree,
             history_tree,
@@ -240,12 +257,20 @@ impl Chain {
         // See https://github.com/ZcashFoundation/zebra/issues/2378
         for block in forked.blocks.values() {
             for transaction in block.block.transactions.iter() {
+                for sprout_note_commitment in transaction.sprout_note_commitments() {
+                    forked
+                        .sprout_note_commitment_tree
+                        .append(*sprout_note_commitment)
+                        .expect("must work since it was already appended before the fork");
+                }
+
                 for sapling_note_commitment in transaction.sapling_note_commitments() {
                     forked
                         .sapling_note_commitment_tree
                         .append(*sapling_note_commitment)
                         .expect("must work since it was already appended before the fork");
                 }
+
                 for orchard_note_commitment in transaction.orchard_note_commitments() {
                     forked
                         .orchard_note_commitment_tree
@@ -256,15 +281,16 @@ impl Chain {
 
             // Note that anchors don't need to be recreated since they are already
             // handled in revert_chain_state_with.
-
             let sapling_root = forked
                 .sapling_anchors_by_height
                 .get(&block.height)
                 .expect("Sapling anchors must exist for pre-fork blocks");
+
             let orchard_root = forked
                 .orchard_anchors_by_height
                 .get(&block.height)
                 .expect("Orchard anchors must exist for pre-fork blocks");
+
             forked.history_tree.push(
                 self.network,
                 block.block.clone(),
@@ -339,6 +365,7 @@ impl Chain {
     /// Useful when forking, where the trees are rebuilt anyway.
     fn with_trees(
         &self,
+        sprout_note_commitment_tree: sprout::tree::NoteCommitmentTree,
         sapling_note_commitment_tree: sapling::tree::NoteCommitmentTree,
         orchard_note_commitment_tree: orchard::tree::NoteCommitmentTree,
         history_tree: HistoryTree,
@@ -350,10 +377,13 @@ impl Chain {
             tx_by_hash: self.tx_by_hash.clone(),
             created_utxos: self.created_utxos.clone(),
             spent_utxos: self.spent_utxos.clone(),
+            sprout_note_commitment_tree,
             sapling_note_commitment_tree,
             orchard_note_commitment_tree,
+            sprout_anchors: self.sprout_anchors.clone(),
             sapling_anchors: self.sapling_anchors.clone(),
             orchard_anchors: self.orchard_anchors.clone(),
+            sprout_anchors_by_height: self.sprout_anchors_by_height.clone(),
             sapling_anchors_by_height: self.sapling_anchors_by_height.clone(),
             orchard_anchors_by_height: self.orchard_anchors_by_height.clone(),
             sprout_nullifiers: self.sprout_nullifiers.clone(),
@@ -487,6 +517,9 @@ impl UpdateWith<ContextuallyValidBlock> for Chain {
         // 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();
+        self.sprout_anchors.insert(sprout_root);
+        self.sprout_anchors_by_height.insert(height, sprout_root);
         let sapling_root = self.sapling_note_commitment_tree.root();
         self.sapling_anchors.insert(sapling_root);
         self.sapling_anchors_by_height.insert(height, sapling_root);
@@ -593,6 +626,15 @@ impl UpdateWith<ContextuallyValidBlock> for Chain {
             self.revert_chain_with(orchard_shielded_data, position);
         }
 
+        let anchor = self
+            .sprout_anchors_by_height
+            .remove(&height)
+            .expect("Sprout anchor must be present if block was added to chain");
+        assert!(
+            self.sprout_anchors.remove(&anchor),
+            "Sprout anchor must be present if block was added to chain"
+        );
+
         let anchor = self
             .sapling_anchors_by_height
             .remove(&height)
@@ -601,6 +643,7 @@ impl UpdateWith<ContextuallyValidBlock> for Chain {
             self.sapling_anchors.remove(&anchor),
             "Sapling anchor must be present if block was added to chain"
         );
+
         let anchor = self
             .orchard_anchors_by_height
             .remove(&height)
@@ -673,6 +716,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(),
@@ -711,6 +758,9 @@ where
         sapling_shielded_data: &Option<sapling::ShieldedData<AnchorV>>,
     ) -> Result<(), ValidateContextError> {
         if let Some(sapling_shielded_data) = sapling_shielded_data {
+            // The `_u` here indicates that the Sapling note commitment is
+            // specified only by the `u`-coordinate of the Jubjub curve
+            // point `(u, v)`.
             for cm_u in sapling_shielded_data.note_commitments() {
                 self.sapling_note_commitment_tree.append(*cm_u)?;
             }

zebra-state/src/service/non_finalized_state/tests/prop.rs

@@ -45,7 +45,7 @@ fn push_genesis_chain() -> Result<()> {
         |((chain, count, network, empty_tree) in PreparedChain::default())| {
             prop_assert!(empty_tree.is_none());
 
-            let mut only_chain = Chain::new(network, Default::default(), Default::default(), empty_tree, ValueBalance::zero());
+            let mut only_chain = Chain::new(network, Default::default(), Default::default(), Default::default(), empty_tree, ValueBalance::zero());
             // contains the block value pool changes and chain value pool balances for each height
             let mut chain_values = BTreeMap::new();
 
@@ -97,7 +97,7 @@ fn push_history_tree_chain() -> Result<()> {
         let count = std::cmp::min(count, chain.len() - 1);
         let chain = &chain[1..];
 
-        let mut only_chain = Chain::new(network, Default::default(), Default::default(), finalized_tree, ValueBalance::zero());
+        let mut only_chain = Chain::new(network, Default::default(), Default::default(), Default::default(), finalized_tree, ValueBalance::zero());
 
         for block in chain
             .iter()
@@ -112,13 +112,19 @@ fn push_history_tree_chain() -> Result<()> {
     Ok(())
 }
 
-/// Check that a forked genesis chain is the same as a chain that had the same blocks appended.
+/// Checks that a forked genesis chain is the same as a chain that had the same
+/// blocks appended.
 ///
-/// Also check that:
+/// In other words, this test checks that we get the same chain if we:
-/// - there are no transparent spends in the chain from the genesis block,
+/// - fork the original chain, then push some blocks, or
-///   because genesis transparent outputs are ignored
+/// - push the same blocks to the original chain.
-/// - transactions only spend transparent outputs from earlier in the block or chain
+///
-/// - chain value balances are non-negative
+/// Also checks that:
+/// - There are no transparent spends in the chain from the genesis block,
+///   because genesis transparent outputs are ignored.
+/// - Transactions only spend transparent outputs from earlier in the block or
+///   chain.
+/// - Chain value balances are non-negative.
 #[test]
 fn forked_equals_pushed_genesis() -> Result<()> {
     zebra_test::init();
@@ -129,73 +135,90 @@ fn forked_equals_pushed_genesis() -> Result<()> {
                                .and_then(|v| v.parse().ok())
                                .unwrap_or(DEFAULT_PARTIAL_CHAIN_PROPTEST_CASES)),
     |((chain, fork_at_count, network, empty_tree) in PreparedChain::default())| {
-
         prop_assert!(empty_tree.is_none());
 
-        // use `fork_at_count` as the fork tip
+        // This chain will be used to check if the blocks in the forked chain
-        let fork_tip_hash = chain[fork_at_count - 1].hash;
+        // correspond to the blocks in the original chain before the fork.
-
+        let mut partial_chain = Chain::new(
-        let mut full_chain = Chain::new(network, Default::default(), Default::default(), empty_tree.clone(), ValueBalance::zero());
+            network,
-        let mut partial_chain = Chain::new(network, Default::default(), Default::default(), empty_tree.clone(), ValueBalance::zero());
+            Default::default(),
-
+            Default::default(),
+            Default::default(),
+            empty_tree.clone(),
+            ValueBalance::zero(),
+        );
         for block in chain.iter().take(fork_at_count).cloned() {
-            let block =
+            let block = ContextuallyValidBlock::with_block_and_spent_utxos(
-                ContextuallyValidBlock::with_block_and_spent_utxos(
+                block,
-                    block,
+                partial_chain.unspent_utxos(),
-                    partial_chain.unspent_utxos(),
+            )?;
-                )?;
+            partial_chain = partial_chain
-            partial_chain = partial_chain.push(block).expect("partial chain push is valid");
+                .push(block)
+                .expect("partial chain push is valid");
         }
 
+        // This chain will be forked.
+        let mut full_chain = Chain::new(
+            network,
+            Default::default(),
+            Default::default(),
+            Default::default(),
+            empty_tree.clone(),
+            ValueBalance::zero(),
+        );
         for block in chain.iter().cloned() {
             let block =
-                ContextuallyValidBlock::with_block_and_spent_utxos(
+                ContextuallyValidBlock::with_block_and_spent_utxos(block, full_chain.unspent_utxos())?;
-                    block,
+            full_chain = full_chain
-                    full_chain.unspent_utxos(),
+                .push(block.clone())
-                )?;
+                .expect("full chain push is valid");
-                full_chain = full_chain.push(block.clone()).expect("full chain push is valid");
 
-                // check some other properties of generated chains
+            // Check some other properties of generated chains.
-                if block.height == block::Height(0) {
+            if block.height == block::Height(0) {
-                    prop_assert_eq!(
+                prop_assert_eq!(
-                        block
+                    block
-                            .block
+                        .block
-                            .transactions
+                        .transactions
-                            .iter()
+                        .iter()
-                            .flat_map(|t| t.inputs())
+                        .flat_map(|t| t.inputs())
-                            .filter_map(|i| i.outpoint())
+                        .filter_map(|i| i.outpoint())
-                            .count(),
+                        .count(),
-                        0,
+                    0,
-                        "unexpected transparent prevout input at height {:?}: \
+                    "unexpected transparent prevout input at height {:?}: \
-                         genesis transparent outputs must be ignored, \
+                            genesis transparent outputs must be ignored, \
-                         so there can not be any spends in the genesis block",
+                            so there can not be any spends in the genesis block",
-                        block.height,
+                    block.height,
-                    );
+                );
-                }
             }
+        }
 
+        // Use [`fork_at_count`] as the fork tip.
+        let fork_tip_hash = chain[fork_at_count - 1].hash;
+
+        // Fork the chain.
         let mut forked = full_chain
             .fork(
                 fork_tip_hash,
                 Default::default(),
                 Default::default(),
+                Default::default(),
                 empty_tree,
             )
             .expect("fork works")
             .expect("hash is present");
 
-        // the first check is redundant, but it's useful for debugging
+        // This check is redundant, but it's useful for debugging.
         prop_assert_eq!(forked.blocks.len(), partial_chain.blocks.len());
+
+        // Check that the entire internal state of the forked chain corresponds to the state of
+        // the original chain.
         prop_assert!(forked.eq_internal_state(&partial_chain));
 
         // Re-add blocks to the fork and check if we arrive at the
-        // same original full chain
+        // same original full chain.
         for block in chain.iter().skip(fork_at_count).cloned() {
             let block =
-                ContextuallyValidBlock::with_block_and_spent_utxos(
+                ContextuallyValidBlock::with_block_and_spent_utxos(block, forked.unspent_utxos())?;
-                    block,
-                    forked.unspent_utxos(),
-                )?;
             forked = forked.push(block).expect("forked chain push is valid");
         }
 
@@ -229,8 +252,8 @@ fn forked_equals_pushed_history_tree() -> Result<()> {
         // use `fork_at_count` as the fork tip
         let fork_tip_hash = chain[fork_at_count - 1].hash;
 
-        let mut full_chain = Chain::new(network, Default::default(), Default::default(), finalized_tree.clone(), ValueBalance::zero());
+        let mut full_chain = Chain::new(network, Default::default(), Default::default(), Default::default(), finalized_tree.clone(), ValueBalance::zero());
-        let mut partial_chain = Chain::new(network, Default::default(), Default::default(), finalized_tree.clone(), ValueBalance::zero());
+        let mut partial_chain = Chain::new(network, Default::default(), Default::default(), Default::default(), finalized_tree.clone(), ValueBalance::zero());
 
         for block in chain
             .iter()
@@ -250,6 +273,7 @@ fn forked_equals_pushed_history_tree() -> Result<()> {
                 fork_tip_hash,
                 Default::default(),
                 Default::default(),
+                Default::default(),
                 finalized_tree,
             )
             .expect("fork works")
@@ -297,7 +321,7 @@ fn finalized_equals_pushed_genesis() -> Result<()> {
 
         let fake_value_pool = ValueBalance::<NonNegative>::fake_populated_pool();
 
-        let mut full_chain = Chain::new(network, Default::default(), Default::default(), empty_tree, fake_value_pool);
+        let mut full_chain = Chain::new(network, Default::default(), Default::default(), Default::default(), empty_tree, fake_value_pool);
         for block in chain
             .iter()
             .take(finalized_count)
@@ -307,6 +331,7 @@ fn finalized_equals_pushed_genesis() -> Result<()> {
 
         let mut partial_chain = Chain::new(
             network,
+            full_chain.sprout_note_commitment_tree.clone(),
             full_chain.sapling_note_commitment_tree.clone(),
             full_chain.orchard_note_commitment_tree.clone(),
             full_chain.history_tree.clone(),
@@ -366,7 +391,7 @@ fn finalized_equals_pushed_history_tree() -> Result<()> {
 
         let fake_value_pool = ValueBalance::<NonNegative>::fake_populated_pool();
 
-        let mut full_chain = Chain::new(network, Default::default(), Default::default(), finalized_tree, fake_value_pool);
+        let mut full_chain = Chain::new(network, Default::default(), Default::default(), Default::default(), finalized_tree, fake_value_pool);
         for block in chain
             .iter()
             .take(finalized_count)
@@ -376,11 +401,13 @@ fn finalized_equals_pushed_history_tree() -> Result<()> {
 
         let mut partial_chain = Chain::new(
             network,
+            full_chain.sprout_note_commitment_tree.clone(),
             full_chain.sapling_note_commitment_tree.clone(),
             full_chain.orchard_note_commitment_tree.clone(),
             full_chain.history_tree.clone(),
             full_chain.chain_value_pools,
         );
+
         for block in chain
             .iter()
             .skip(finalized_count)
@@ -528,8 +555,8 @@ fn different_blocks_different_chains() -> Result<()> {
         } else {
             Default::default()
         };
-        let chain1 = Chain::new(Network::Mainnet, Default::default(), Default::default(), finalized_tree1, ValueBalance::fake_populated_pool());
+        let chain1 = Chain::new(Network::Mainnet, Default::default(), Default::default(), Default::default(), finalized_tree1, ValueBalance::fake_populated_pool());
-        let chain2 = Chain::new(Network::Mainnet, Default::default(), Default::default(), finalized_tree2, ValueBalance::fake_populated_pool());
+        let chain2 = Chain::new(Network::Mainnet, Default::default(), Default::default(), Default::default(), finalized_tree2, ValueBalance::fake_populated_pool());
 
         let block1 = vec1[1].clone().prepare().test_with_zero_spent_utxos();
         let block2 = vec2[1].clone().prepare().test_with_zero_spent_utxos();
@@ -564,6 +591,7 @@ fn different_blocks_different_chains() -> Result<()> {
                 chain1.spent_utxos = chain2.spent_utxos.clone();
 
                 // note commitment trees
+                chain1.sprout_note_commitment_tree = chain2.sprout_note_commitment_tree.clone();
                 chain1.sapling_note_commitment_tree = chain2.sapling_note_commitment_tree.clone();
                 chain1.orchard_note_commitment_tree = chain2.orchard_note_commitment_tree.clone();
 
@@ -571,6 +599,8 @@ fn different_blocks_different_chains() -> Result<()> {
                 chain1.history_tree = chain2.history_tree.clone();
 
                 // anchors
+                chain1.sprout_anchors = chain2.sprout_anchors.clone();
+                chain1.sprout_anchors_by_height = chain2.sprout_anchors_by_height.clone();
                 chain1.sapling_anchors = chain2.sapling_anchors.clone();
                 chain1.sapling_anchors_by_height = chain2.sapling_anchors_by_height.clone();
                 chain1.orchard_anchors = chain2.orchard_anchors.clone();

zebra-state/src/service/non_finalized_state/tests/vectors.rs

@@ -28,6 +28,7 @@ fn construct_empty() {
         Default::default(),
         Default::default(),
         Default::default(),
+        Default::default(),
         ValueBalance::zero(),
     );
 }
@@ -43,8 +44,10 @@ fn construct_single() -> Result<()> {
         Default::default(),
         Default::default(),
         Default::default(),
+        Default::default(),
         ValueBalance::fake_populated_pool(),
     );
+
     chain = chain.push(block.prepare().test_with_zero_spent_utxos())?;
 
     assert_eq!(1, chain.blocks.len());
@@ -71,6 +74,7 @@ fn construct_many() -> Result<()> {
         Default::default(),
         Default::default(),
         Default::default(),
+        Default::default(),
         ValueBalance::fake_populated_pool(),
     );
 
@@ -96,6 +100,7 @@ fn ord_matches_work() -> Result<()> {
         Default::default(),
         Default::default(),
         Default::default(),
+        Default::default(),
         ValueBalance::fake_populated_pool(),
     );
     lesser_chain = lesser_chain.push(less_block.prepare().test_with_zero_spent_utxos())?;
@@ -105,6 +110,7 @@ fn ord_matches_work() -> Result<()> {
         Default::default(),
         Default::default(),
         Default::default(),
+        Default::default(),
         ValueBalance::zero(),
     );
     bigger_chain = bigger_chain.push(more_block.prepare().test_with_zero_spent_utxos())?;