[zcash#1411] Refactor AccountBalance to use Balance for transparent f…

…unds

closes zcash#1411

- Adds min_confirmations to transparent account balances
- Adds a new `transparent_balance_spendability` test to verify
transparent funds spendability.
- cargo clippy + cargo fmt

Refactor: Extract Method on lambda `check_balance`

zcash_client_backend/CHANGELOG.md

@@ -32,6 +32,9 @@ and this library adheres to Rust's notion of
 - MSRV is now 1.77.0.
 - Migrated to `arti-client 0.23`.
 - `zcash_client_backend::data_api`:
+  - `AccountBalance` was refactored to use `Balance` for transparent funds (issue #1411).
+    `AccountBalance` now has an `unshielded_balance()` that uses `Balance` and replaces the
+    (now deleted) `unshielded` non-negative amount.
   - `InputSource` has an added method `get_wallet_metadata`
   - `error::Error` has additional variant `Error::Change`. This necessitates
     the addition of two type parameters to the `Error` type,
@@ -86,6 +89,15 @@ and this library adheres to Rust's notion of
 
 ### Removed
 - `zcash_client_backend::data_api`:
+  - `AccountBalance::unshielded`. `AccountBalance` no longer provides the `unshielded`
+    method returning a `NonNegativeAmount` for the total of transparent funds.
+    Instead use `unshielded_balance` which provides a `Balance` value.
+  - `zcash_client_backend::AccountBalance::add_unshielded_value`. Instead use
+    `AccountBalance::with_unshielded_balance_mut` with a closure that calls
+    the appropriate `add_*_value` method(s) of `Balance` on its argument.
+    Note that the appropriate method(s) depend on whether the funds are
+    spendable, pending change, or pending non-change (previously, only the
+    total unshielded value was tracked).
   - `WalletSummary::scan_progress` and `WalletSummary::recovery_progress` have
     been removed. Use `WalletSummary::progress` instead.
   - `testing::input_selector` use explicit `InputSelector` constructors

zcash_client_backend/src/data_api.rs

@@ -231,27 +231,23 @@ pub struct AccountBalance {
     /// The value of unspent Orchard outputs belonging to the account.
     orchard_balance: Balance,
 
-    /// The value of all unspent transparent outputs belonging to the account, irrespective of
-    /// confirmation depth.
-    ///
-    /// Unshielded balances are not subject to confirmation-depth constraints, because the only
-    /// possible operation on a transparent balance is to shield it, it is possible to create a
-    /// zero-conf transaction to perform that shielding, and the resulting shielded notes will be
-    /// subject to normal confirmation rules.
-    unshielded: NonNegativeAmount,
+    /// The value of all unspent transparent outputs belonging to the account.
+    unshielded_balance: Balance,
 }
 
 impl AccountBalance {
     /// The [`Balance`] value having zero values for all its fields.
     pub const ZERO: Self = Self {
         sapling_balance: Balance::ZERO,
         orchard_balance: Balance::ZERO,
-        unshielded: NonNegativeAmount::ZERO,
+        unshielded_balance: Balance::ZERO,
     };
 
     fn check_total(&self) -> Result<NonNegativeAmount, BalanceError> {
-        (self.sapling_balance.total() + self.orchard_balance.total() + self.unshielded)
-            .ok_or(BalanceError::Overflow)
+        (self.sapling_balance.total()
+            + self.orchard_balance.total()
+            + self.unshielded_balance.total())
+        .ok_or(BalanceError::Overflow)
     }
 
     /// Returns the [`Balance`] of Sapling funds in the account.
@@ -289,22 +285,36 @@ impl AccountBalance {
     }
 
     /// Returns the total value of unspent transparent transaction outputs belonging to the wallet.
+    #[deprecated(
+        note = "this function is deprecated. Please use [`AccountBalance::unshielded_balance`] instead."
+    )]
     pub fn unshielded(&self) -> NonNegativeAmount {
-        self.unshielded
+        self.unshielded_balance.total()
+    }
+
+    /// Returns the [`Balance`] of unshielded funds in the account.
+    pub fn unshielded_balance(&self) -> &Balance {
+        &self.unshielded_balance
     }
 
-    /// Adds the specified value to the unshielded total, checking for overflow of
-    /// the total account balance.
-    pub fn add_unshielded_value(&mut self, value: NonNegativeAmount) -> Result<(), BalanceError> {
-        self.unshielded = (self.unshielded + value).ok_or(BalanceError::Overflow)?;
+    /// Provides a `mutable reference to the [`Balance`] of transparent funds in the account
+    /// to the specified callback, checking invariants after the callback's action has been
+    /// evaluated.
+    pub fn with_unshielded_balance_mut<A, E: From<BalanceError>>(
+        &mut self,
+        f: impl FnOnce(&mut Balance) -> Result<A, E>,
+    ) -> Result<A, E> {
+        let result = f(&mut self.unshielded_balance)?;
         self.check_total()?;
-        Ok(())
+        Ok(result)
     }
 
     /// Returns the total value of funds belonging to the account.
     pub fn total(&self) -> NonNegativeAmount {
-        (self.sapling_balance.total() + self.orchard_balance.total() + self.unshielded)
-            .expect("Account balance cannot overflow MAX_MONEY")
+        (self.sapling_balance.total()
+            + self.orchard_balance.total()
+            + self.unshielded_balance.total())
+        .expect("Account balance cannot overflow MAX_MONEY")
     }
 
     /// Returns the total value of shielded (Sapling and Orchard) funds that may immediately be
@@ -1028,8 +1038,9 @@ pub trait WalletRead {
     /// or `Ok(None)` if the wallet has no initialized accounts.
     fn get_wallet_birthday(&self) -> Result<Option<BlockHeight>, Self::Error>;
 
-    /// Returns the wallet balances and sync status for an account given the specified minimum
-    /// number of confirmations, or `Ok(None)` if the wallet has no balance data available.
+    /// Returns a [`WalletSummary`] that represents the sync status, and the wallet balances
+    /// given the specified minimum number of confirmations for all accounts known to the
+    /// wallet; or `Ok(None)` if the wallet has no summary data available.
     fn get_wallet_summary(
         &self,
         min_confirmations: u32,

zcash_client_backend/src/data_api/testing/transparent.rs

@@ -4,7 +4,7 @@ use crate::{
             AddressType, DataStoreFactory, ShieldedProtocol, TestBuilder, TestCache, TestState,
         },
         wallet::input_selection::GreedyInputSelector,
-        Account as _, InputSource, WalletRead, WalletWrite,
+        Account as _, Balance, InputSource, WalletRead, WalletWrite,
     },
     fees::{standard, DustOutputPolicy, StandardFeeRule},
     wallet::WalletTransparentOutput,
@@ -13,8 +13,67 @@ use assert_matches::assert_matches;
 use sapling::zip32::ExtendedSpendingKey;
 use zcash_primitives::{
     block::BlockHash,
+    legacy::TransparentAddress,
     transaction::components::{amount::NonNegativeAmount, OutPoint, TxOut},
 };
+use zcash_protocol::local_consensus::LocalNetwork;
+
+use super::TestAccount;
+
+fn check_balance<DSF>(
+    st: &TestState<impl TestCache, <DSF as DataStoreFactory>::DataStore, LocalNetwork>,
+    account: &TestAccount<<DSF as DataStoreFactory>::Account>,
+    taddr: &TransparentAddress,
+    min_confirmations: u32,
+    expected: &Balance,
+) where
+    DSF: DataStoreFactory,
+{
+    // Check the wallet summary returns the expected transparent balance.
+    let summary = st
+        .wallet()
+        .get_wallet_summary(min_confirmations)
+        .unwrap()
+        .unwrap();
+    let balance = summary.account_balances().get(&account.id()).unwrap();
+
+    #[allow(deprecated)]
+    let old_unshielded_value = balance.unshielded();
+    assert_eq!(old_unshielded_value, expected.total());
+    assert_eq!(balance.unshielded_balance(), expected);
+
+    // Check the older APIs for consistency.
+    let mempool_height = st.wallet().chain_height().unwrap().unwrap() + 1;
+    assert_eq!(
+        st.wallet()
+            .get_transparent_balances(account.id(), mempool_height)
+            .unwrap()
+            .get(taddr)
+            .cloned()
+            .unwrap_or(NonNegativeAmount::ZERO),
+        expected.total(),
+    );
+    assert_eq!(
+        st.wallet()
+            .get_spendable_transparent_outputs(taddr, mempool_height, 0)
+            .unwrap()
+            .into_iter()
+            .map(|utxo| utxo.value())
+            .sum::<Option<NonNegativeAmount>>(),
+        Some(expected.spendable_value()),
+    );
+
+    // we currently treat min_confirmation the same regardless they are 0 or 1
+    assert_eq!(
+        st.wallet()
+            .get_spendable_transparent_outputs(taddr, mempool_height, 1)
+            .unwrap()
+            .into_iter()
+            .map(|utxo| utxo.value())
+            .sum::<Option<NonNegativeAmount>>(),
+        Some(expected.spendable_value()),
+    );
+}
 
 pub fn put_received_transparent_utxo<DSF>(dsf: DSF)
 where
@@ -136,46 +195,10 @@ where
     }
     st.scan_cached_blocks(start_height, 10);
 
-    let check_balance = |st: &TestState<_, DSF::DataStore, _>, min_confirmations: u32, expected| {
-        // Check the wallet summary returns the expected transparent balance.
-        let summary = st
-            .wallet()
-            .get_wallet_summary(min_confirmations)
-            .unwrap()
-            .unwrap();
-        let balance = summary.account_balances().get(&account.id()).unwrap();
-        // TODO: in the future, we will distinguish between available and total
-        // balance according to `min_confirmations`
-        assert_eq!(balance.unshielded(), expected);
-
-        // Check the older APIs for consistency.
-        let mempool_height = st.wallet().chain_height().unwrap().unwrap() + 1;
-        assert_eq!(
-            st.wallet()
-                .get_transparent_balances(account.id(), mempool_height)
-                .unwrap()
-                .get(taddr)
-                .cloned()
-                .unwrap_or(NonNegativeAmount::ZERO),
-            expected,
-        );
-        assert_eq!(
-            st.wallet()
-                .get_spendable_transparent_outputs(taddr, mempool_height, 0)
-                .unwrap()
-                .into_iter()
-                .map(|utxo| utxo.value())
-                .sum::<Option<NonNegativeAmount>>(),
-            Some(expected),
-        );
-    };
-
     // The wallet starts out with zero balance.
-    // TODO: Once we have refactored `get_wallet_summary` to distinguish between available
-    // and total balance, we should perform additional checks against available balance;
-    // we use minconf 0 here because all transparent funds are considered shieldable,
-    // irrespective of confirmation depth.
-    check_balance(&st, 0, NonNegativeAmount::ZERO);
+    check_balance::<DSF>(&st, &account, taddr, 0, &Balance::ZERO);
+    // we currently treat min_confirmation the same regardless they are 0 or 1
+    check_balance::<DSF>(&st, &account, taddr, 1, &Balance::ZERO);
 
     // Create a fake transparent output.
     let value = NonNegativeAmount::from_u64(100000).unwrap();
@@ -192,7 +215,14 @@ where
         .unwrap();
 
     // The wallet should detect the balance as available
-    check_balance(&st, 0, value);
+    let mut zero_or_one_conf_value = Balance::ZERO;
+
+    // add the spendable value to the expected balance
+    zero_or_one_conf_value.add_spendable_value(value).unwrap();
+
+    check_balance::<DSF>(&st, &account, taddr, 0, &zero_or_one_conf_value);
+    // we currently treat min_confirmation the same regardless they are 0 or 1
+    check_balance::<DSF>(&st, &account, taddr, 1, &zero_or_one_conf_value);
 
     // Shield the output.
     let input_selector = GreedyInputSelector::new();
@@ -216,14 +246,18 @@ where
 
     // The wallet should have zero transparent balance, because the shielding
     // transaction can be mined.
-    check_balance(&st, 0, NonNegativeAmount::ZERO);
+    check_balance::<DSF>(&st, &account, taddr, 0, &Balance::ZERO);
+    // we currently treat min_confirmation the same regardless they are 0 or 1
+    check_balance::<DSF>(&st, &account, taddr, 1, &Balance::ZERO);
 
     // Mine the shielding transaction.
     let (mined_height, _) = st.generate_next_block_including(txid);
     st.scan_cached_blocks(mined_height, 1);
 
     // The wallet should still have zero transparent balance.
-    check_balance(&st, 0, NonNegativeAmount::ZERO);
+    check_balance::<DSF>(&st, &account, taddr, 0, &Balance::ZERO);
+    // we currently treat min_confirmation the same regardless they are 0 or 1
+    check_balance::<DSF>(&st, &account, taddr, 1, &Balance::ZERO);
 
     // Unmine the shielding transaction via a reorg.
     st.wallet_mut()
@@ -232,7 +266,9 @@ where
     assert_eq!(st.wallet().chain_height().unwrap(), Some(mined_height - 1));
 
     // The wallet should still have zero transparent balance.
-    check_balance(&st, 0, NonNegativeAmount::ZERO);
+    check_balance::<DSF>(&st, &account, taddr, 0, &Balance::ZERO);
+    // we currently treat min_confirmation the same regardless they are 0 or 1
+    check_balance::<DSF>(&st, &account, taddr, 1, &Balance::ZERO);
 
     // Expire the shielding transaction.
     let expiry_height = st
@@ -243,5 +279,90 @@ where
         .expiry_height();
     st.wallet_mut().update_chain_tip(expiry_height).unwrap();
 
-    check_balance(&st, 0, value);
+    check_balance::<DSF>(&st, &account, taddr, 0, &zero_or_one_conf_value);
+    // we currently treat min_confirmation the same regardless they are 0 or 1
+    check_balance::<DSF>(&st, &account, taddr, 1, &zero_or_one_conf_value);
+}
+
+/// This test attempts to verify that transparent funds spendability is
+/// accounted for properly given the different minimum confirmations values
+/// that can be set when querying for balances.
+pub fn transparent_balance_spendability<DSF>(dsf: DSF, cache: impl TestCache)
+where
+    DSF: DataStoreFactory,
+{
+    let mut st = TestBuilder::new()
+        .with_data_store_factory(dsf)
+        .with_block_cache(cache)
+        .with_account_from_sapling_activation(BlockHash([0; 32]))
+        .build();
+
+    let account = st.test_account().cloned().unwrap();
+    let uaddr = st
+        .wallet()
+        .get_current_address(account.id())
+        .unwrap()
+        .unwrap();
+    let taddr = uaddr.transparent().unwrap();
+
+    // Initialize the wallet with chain data that has no shielded notes for us.
+    let not_our_key = ExtendedSpendingKey::master(&[]).to_diversifiable_full_viewing_key();
+    let not_our_value = NonNegativeAmount::const_from_u64(10000);
+    let (start_height, _, _) =
+        st.generate_next_block(&not_our_key, AddressType::DefaultExternal, not_our_value);
+    for _ in 1..10 {
+        st.generate_next_block(&not_our_key, AddressType::DefaultExternal, not_our_value);
+    }
+    st.scan_cached_blocks(start_height, 10);
+
+    // The wallet starts out with zero balance.
+    check_balance::<DSF>(
+        &st as &TestState<_, DSF::DataStore, _>,
+        &account,
+        taddr,
+        0,
+        &Balance::ZERO,
+    );
+    // we currently treat min_confirmation the same regardless they are 0 or 1
+    check_balance::<DSF>(&st, &account, taddr, 1, &Balance::ZERO);
+
+    // Create a fake transparent output.
+    let value = NonNegativeAmount::from_u64(100000).unwrap();
+    let txout = TxOut {
+        value,
+        script_pubkey: taddr.script(),
+    };
+
+    // Pretend the output was received in the chain tip.
+    let height = st.wallet().chain_height().unwrap().unwrap();
+    let utxo = WalletTransparentOutput::from_parts(OutPoint::fake(), txout, Some(height)).unwrap();
+    st.wallet_mut()
+        .put_received_transparent_utxo(&utxo)
+        .unwrap();
+
+    // The wallet should detect the balance as available
+    let mut zero_or_one_conf_value = Balance::ZERO;
+
+    // add the spendable value to the expected balance
+    zero_or_one_conf_value.add_spendable_value(value).unwrap();
+
+    check_balance::<DSF>(&st, &account, taddr, 0, &zero_or_one_conf_value);
+    // we currently treat min_confirmation the same regardless they are 0 or 1
+    check_balance::<DSF>(&st, &account, taddr, 1, &zero_or_one_conf_value);
+
+    // now if we increase the number of confirmations our spendable balance should
+    // be zero and the total balance equal to `value`
+    let mut not_confirmed_yet_value = Balance::ZERO;
+
+    not_confirmed_yet_value
+        .add_pending_spendable_value(value)
+        .unwrap();
+
+    check_balance::<DSF>(&st, &account, taddr, 2, &not_confirmed_yet_value);
+
+    st.generate_empty_block();
+
+    // now we generate one more block and the balance should be the same as when the
+    // check_balance function was called with zero or one confirmation.
+    check_balance::<DSF>(&st, &account, taddr, 2, &zero_or_one_conf_value);
 }