Skip to content

Commit

Permalink
zcash_client_backend: Allow change strategies to act based on wallet …
Browse files Browse the repository at this point in the history
…balance.
  • Loading branch information
nuttycom committed Oct 17, 2024
1 parent e6c1fa3 commit 36a54aa
Show file tree
Hide file tree
Showing 4 changed files with 64 additions and 14 deletions.
29 changes: 29 additions & 0 deletions zcash_client_backend/src/data_api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -804,20 +804,28 @@ impl<NoteRef> SpendableNotes<NoteRef> {
/// the wallet.
pub struct WalletMeta {
sapling_note_count: usize,
sapling_total_value: NonNegativeAmount,
#[cfg(feature = "orchard")]
orchard_note_count: usize,
#[cfg(feature = "orchard")]
orchard_total_value: NonNegativeAmount,
}

impl WalletMeta {
/// Constructs a new [`WalletMeta`] value from its constituent parts.
pub fn new(
sapling_note_count: usize,
sapling_total_value: NonNegativeAmount,
#[cfg(feature = "orchard")] orchard_note_count: usize,
#[cfg(feature = "orchard")] orchard_total_value: NonNegativeAmount,
) -> Self {
Self {
sapling_note_count,
sapling_total_value,
#[cfg(feature = "orchard")]
orchard_note_count,
#[cfg(feature = "orchard")]
orchard_total_value,
}
}

Expand All @@ -838,13 +846,24 @@ impl WalletMeta {
self.sapling_note_count

Check warning on line 846 in zcash_client_backend/src/data_api.rs

View check run for this annotation

Codecov / codecov/patch

zcash_client_backend/src/data_api.rs#L845-L846

Added lines #L845 - L846 were not covered by tests
}

/// Returns the total value of Sapling notes represented by [`Self::sapling_note_count`].
pub fn sapling_total_value(&self) -> NonNegativeAmount {
self.sapling_total_value

Check warning on line 851 in zcash_client_backend/src/data_api.rs

View check run for this annotation

Codecov / codecov/patch

zcash_client_backend/src/data_api.rs#L850-L851

Added lines #L850 - L851 were not covered by tests
}

/// Returns the number of unspent Orchard notes belonging to the account for which this was
/// generated.
#[cfg(feature = "orchard")]
pub fn orchard_note_count(&self) -> usize {
self.orchard_note_count
}

/// Returns the total value of Orchard notes represented by [`Self::orchard_note_count`].
#[cfg(feature = "orchard")]
pub fn orchard_total_value(&self) -> NonNegativeAmount {
self.orchard_total_value
}

/// Returns the total number of unspent shielded notes belonging to the account for which this
/// was generated.
pub fn total_note_count(&self) -> usize {
Expand All @@ -855,6 +874,16 @@ impl WalletMeta {

self.sapling_note_count + orchard_note_count
}

/// Returns the total value of shielded notes represented by [`Self::total_note_count`]
pub fn total_value(&self) -> NonNegativeAmount {

Check warning on line 879 in zcash_client_backend/src/data_api.rs

View check run for this annotation

Codecov / codecov/patch

zcash_client_backend/src/data_api.rs#L879

Added line #L879 was not covered by tests
#[cfg(feature = "orchard")]
let orchard_value = self.orchard_total_value;

Check warning on line 881 in zcash_client_backend/src/data_api.rs

View check run for this annotation

Codecov / codecov/patch

zcash_client_backend/src/data_api.rs#L881

Added line #L881 was not covered by tests
#[cfg(not(feature = "orchard"))]
let orchard_value = NonNegativeAmount::ZERO;

Check warning on line 883 in zcash_client_backend/src/data_api.rs

View check run for this annotation

Codecov / codecov/patch

zcash_client_backend/src/data_api.rs#L883

Added line #L883 was not covered by tests

(self.sapling_total_value + orchard_value).expect("Does not overflow Zcash maximum value.")

Check warning on line 885 in zcash_client_backend/src/data_api.rs

View check run for this annotation

Codecov / codecov/patch

zcash_client_backend/src/data_api.rs#L885

Added line #L885 was not covered by tests
}
}

/// A trait representing the capability to query a data store for unspent transaction outputs
Expand Down
12 changes: 9 additions & 3 deletions zcash_client_backend/src/fees/zip317.rs
Original file line number Diff line number Diff line change
Expand Up @@ -303,7 +303,7 @@ mod tests {

{
// spend a single Sapling note and produce 5 outputs
let balance = |existing_notes| {
let balance = |existing_notes, total| {
change_strategy.compute_balance(
&Network::TestNetwork,
Network::TestNetwork
Expand All @@ -326,14 +326,17 @@ mod tests {
None,
Some(&WalletMeta::new(
existing_notes,
total,
#[cfg(feature = "orchard")]
0,
#[cfg(feature = "orchard")]
NonNegativeAmount::ZERO,
)),
)
};

assert_matches!(
balance(0),
balance(0, NonNegativeAmount::ZERO),
Ok(balance) if
balance.proposed_change() == [
ChangeValue::sapling(NonNegativeAmount::const_from_u64(129_4000), None),
Expand All @@ -346,7 +349,7 @@ mod tests {
);

assert_matches!(
balance(2),
balance(2, NonNegativeAmount::const_from_u64(100_0000)),
Ok(balance) if
balance.proposed_change() == [
ChangeValue::sapling(NonNegativeAmount::const_from_u64(216_0000), None),
Expand Down Expand Up @@ -382,8 +385,11 @@ mod tests {
None,
Some(&WalletMeta::new(
0,
NonNegativeAmount::ZERO,
#[cfg(feature = "orchard")]
0,
#[cfg(feature = "orchard")]
NonNegativeAmount::ZERO,
)),
);

Expand Down
13 changes: 8 additions & 5 deletions zcash_client_sqlite/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -128,7 +128,7 @@ pub mod error;
pub mod wallet;
use wallet::{
commitment_tree::{self, put_shard_roots},
common::count_outputs,
common::spendable_notes_meta,
SubtreeProgressEstimator,
};

Expand Down Expand Up @@ -354,7 +354,7 @@ impl<C: Borrow<rusqlite::Connection>, P: consensus::Parameters> InputSource for
min_value: NonNegativeAmount,
exclude: &[Self::NoteRef],
) -> Result<WalletMeta, Self::Error> {
let sapling_note_count = count_outputs(
let sapling_pool_meta = spendable_notes_meta(
self.conn.borrow(),
account_id,
min_value,
Expand All @@ -363,7 +363,7 @@ impl<C: Borrow<rusqlite::Connection>, P: consensus::Parameters> InputSource for
)?;

#[cfg(feature = "orchard")]
let orchard_note_count = count_outputs(
let orchard_pool_meta = spendable_notes_meta(
self.conn.borrow(),
account_id,
min_value,
Expand All @@ -372,9 +372,12 @@ impl<C: Borrow<rusqlite::Connection>, P: consensus::Parameters> InputSource for
)?;

Ok(WalletMeta::new(
sapling_note_count,
sapling_pool_meta.note_count,
sapling_pool_meta.total_value,
#[cfg(feature = "orchard")]
orchard_note_count,
orchard_pool_meta.note_count,
#[cfg(feature = "orchard")]
orchard_pool_meta.total_value,
))
}
}
Expand Down
24 changes: 18 additions & 6 deletions zcash_client_sqlite/src/wallet/common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -226,13 +226,18 @@ where
.collect::<Result<_, _>>()
}

pub(crate) fn count_outputs(
pub(crate) struct PoolMeta {
pub(crate) note_count: usize,
pub(crate) total_value: NonNegativeAmount,
}

pub(crate) fn spendable_notes_meta(
conn: &rusqlite::Connection,
account: AccountId,
min_value: NonNegativeAmount,
exclude: &[ReceivedNoteId],
protocol: ShieldedProtocol,
) -> Result<usize, rusqlite::Error> {
) -> Result<PoolMeta, SqliteClientError> {
let (table_prefix, _, _) = per_protocol_names(protocol);

let excluded: Vec<Value> = exclude
Expand All @@ -247,9 +252,9 @@ pub(crate) fn count_outputs(
.collect();
let excluded_ptr = Rc::new(excluded);

conn.query_row(
let (note_count, total_value) = conn.query_row(
&format!(
"SELECT COUNT(*)
"SELECT COUNT(*), SUM(value)
FROM {table_prefix}_received_notes
INNER JOIN accounts
ON accounts.id = {table_prefix}_received_notes.account_id
Expand All @@ -274,6 +279,13 @@ pub(crate) fn count_outputs(
":min_value": u64::from(min_value),
":exclude": &excluded_ptr
],
|row| row.get(0),
)
|row| Ok((row.get::<_, usize>(0)?, row.get::<_, Option<i64>>(1)?)),
)?;

Ok(PoolMeta {
note_count,
total_value: total_value.map_or(Ok(NonNegativeAmount::ZERO), |v| {
NonNegativeAmount::from_nonnegative_i64(v).map_err(SqliteClientError::BalanceError)
})?,
})
}

0 comments on commit 36a54aa

Please sign in to comment.