WIP: Add get_wallet_summary to WalletRead

The intent of this API is to provide a single API which returns in a
single call:

* per-account balances, including pending values
* wallet sync progress

Fixes #865
Fixes #900

zcash_client_backend/src/data_api.rs

@@ -1,6 +1,6 @@
 //! Interfaces for wallet data persistence & low-level wallet utilities.
 
-use std::collections::HashMap;
+use std::collections::{BTreeMap, HashMap};
 use std::fmt::Debug;
 use std::num::NonZeroU32;
 
@@ -14,7 +14,10 @@ use zcash_primitives::{
     memo::{Memo, MemoBytes},
     sapling::{self, Node, NOTE_COMMITMENT_TREE_DEPTH},
     transaction::{
-        components::{amount::Amount, OutPoint},
+        components::{
+            amount::{Amount, NonNegativeAmount},
+            OutPoint,
+        },
         Transaction, TxId,
     },
     zip32::{AccountId, ExtendedFullViewingKey},
@@ -44,6 +47,95 @@ pub enum NullifierQuery {
     All,
 }
 
+#[derive(Debug, Clone)]
+pub struct AccountBalance {
+    pub spendable_value: NonNegativeAmount,
+    pub pending_change_value: NonNegativeAmount,
+    pub value_pending_spendability: NonNegativeAmount,
+    pub unshielded: NonNegativeAmount,
+}
+
+impl AccountBalance {
+    pub fn total(&self) -> NonNegativeAmount {
+        (self.spendable_value
+            + self.pending_change_value
+            + self.value_pending_spendability
+            + self.unshielded)
+            .expect("Account balance cannot overflow MAX_MONEY")
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct Ratio<T> {
+    numerator: T,
+    denominator: T,
+}
+
+impl<T> Ratio<T> {
+    pub fn new(numerator: T, denominator: T) -> Self {
+        Self {
+            numerator,
+            denominator,
+        }
+    }
+
+    pub fn numerator(&self) -> &T {
+        &self.numerator
+    }
+
+    pub fn denominator(&self) -> &T {
+        &self.denominator
+    }
+}
+
+/// A type representing the potentially-spendable value of unspent outputs in the wallet.
+///
+/// The balances reported using this data structure may overestimate the total spendable value of
+/// the wallet, in the case that the spend of a previously received shielded note has not yet been
+/// detected by the process of scanning the chain. The balances reported using this data structure
+/// can only be certain to be unspent in the case that [`Self::chain_tip_height`] is equal to
+/// [`Self::fully_scanned_height]`, and even in this circumstance it is possible that a newly
+/// created transaction could conflict with a not-yet-mined transaction in the mempool.
+#[derive(Debug, Clone)]
+pub struct WalletSummary {
+    account_balances: BTreeMap<AccountId, AccountBalance>,
+    chain_tip_height: BlockHeight,
+    fully_scanned_height: BlockHeight,
+    sapling_scan_progress: Ratio<u64>,
+}
+
+impl WalletSummary {
+    pub fn new(
+        account_balances: BTreeMap<AccountId, AccountBalance>,
+        chain_tip_height: BlockHeight,
+        fully_scanned_height: BlockHeight,
+        sapling_scan_progress: Ratio<u64>,
+    ) -> Self {
+        Self {
+            account_balances,
+            chain_tip_height,
+            fully_scanned_height,
+            sapling_scan_progress,
+        }
+    }
+
+    pub fn account_balances(&self) -> &BTreeMap<AccountId, AccountBalance> {
+        &self.account_balances
+    }
+
+    pub fn chain_tip_height(&self) -> BlockHeight {
+        self.chain_tip_height
+    }
+
+    pub fn fully_scanned_height(&self) -> BlockHeight {
+        self.fully_scanned_height
+    }
+
+    pub fn sapling_scan_progress(&self) -> Ratio<u64> {
+        self.sapling_scan_progress
+    }
+}
+
 /// Read-only operations required for light wallet functions.
 ///
 /// This trait defines the read-only portion of the storage interface atop which
@@ -155,15 +247,12 @@ pub trait WalletRead {
         extfvk: &ExtendedFullViewingKey,
     ) -> Result<bool, Self::Error>;
 
-    /// Returns the wallet balance for an account as of the specified block height.
-    ///
-    /// This may be used to obtain a balance that ignores notes that have been received so recently
-    /// that they are not yet deemed spendable.
-    fn get_balance_at(
+    /// Returns the wallet balances and sync status for an account as of the specified block
+    /// height, or `Ok(None)` if the wallet has no balance data available.
+    fn get_wallet_summary(
         &self,
-        account: AccountId,
-        anchor_height: BlockHeight,
-    ) -> Result<Amount, Self::Error>;
+        anchor_height: Option<BlockHeight>,
+    ) -> Result<Option<WalletSummary>, Self::Error>;
 
     /// Returns the memo for a note.
     ///
@@ -694,7 +783,7 @@ pub mod testing {
     use super::{
         chain::CommitmentTreeRoot, scanning::ScanRange, AccountBirthday, BlockMetadata,
         DecryptedTransaction, NoteId, NullifierQuery, ScannedBlock, SentTransaction,
-        WalletCommitmentTrees, WalletRead, WalletWrite, SAPLING_SHARD_HEIGHT,
+        WalletCommitmentTrees, WalletRead, WalletSummary, WalletWrite, SAPLING_SHARD_HEIGHT,
     };
 
     pub struct MockWalletDb {
@@ -800,12 +889,11 @@ pub mod testing {
             Ok(false)
         }
 
-        fn get_balance_at(
+        fn get_wallet_summary(
             &self,
-            _account: AccountId,
-            _anchor_height: BlockHeight,
-        ) -> Result<Amount, Self::Error> {
-            Ok(Amount::zero())
+            _anchor_height: Option<BlockHeight>,
+        ) -> Result<Option<WalletSummary>, Self::Error> {
+            Ok(None)
         }
 
         fn get_memo(&self, _id_note: NoteId) -> Result<Option<Memo>, Self::Error> {

zcash_client_sqlite/src/lib.rs

@@ -66,7 +66,7 @@ use zcash_client_backend::{
         scanning::{ScanPriority, ScanRange},
         AccountBirthday, BlockMetadata, DecryptedTransaction, NoteId, NullifierQuery, PoolType,
         Recipient, ScannedBlock, SentTransaction, ShieldedProtocol, WalletCommitmentTrees,
-        WalletRead, WalletWrite, SAPLING_SHARD_HEIGHT,
+        WalletRead, WalletSummary, WalletWrite, SAPLING_SHARD_HEIGHT,
     },
     keys::{UnifiedFullViewingKey, UnifiedSpendingKey},
     proto::compact_formats::CompactBlock,
@@ -243,12 +243,11 @@ impl<C: Borrow<rusqlite::Connection>, P: consensus::Parameters> WalletRead for W
         wallet::is_valid_account_extfvk(self.conn.borrow(), &self.params, account, extfvk)
     }
 
-    fn get_balance_at(
+    fn get_wallet_summary(
         &self,
-        account: AccountId,
-        anchor_height: BlockHeight,
-    ) -> Result<Amount, Self::Error> {
-        wallet::get_balance_at(self.conn.borrow(), account, anchor_height)
+        anchor_height: Option<BlockHeight>,
+    ) -> Result<Option<WalletSummary>, Self::Error> {
+        wallet::get_wallet_summary(self.conn.borrow(), anchor_height)
     }
 
     fn get_memo(&self, note_id: NoteId) -> Result<Option<Memo>, Self::Error> {

zcash_client_sqlite/src/wallet.rs

@@ -68,11 +68,13 @@ use incrementalmerkletree::Retention;
 use rusqlite::{self, named_params, OptionalExtension, ToSql};
 use shardtree::ShardTree;
 use std::cmp;
-use std::collections::HashMap;
+use std::collections::{BTreeMap, HashMap};
 use std::convert::TryFrom;
 use std::io::{self, Cursor};
 use std::num::NonZeroU32;
 use tracing::debug;
+use zcash_client_backend::data_api::{AccountBalance, Ratio, WalletSummary};
+use zcash_primitives::transaction::components::amount::NonNegativeAmount;
 
 use zcash_client_backend::data_api::{
     scanning::{ScanPriority, ScanRange},
@@ -106,7 +108,7 @@ use crate::{
 };
 use crate::{SAPLING_TABLES_PREFIX, VERIFY_LOOKAHEAD};
 
-use self::scanning::replace_queue_entries;
+use self::scanning::{parse_priority_code, replace_queue_entries};
 
 #[cfg(feature = "transparent-inputs")]
 use {
@@ -519,28 +521,152 @@ pub(crate) fn get_balance(
     }
 }
 
-/// Returns the verified balance for the account at the specified height,
-/// This may be used to obtain a balance that ignores notes that have been
-/// received so recently that they are not yet deemed spendable.
-pub(crate) fn get_balance_at(
+/// Returns the spendable balance for the account at the specified height.
+///
+/// This may be used to obtain a balance that ignores notes that have been detected so recently
+/// that they are not yet spendable, or for which it is not yet possible to construct witnesses.
+pub(crate) fn get_wallet_summary(
     conn: &rusqlite::Connection,
-    account: AccountId,
-    anchor_height: BlockHeight,
-) -> Result<Amount, SqliteClientError> {
-    let balance = conn.query_row(
-        "SELECT SUM(value) FROM sapling_received_notes
-        INNER JOIN transactions ON transactions.id_tx = sapling_received_notes.tx
-        WHERE account = ? AND spent IS NULL AND transactions.block <= ?",
-        [u32::from(account), u32::from(anchor_height)],
-        |row| row.get(0).or(Ok(0)),
+    anchor_height: Option<BlockHeight>,
+) -> Result<Option<WalletSummary>, SqliteClientError> {
+    let chain_tip_height = match scan_queue_extrema(conn)? {
+        Some((_, max)) => max,
+        None => {
+            println!("Scan queue is empty");
+            return Ok(None);
+        }
+    };
+
+    let birthday =
+        wallet_birthday(conn)?.expect("If a scan range exists, we know the wallet birthday.");
+
+    let fully_scanned_height = block_fully_scanned(conn)?.map_or(birthday, |m| m.block_height());
+    let summary_height = anchor_height.unwrap_or(chain_tip_height);
+
+    let sapling_scan_progress = conn.query_row(
+        "WITH
+            fully_scanned AS (
+                SELECT SUM(sapling_output_count) AS notes
+                FROM blocks
+                WHERE height <= :fully_scanned_height
+                AND height <= :summary_height
+            ),
+            scan_incomplete AS (
+                SELECT SUM(sapling_output_count) AS notes
+                FROM blocks
+                WHERE height > :fully_scanned_height
+                AND height <= :summary_height
+            )
+        SELECT 
+            IFNULL(fully_scanned.notes, 0), 
+            IFNULL(scan_incomplete.notes, 0)
+        FROM fully_scanned LEFT OUTER JOIN scan_incomplete",
+        named_params![
+            ":fully_scanned_height": u32::from(fully_scanned_height),
+            ":summary_height": u32::from(summary_height)
+        ],
+        |row| {
+            let fully_scanned: u64 = row.get(0)?;
+            let scan_incomplete: u64 = row.get(0)?;
+            Ok(Ratio::new(scan_incomplete, scan_incomplete + fully_scanned))
+        },
     )?;
 
-    match Amount::from_i64(balance) {
-        Ok(amount) if !amount.is_negative() => Ok(amount),
-        _ => Err(SqliteClientError::CorruptedData(
-            "Sum of values in sapling_received_notes is out of range".to_string(),
-        )),
+    // If any of the shard below the anchor (or chain tip if no anchor is specified) is unscanned
+    // above the wallet birthday height, none of our balance is currently spendable.
+    let any_spendable = conn
+        .query_row(
+            "SELECT 1 FROM v_sapling_shard_unscanned_ranges
+             WHERE :summary_height
+                BETWEEN subtree_start_height
+                AND IFNULL(subtree_end_height, :summary_height)
+             AND block_range_start <= :summary_height",
+            named_params![":summary_height": u32::from(summary_height)],
+            |_| Ok(true),
+        )
+        .optional()?
+        .is_none();
+
+    let mut stmt_select_notes = conn.prepare_cached(
+        "SELECT n.account, n.value, n.is_change, scan_state.max_priority, t.block, t.expiry_height
+        FROM sapling_received_notes n
+        LEFT OUTER JOIN transactions t ON t.id_tx = n.tx
+        INNER JOIN v_sapling_shards_scan_state scan_state
+            ON  n.commitment_tree_position >= scan_state.start_position
+            AND n.commitment_tree_position <  scan_state.end_position_exclusive
+        WHERE n.spent IS NULL",
+    )?;
+
+    let mut account_balances: BTreeMap<AccountId, AccountBalance> = BTreeMap::new();
+    let mut rows = stmt_select_notes.query([])?;
+    while let Some(row) = rows.next()? {
+        let account = row.get::<_, u32>(0).map(AccountId::from)?;
+
+        let value_raw = row.get::<_, i64>(1)?;
+        let value = NonNegativeAmount::from_nonnegative_i64(value_raw).map_err(|_| {
+            SqliteClientError::CorruptedData(format!("Negative received note value: {}", value_raw))
+        })?;
+
+        let is_change = row.get::<_, bool>(2)?;
+
+        let max_priority_raw = row.get::<_, i64>(3)?;
+        let max_priority = parse_priority_code(max_priority_raw).ok_or_else(|| {
+            SqliteClientError::CorruptedData(format!(
+                "Priority code {} not recognized.",
+                max_priority_raw
+            ))
+        })?;
+
+        let received_height = row
+            .get::<_, Option<u32>>(4)
+            .map(|opt| opt.map(BlockHeight::from))?;
+
+        let mut spendable_value = NonNegativeAmount::zero();
+        let mut pending_change_value = NonNegativeAmount::zero();
+        let mut value_pending_spendability = NonNegativeAmount::zero();
+
+        let is_spendable = any_spendable
+            && received_height.iter().any(|h| h <= &summary_height)
+            && max_priority <= ScanPriority::Scanned;
+
+        if is_spendable {
+            spendable_value = value;
+        } else if is_change {
+            pending_change_value = value;
+        } else {
+            value_pending_spendability = value;
+        }
+
+        account_balances
+            .entry(account)
+            .and_modify(|bal| {
+                bal.spendable_value = (bal.spendable_value + spendable_value)
+                    .expect("Spendable value cannot overflow");
+                bal.pending_change_value = (bal.pending_change_value + pending_change_value)
+                    .expect("Pending change value cannot overflow");
+                bal.value_pending_spendability = (bal.value_pending_spendability
+                    + value_pending_spendability)
+                    .expect("Value pending spendability cannot overflow");
+            })
+            .or_insert(AccountBalance {
+                spendable_value,
+                pending_change_value,
+                value_pending_spendability,
+                unshielded: NonNegativeAmount::zero(),
+            });
     }
+
+    // TODO: Add unshielded balances
+
+    let summary = WalletSummary::new(
+        account_balances,
+        chain_tip_height,
+        fully_scanned_height,
+        sapling_scan_progress,
+    );
+    println!("Summary: {:?}", summary);
+
+    Ok(Some(summary))
 }
 
 /// Returns the memo for a received note, if the note is known to the wallet.