Add value pools to FinalizedState (#2599)

* add value pools to the database

* remove redundant genesis block check

* use update_with_chain_value_pool_change()

* remove constrains

* remove height from the database

* remove calls to chain_value_pool_change

* clippy

* use the "correct" value balances

* bump the database format

* remove everything that is not finalized state

* clippy

* rustfmt

* use all spent utxos

* add new_outputs utxos to all_utxos_spent_by_block

* remove panic

* add finalized state value pool test

* clippy

* clippy 2

* move import

* fix import

* rustfmt

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

zebra-chain/src/value_balance.rs

@@ -140,7 +140,7 @@ where
     /// value pool.
     ///
     /// See `update_with_block` for details.
-    pub(crate) fn update_with_chain_value_pool_change(
+    pub fn update_with_chain_value_pool_change(
         self,
         chain_value_pool_change: ValueBalance<NegativeAllowed>,
     ) -> Result<ValueBalance<C>, ValueBalanceError> {

zebra-consensus/src/block.rs

@@ -203,9 +203,10 @@ where
             metrics::gauge!("zcash.chain.verified.block.height", height.0 as _);
             metrics::counter!("zcash.chain.verified.block.total", 1);
 
-            // Finally, submit the block for contextual verification.
             let new_outputs = Arc::try_unwrap(known_utxos)
                 .expect("all verification tasks using known_utxos are complete");
+
+            // Finally, submit the block for contextual verification.
             let prepared_block = zs::PreparedBlock {
                 block,
                 hash,

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 = 9;
+pub const DATABASE_FORMAT_VERSION: u32 = 10;
 
 /// The maximum number of blocks to check for NU5 transactions,
 /// before we assume we are on a pre-NU5 legacy chain.

zebra-state/src/service/finalized_state.rs

@@ -5,9 +5,10 @@ mod disk_format;
 #[cfg(test)]
 mod tests;
 
-use std::{collections::HashMap, convert::TryInto, path::Path, sync::Arc};
+use std::{borrow::Borrow, collections::HashMap, convert::TryInto, path::Path, sync::Arc};
 
 use zebra_chain::{
+    amount::NonNegative,
     block::{self, Block},
     history_tree::{HistoryTree, NonEmptyHistoryTree},
     orchard,
@@ -15,6 +16,7 @@ use zebra_chain::{
     sapling, sprout,
     transaction::{self, Transaction},
     transparent,
+    value_balance::ValueBalance,
 };
 
 use crate::{BoxError, Config, FinalizedBlock, HashOrHeight};
@@ -64,6 +66,7 @@ impl FinalizedState {
                 db_options.clone(),
             ),
             rocksdb::ColumnFamilyDescriptor::new("history_tree", db_options.clone()),
+            rocksdb::ColumnFamilyDescriptor::new("tip_chain_value_pool", db_options.clone()),
         ];
         let db_result = rocksdb::DB::open_cf_descriptors(&db_options, &path, column_families);
 
@@ -240,6 +243,8 @@ impl FinalizedState {
             self.db.cf_handle("orchard_note_commitment_tree").unwrap();
         let history_tree_cf = self.db.cf_handle("history_tree").unwrap();
 
+        let tip_chain_value_pool = self.db.cf_handle("tip_chain_value_pool").unwrap();
+
         // Assert that callers (including unit tests) get the chain order correct
         if self.is_empty(hash_by_height) {
             assert_eq!(
@@ -310,10 +315,13 @@ impl FinalizedState {
             }
 
             // Index all new transparent outputs
-            for (outpoint, utxo) in new_outputs.into_iter() {
+            for (outpoint, utxo) in new_outputs.borrow().iter() {
                 batch.zs_insert(utxo_by_outpoint, outpoint, utxo);
             }
 
+            // Create a map for all the utxos spent by the block
+            let mut all_utxos_spent_by_block = HashMap::new();
+
             // Index each transaction, spent inputs, nullifiers
             for (transaction_index, (transaction, transaction_hash)) in block
                 .transactions
@@ -329,10 +337,13 @@ impl FinalizedState {
                 };
                 batch.zs_insert(tx_by_hash, transaction_hash, transaction_location);
 
-                // Mark all transparent inputs as spent
+                // Mark all transparent inputs as spent, collect them as well.
                 for input in transaction.inputs() {
                     match input {
                         transparent::Input::PrevOut { outpoint, .. } => {
+                            if let Some(utxo) = self.utxo(outpoint) {
+                                all_utxos_spent_by_block.insert(*outpoint, utxo);
+                            }
                             batch.delete_cf(utxo_by_outpoint, outpoint.as_bytes());
                         }
                         // Coinbase inputs represent new coins,
@@ -390,6 +401,14 @@ impl FinalizedState {
                 batch.zs_insert(history_tree_cf, height, history_tree);
             }
 
+            // Some utxos are spent in the same block so they will be in `new_outputs`.
+            all_utxos_spent_by_block.extend(new_outputs);
+
+            let current_pool = self.current_value_pool();
+            let new_pool =
+                current_pool.update_with_block(block.borrow(), &all_utxos_spent_by_block)?;
+            batch.zs_insert(tip_chain_value_pool, (), new_pool);
+
             Ok(batch)
         };
 
@@ -572,6 +591,14 @@ impl FinalizedState {
     pub fn path(&self) -> &Path {
         self.db.path()
     }
+
+    /// Returns the stored `ValueBalance` for the best chain at the finalized tip height.
+    pub fn current_value_pool(&self) -> ValueBalance<NonNegative> {
+        let value_pool_cf = self.db.cf_handle("tip_chain_value_pool").unwrap();
+        self.db
+            .zs_get(value_pool_cf, &())
+            .unwrap_or_else(ValueBalance::zero)
+    }
 }
 
 // Drop isn't guaranteed to run, such as when we panic, or if someone stored

zebra-state/src/service/tests.rs

@@ -8,6 +8,7 @@ use zebra_chain::{
     parameters::{Network, NetworkUpgrade},
     serialization::{ZcashDeserialize, ZcashDeserializeInto},
     transaction, transparent,
+    value_balance::ValueBalance,
 };
 use zebra_test::{prelude::*, transcript::Transcript};
 
@@ -314,6 +315,36 @@ proptest! {
             prop_assert_eq!(*best_tip_height.borrow(), Some(expected_height));
         }
     }
+
+    /// Test that the value pool is updated accordingly.
+    ///
+    /// 1. Generate a finalized chain and some non-finalized blocks.
+    /// 2. Check that initially the value pool is empty.
+    /// 3. Commit the finalized blocks and check that the value pool is updated accordingly.
+    /// 4. TODO: Commit the non-finalized blocks and check that the value pool is also updated
+    ///    accordingly.
+    #[test]
+    fn value_pool_is_updated(
+        (network, finalized_blocks, _non_finalized_blocks)
+            in continuous_empty_blocks_from_test_vectors(),
+    ) {
+        zebra_test::init();
+
+        let (mut state_service, _) = StateService::new(Config::ephemeral(), network);
+
+        prop_assert_eq!(state_service.disk.current_value_pool(), ValueBalance::zero());
+
+        let mut expected_value_pool = Ok(ValueBalance::zero());
+        for block in finalized_blocks {
+            let utxos = &block.new_outputs;
+            let block_value_pool = &block.block.chain_value_pool_change(utxos)?;
+            expected_value_pool += *block_value_pool;
+
+            state_service.queue_and_commit_finalized(block);
+        }
+
+        prop_assert_eq!(state_service.disk.current_value_pool(), expected_value_pool?.constrain()?);
+    }
 }
 
 /// Test strategy to generate a chain split in two from the test vectors.