Skip to content

Commit

Permalink
zcash_client_sqlite: Add Orchard support to get_wallet_summary
Browse files Browse the repository at this point in the history
  • Loading branch information
@str4d
str4d committed Mar 10, 2024
1 parent a7ec6ad commit 66d7362
Showing 1 changed file with 218 additions and 59 deletions.
277 changes: 218 additions & 59 deletions zcash_client_sqlite/src/wallet.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -723,6 +723,14 @@ pub(crate) trait ScanProgress {
fully_scanned_height: BlockHeight,
chain_tip_height: BlockHeight,
) -> Result<Option<Ratio<u64>>, SqliteClientError>;

fn orchard_scan_progress(
&self,
conn: &rusqlite::Connection,
birthday_height: BlockHeight,
fully_scanned_height: BlockHeight,
chain_tip_height: BlockHeight,
) -> Result<Option<Ratio<u64>>, SqliteClientError>;
}

#[derive(Debug)]
Expand Down Expand Up @@ -804,6 +812,82 @@ impl ScanProgress for SubtreeScanProgress {
.flatten())
}
}

#[tracing::instrument(skip(conn))]
fn orchard_scan_progress(
&self,
conn: &rusqlite::Connection,
birthday_height: BlockHeight,
fully_scanned_height: BlockHeight,
chain_tip_height: BlockHeight,
) -> Result<Option<Ratio<u64>>, SqliteClientError> {
if fully_scanned_height == chain_tip_height {
// Compute the total blocks scanned since the wallet birthday
conn.query_row(
"SELECT SUM(orchard_action_count)
FROM blocks
WHERE height >= :birthday_height",
named_params![":birthday_height": u32::from(birthday_height)],
|row| {
let scanned = row.get::<_, Option<u64>>(0)?;
Ok(scanned.map(|n| Ratio::new(n, n)))
},
)
.map_err(SqliteClientError::from)
} else {
let start_height = birthday_height;
// Compute the starting number of notes directly from the blocks table
let start_size = conn.query_row(
"SELECT MAX(orchard_commitment_tree_size)
FROM blocks
WHERE height <= :start_height",
named_params![":start_height": u32::from(start_height)],
|row| row.get::<_, Option<u64>>(0),
)?;

// Compute the total blocks scanned so far above the starting height
let scanned_count = conn.query_row(
"SELECT SUM(orchard_action_count)
FROM blocks
WHERE height > :start_height",
named_params![":start_height": u32::from(start_height)],
|row| row.get::<_, Option<u64>>(0),
)?;

// We don't have complete information on how many actions will exist in the shard at
// the chain tip without having scanned the chain tip block, so we overestimate by
// computing the maximum possible number of notes directly from the shard indices.
//
// TODO: it would be nice to be able to reliably have the size of the commitment tree
// at the chain tip without having to have scanned that block.
Ok(conn
.query_row(
"SELECT MIN(shard_index), MAX(shard_index)
FROM orchard_tree_shards
WHERE subtree_end_height > :start_height
OR subtree_end_height IS NULL",
named_params![":start_height": u32::from(start_height)],
|row| {
let min_tree_size = row
.get::<_, Option<u64>>(0)?
.map(|min| min << ORCHARD_SHARD_HEIGHT);
let max_idx = row.get::<_, Option<u64>>(1)?;
Ok(start_size
.or(min_tree_size)
.zip(max_idx)
.map(|(min_tree_size, max)| {
let max_tree_size = (max + 1) << ORCHARD_SHARD_HEIGHT;
Ratio::new(
scanned_count.unwrap_or(0),
max_tree_size - min_tree_size,
)
}))
},
)
.optional()?
.flatten())
}
}
}

/// Returns the spendable balance for the account at the specified height.
Expand Down Expand Up @@ -841,6 +925,25 @@ pub(crate) fn get_wallet_summary<P: consensus::Parameters>(
chain_tip_height,
)?;

let orchard_scan_progress = progress.orchard_scan_progress(
tx,
birthday_height,
fully_scanned_height,
chain_tip_height,
)?;

// Treat Sapling and Orchard outputs as having the same cost to scan.
let scan_progress = sapling_scan_progress
.zip(orchard_scan_progress)
.map(|(s, o)| {
Ratio::new(
s.numerator() + o.numerator(),
s.denominator() + o.denominator(),
)
})
.or(sapling_scan_progress)
.or(orchard_scan_progress);

// If the shard containing the summary height contains any unscanned ranges that start below or
// including that height, none of our balance is currently spendable.
#[tracing::instrument(skip_all)]
Expand Down Expand Up @@ -871,78 +974,134 @@ pub(crate) fn get_wallet_summary<P: consensus::Parameters>(
})
.collect::<Result<HashMap<AccountId, AccountBalance>, _>>()?;

