Refactor value balances to support generated data modification (#2596)

Co-authored-by: Alfredo Garcia <oxarbitrage@gmail.com>

zebra-chain/src/sprout/joinsplit.rs

@@ -3,7 +3,7 @@ use std::io;
 use serde::{Deserialize, Serialize};
 
 use crate::{
-    amount::{Amount, NonNegative},
+    amount::{Amount, NegativeAllowed, NonNegative},
     block::MAX_BLOCK_BYTES,
     primitives::{x25519, Bctv14Proof, Groth16Proof, ZkSnarkProof},
     serialization::{
@@ -70,6 +70,27 @@ impl<P: ZkSnarkProof> ZcashSerialize for JoinSplit<P> {
     }
 }
 
+impl<P: ZkSnarkProof> JoinSplit<P> {
+    /// Return the sprout value balance,
+    /// the change in the transaction value pool due to this sprout [`JoinSplit`].
+    ///
+    /// https://zebra.zfnd.org/dev/rfcs/0012-value-pools.html#definitions
+    ///
+    /// See [`Transaction::sprout_value_balance`] for details.
+    pub fn value_balance(&self) -> Amount<NegativeAllowed> {
+        let vpub_new = self
+            .vpub_new
+            .constrain()
+            .expect("constrain::NegativeAllowed is always valid");
+        let vpub_old = self
+            .vpub_old
+            .constrain()
+            .expect("constrain::NegativeAllowed is always valid");
+
+        (vpub_new - vpub_old).expect("subtraction of two valid amounts is a valid NegativeAllowed")
+    }
+}
+
 impl<P: ZkSnarkProof> ZcashDeserialize for JoinSplit<P> {
     fn zcash_deserialize<R: io::Read>(mut reader: R) -> Result<Self, SerializationError> {
         Ok(JoinSplit::<P> {

zebra-chain/src/transaction.rs

@@ -37,7 +37,7 @@ use crate::{
     value_balance::{ValueBalance, ValueBalanceError},
 };
 
-use std::collections::HashMap;
+use std::{collections::HashMap, iter};
 
 /// A Zcash transaction.
 ///
@@ -927,19 +927,16 @@ impl Transaction {
         }
     }
 
-    /// Return the sprout value balance,
+    /// Return a list of sprout value balances,
-    /// the change in the transaction value pool due to sprout [`JoinSplit`]s.
+    /// the changes in the transaction value pool due to each sprout [`JoinSplit`].
     ///
-    /// The sum of all sprout `vpub_new` fields, minus the sum of all `vpub_old` fields.
+    /// Each value balance is the sprout `vpub_new` field, minus the `vpub_old` field.
     ///
-    /// Positive values are added to this transaction's value pool,
+    /// See `sprout_value_balance` for details.
-    /// and removed from the sprout chain value pool.
+    fn sprout_joinsplit_value_balances(
-    /// Negative values are removed from this transaction,
+        &self,
-    /// and added to the sprout pool.
+    ) -> impl Iterator<Item = ValueBalance<NegativeAllowed>> + '_ {
-    ///
+        let joinsplit_value_balances = match self {
-    /// https://zebra.zfnd.org/dev/rfcs/0012-value-pools.html#definitions
-    fn sprout_value_balance(&self) -> Result<ValueBalance<NegativeAllowed>, ValueBalanceError> {
-        let joinsplit_value_balance = match self {
             Transaction::V2 {
                 joinsplit_data: Some(joinsplit_data),
                 ..
@@ -947,16 +944,11 @@ impl Transaction {
             | Transaction::V3 {
                 joinsplit_data: Some(joinsplit_data),
                 ..
-            } => joinsplit_data
+            } => joinsplit_data.joinsplit_value_balances(),
-                .value_balance()
-                .map_err(ValueBalanceError::Sprout)?,
             Transaction::V4 {
                 joinsplit_data: Some(joinsplit_data),
                 ..
-            } => joinsplit_data
+            } => joinsplit_data.joinsplit_value_balances(),
-                .value_balance()
-                .map_err(ValueBalanceError::Sprout)?,
-
             Transaction::V1 { .. }
             | Transaction::V2 {
                 joinsplit_data: None,
@@ -970,10 +962,25 @@ impl Transaction {
                 joinsplit_data: None,
                 ..
             }
-            | Transaction::V5 { .. } => Amount::zero(),
+            | Transaction::V5 { .. } => Box::new(iter::empty()),
         };
 
-        Ok(ValueBalance::from_sprout_amount(joinsplit_value_balance))
+        joinsplit_value_balances.map(ValueBalance::from_sprout_amount)
+    }
+
+    /// Return the sprout value balance,
+    /// the change in the transaction value pool due to sprout [`JoinSplit`]s.
+    ///
+    /// The sum of all sprout `vpub_new` fields, minus the sum of all `vpub_old` fields.
+    ///
+    /// Positive values are added to this transaction's value pool,
+    /// and removed from the sprout chain value pool.
+    /// Negative values are removed from this transaction,
+    /// and added to the sprout pool.
+    ///
+    /// https://zebra.zfnd.org/dev/rfcs/0012-value-pools.html#definitions
+    fn sprout_value_balance(&self) -> Result<ValueBalance<NegativeAllowed>, ValueBalanceError> {
+        self.sprout_joinsplit_value_balances().sum()
     }
 
     /// Return the sapling value balance,

zebra-chain/src/transaction/joinsplit.rs

@@ -1,7 +1,7 @@
 use serde::{Deserialize, Serialize};
 
 use crate::{
-    amount::{Amount, Error},
+    amount::{self, Amount, NegativeAllowed},
     primitives::{ed25519, ZkSnarkProof},
     sprout::{self, JoinSplit, Nullifier},
 };
@@ -72,10 +72,18 @@ impl<P: ZkSnarkProof> JoinSplitData<P> {
     /// https://zebra.zfnd.org/dev/rfcs/0012-value-pools.html#definitions
     ///
     /// See [`Transaction::sprout_value_balance`] for details.
-    pub fn value_balance(&self) -> Result<Amount, Error> {
+    pub fn value_balance(&self) -> Result<Amount<NegativeAllowed>, amount::Error> {
-        self.joinsplits()
+        self.joinsplit_value_balances().sum()
-            .flat_map(|j| j.vpub_new.constrain() - j.vpub_old.constrain()?)
+    }
-            .sum()
+
+    /// Return a list of sprout value balances,
+    /// the changes in the transaction value pool due to each sprout [`JoinSplit`].
+    ///
+    /// See [`Transaction::sprout_value_balance`] for details.
+    pub fn joinsplit_value_balances(
+        &self,
+    ) -> Box<dyn Iterator<Item = Amount<NegativeAllowed>> + '_> {
+        Box::new(self.joinsplits().map(JoinSplit::value_balance))
     }
 
     /// Collect the Sprout note commitments  for this transaction, if it contains [`Output`]s,

zebra-chain/src/value_balance.rs

@@ -48,26 +48,38 @@ where
         (self.transparent + self.sprout + self.sapling + self.orchard)?.constrain::<NonNegative>()
     }
 
-    /// Update this value balance with the chain value pool change from `block`.
+    /// Returns this value balance, updated with the chain value pool change from `block`.
     ///
     /// `utxos` must contain the [`Utxo`]s of every input in this block,
     /// including UTXOs created by earlier transactions in this block.
     ///
+    /// "If any of the "Sprout chain value pool balance", "Sapling chain value pool balance",
+    /// or "Orchard chain value pool balance" would become negative in the block chain created
+    /// as a result of accepting a block, then all nodes MUST reject the block as invalid.
+    ///
+    /// Nodes MAY relay transactions even if one or more of them cannot be mined due to the
+    /// aforementioned restriction."
+    ///
+    /// https://zips.z.cash/zip-0209#specification
+    ///
+    /// Zebra also checks that the transparent value pool is non-negative,
+    /// which is a consensus rule derived from Bitcoin.
+    ///
     /// Note: the chain value pool has the opposite sign to the transaction
     /// value pool.
     ///
     /// See [`Block::chain_value_pool_change`] for details.
     pub fn update_with_block(
-        &mut self,
+        self,
         block: impl Borrow<Block>,
         utxos: &HashMap<transparent::OutPoint, transparent::Utxo>,
-    ) -> Result<(), ValueBalanceError> {
+    ) -> Result<ValueBalance<C>, ValueBalanceError> {
         let chain_value_pool_change = block.borrow().chain_value_pool_change(utxos)?;
 
         self.update_with_chain_value_pool_change(chain_value_pool_change)
     }
 
-    /// Update this value balance with the chain value pool change from `transaction`.
+    /// Returns this value balance, updated with the chain value pool change from `transaction`.
     ///
     /// `outputs` must contain the [`Output`]s of every input in this transaction,
     /// including UTXOs created by earlier transactions in its block.
@@ -79,10 +91,10 @@ where
     /// for details.
     #[cfg(any(test, feature = "proptest-impl"))]
     pub fn update_with_transaction(
-        &mut self,
+        self,
         transaction: impl Borrow<Transaction>,
         utxos: &HashMap<transparent::OutPoint, transparent::Output>,
-    ) -> Result<(), ValueBalanceError> {
+    ) -> Result<ValueBalance<C>, ValueBalanceError> {
         use std::ops::Neg;
 
         // the chain pool (unspent outputs) has the opposite sign to
@@ -95,23 +107,56 @@ where
         self.update_with_chain_value_pool_change(chain_value_pool_change)
     }
 
-    /// Update this value balance with a chain value pool change.
+    /// Returns this value balance, updated with the chain value pool change from `input`.
+    ///
+    /// `outputs` must contain the [`Output`] spent by `input`,
+    /// (including UTXOs created by earlier transactions in its block).
+    ///
+    /// Note: the chain value pool has the opposite sign to the transaction
+    /// value pool. Inputs remove value from the chain value pool.
+    ///
+    /// See [`Block::chain_value_pool_change`] and [`Transaction::value_balance`]
+    /// for details.
+    #[cfg(any(test, feature = "proptest-impl"))]
+    pub fn update_with_transparent_input(
+        self,
+        input: impl Borrow<transparent::Input>,
+        utxos: &HashMap<transparent::OutPoint, transparent::Output>,
+    ) -> Result<ValueBalance<C>, ValueBalanceError> {
+        use std::ops::Neg;
+
+        // the chain pool (unspent outputs) has the opposite sign to
+        // transaction value balances (inputs - outputs)
+        let transparent_value_pool_change = input.borrow().value_from_outputs(utxos).neg();
+        let transparent_value_pool_change =
+            ValueBalance::from_transparent_amount(transparent_value_pool_change);
+
+        self.update_with_chain_value_pool_change(transparent_value_pool_change)
+    }
+
+    /// Returns this value balance, updated with a chain value pool change.
     ///
     /// Note: the chain value pool has the opposite sign to the transaction
     /// value pool.
     ///
     /// See `update_with_block` for details.
-    fn update_with_chain_value_pool_change(
+    pub(crate) fn update_with_chain_value_pool_change(
-        &mut self,
+        self,
         chain_value_pool_change: ValueBalance<NegativeAllowed>,
-    ) -> Result<(), ValueBalanceError> {
+    ) -> Result<ValueBalance<C>, ValueBalanceError> {
-        let mut current_value_pool = self
+        let mut chain_value_pool = self
             .constrain::<NegativeAllowed>()
             .expect("conversion from NonNegative to NegativeAllowed is always valid");
-        current_value_pool = (current_value_pool + chain_value_pool_change)?;
+        chain_value_pool = (chain_value_pool + chain_value_pool_change)?;
-        *self = current_value_pool.constrain()?;
 
-        Ok(())
+        // Consensus rule: If any of the "Sprout chain value pool balance",
+        // "Sapling chain value pool balance", or "Orchard chain value pool balance"
+        // would become negative in the block chain created as a result of accepting a block,
+        // then all nodes MUST reject the block as invalid.
+        //
+        // Zebra also checks that the transparent value pool is non-negative,
+        // which is a consensus rule derived from Bitcoin.
+        chain_value_pool.constrain()
     }
 
     /// Creates a [`ValueBalance`] from the given transparent amount.