let sapling_trace = tracing::info_span!("stmt_select_notes").entered();
let mut stmt_select_notes = tx.prepare_cached(
"SELECT n.account_id, n.value, n.is_change, scan_state.max_priority, t.block
FROM sapling_received_notes n
JOIN transactions t ON t.id_tx = n.tx
LEFT OUTER 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
AND (
t.expiry_height IS NULL
OR t.block IS NOT NULL
OR t.expiry_height >= :summary_height
)",
)?;
fn count_notes<F>(
tx: &rusqlite::Transaction,
summary_height: BlockHeight,
any_spendable: bool,
account_balances: &mut HashMap<AccountId, AccountBalance>,
table_prefix: &'static str,
with_pool_balance: F,
) -> Result<(), SqliteClientError>
where
F: Fn(
&mut AccountBalance,
NonNegativeAmount,
NonNegativeAmount,
NonNegativeAmount,
) -> Result<(), SqliteClientError>,
{
let mut stmt_select_notes = tx.prepare_cached(&format!(
"SELECT n.account_id, n.value, n.is_change, scan_state.max_priority, t.block
FROM {table_prefix}_received_notes n
JOIN transactions t ON t.id_tx = n.tx
LEFT OUTER JOIN v_{table_prefix}_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
AND (
t.expiry_height IS NULL
OR t.block IS NOT NULL
OR t.expiry_height >= :summary_height
)",
))?;

let mut rows =
stmt_select_notes.query(named_params![":summary_height": u32::from(summary_height)])?;
while let Some(row) = rows.next()? {
let account = AccountId(row.get::<_, u32>(0)?);
let mut rows =
stmt_select_notes.query(named_params![":summary_height": u32::from(summary_height)])?;
while let Some(row) = rows.next()? {
let account = AccountId(row.get::<_, u32>(0)?);

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 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)?;

// If `max_priority` is null, this means that the note is not positioned; the note
// will not be spendable, so we assign the scan priority to `ChainTip` as a priority
// that is greater than `Scanned`
let max_priority_raw = row.get::<_, Option<i64>>(3)?;
let max_priority = max_priority_raw.map_or_else(
|| Ok(ScanPriority::ChainTip),
|raw| {
parse_priority_code(raw).ok_or_else(|| {
SqliteClientError::CorruptedData(format!(
"Priority code {} not recognized.",
raw
))
})
},
)?;
let is_change = row.get::<_, bool>(2)?;

// If `max_priority` is null, this means that the note is not positioned; the note
// will not be spendable, so we assign the scan priority to `ChainTip` as a priority
// that is greater than `Scanned`
let max_priority_raw = row.get::<_, Option<i64>>(3)?;
let max_priority = max_priority_raw.map_or_else(
|| Ok(ScanPriority::ChainTip),
|raw| {
parse_priority_code(raw).ok_or_else(|| {
SqliteClientError::CorruptedData(format!(
"Priority code {} not recognized.",
raw
))
})
},
)?;

let received_height = row.get::<_, Option<u32>>(4)?.map(BlockHeight::from);
let received_height = row.get::<_, Option<u32>>(4)?.map(BlockHeight::from);

let is_spendable = any_spendable
&& received_height.iter().any(|h| h <= &summary_height)
&& max_priority <= ScanPriority::Scanned;
let is_spendable = any_spendable
&& received_height.iter().any(|h| h <= &summary_height)
&& max_priority <= ScanPriority::Scanned;

let is_pending_change = is_change && received_height.iter().all(|h| h > &summary_height);
let is_pending_change =
is_change && received_height.iter().all(|h| h > &summary_height);

let (spendable_value, change_pending_confirmation, value_pending_spendability) = {
let zero = NonNegativeAmount::ZERO;
if is_spendable {
(value, zero, zero)
} else if is_pending_change {
(zero, value, zero)
} else {
(zero, zero, value)
let (spendable_value, change_pending_confirmation, value_pending_spendability) = {
let zero = NonNegativeAmount::ZERO;
if is_spendable {
(value, zero, zero)
} else if is_pending_change {
(zero, value, zero)
} else {
(zero, zero, value)
}
};

if let Some(balances) = account_balances.get_mut(&account) {
with_pool_balance(
balances,
spendable_value,
change_pending_confirmation,
value_pending_spendability,
)?;
}
};
}
Ok(())
}

if let Some(balances) = account_balances.get_mut(&account) {
let orchard_trace = tracing::info_span!("orchard_balances").entered();
count_notes(
tx,
summary_height,
any_spendable,
&mut account_balances,
ORCHARD_TABLES_PREFIX,
|balances, spendable_value, change_pending_confirmation, value_pending_spendability| {
balances.with_orchard_balance_mut::<_, SqliteClientError>(|bal| {
bal.add_spendable_value(spendable_value)?;
bal.add_pending_change_value(change_pending_confirmation)?;
bal.add_pending_spendable_value(value_pending_spendability)?;
Ok(())
})
},
)?;
drop(orchard_trace);

let sapling_trace = tracing::info_span!("sapling_balances").entered();
count_notes(
tx,
summary_height,
any_spendable,
&mut account_balances,
SAPLING_TABLES_PREFIX,
|balances, spendable_value, change_pending_confirmation, value_pending_spendability| {
balances.with_sapling_balance_mut::<_, SqliteClientError>(|bal| {
bal.add_spendable_value(spendable_value)?;
bal.add_pending_change_value(change_pending_confirmation)?;
bal.add_pending_spendable_value(value_pending_spendability)?;
Ok(())
})?;
}
}
})
},
)?;
drop(sapling_trace);

#[cfg(feature = "transparent-inputs")]
Expand Down Expand Up @@ -1025,7 +1184,7 @@ pub(crate) fn get_wallet_summary<P: consensus::Parameters>(
account_balances,
chain_tip_height,
fully_scanned_height,
sapling_scan_progress,
scan_progress,
next_sapling_subtree_index,
#[cfg(feature = "orchard")]
next_orchard_subtree_index,
Expand Down

0 comments on commit 66d7362

Please sign in to comment.