zcash_client_sqlite: Add Orchard note commitments in `put_blocks` and `put_block`

This commit is contained in:
Kris Nuttycombe 2024-03-07 21:06:34 -07:00 committed by Jack Grigg
parent 1181566401
commit b62763d689
3 changed files with 85 additions and 25 deletions

View File

@ -523,6 +523,8 @@ impl<P: consensus::Parameters> WalletWrite for WalletDb<rusqlite::Connection, P>
) )
}); });
let mut sapling_commitments = vec![]; let mut sapling_commitments = vec![];
#[cfg(feature = "orchard")]
let mut orchard_commitments = vec![];
let mut last_scanned_height = None; let mut last_scanned_height = None;
let mut note_positions = vec![]; let mut note_positions = vec![];
for block in blocks.into_iter() { for block in blocks.into_iter() {
@ -541,6 +543,10 @@ impl<P: consensus::Parameters> WalletWrite for WalletDb<rusqlite::Connection, P>
block.block_time(), block.block_time(),
block.sapling().final_tree_size(), block.sapling().final_tree_size(),
block.sapling().commitments().len().try_into().unwrap(), block.sapling().commitments().len().try_into().unwrap(),
#[cfg(feature = "orchard")]
block.orchard().final_tree_size(),
#[cfg(feature = "orchard")]
block.orchard().commitments().len().try_into().unwrap(),
)?; )?;
for tx in block.transactions() { for tx in block.transactions() {
@ -635,6 +641,8 @@ impl<P: consensus::Parameters> WalletWrite for WalletDb<rusqlite::Connection, P>
last_scanned_height = Some(block.height()); last_scanned_height = Some(block.height());
let block_commitments = block.into_commitments(); let block_commitments = block.into_commitments();
sapling_commitments.extend(block_commitments.sapling.into_iter().map(Some)); sapling_commitments.extend(block_commitments.sapling.into_iter().map(Some));
#[cfg(feature = "orchard")]
orchard_commitments.extend(block_commitments.orchard.into_iter().map(Some));
} }
// Prune the nullifier map of entries we no longer need. // Prune the nullifier map of entries we no longer need.
@ -652,31 +660,63 @@ impl<P: consensus::Parameters> WalletWrite for WalletDb<rusqlite::Connection, P>
{ {
// Create subtrees from the note commitments in parallel. // Create subtrees from the note commitments in parallel.
const CHUNK_SIZE: usize = 1024; const CHUNK_SIZE: usize = 1024;
let subtrees = sapling_commitments {
.par_chunks_mut(CHUNK_SIZE) let sapling_subtrees = sapling_commitments
.enumerate() .par_chunks_mut(CHUNK_SIZE)
.filter_map(|(i, chunk)| { .enumerate()
let start = start_position + (i * CHUNK_SIZE) as u64; .filter_map(|(i, chunk)| {
let end = start + chunk.len() as u64; let start = start_position + (i * CHUNK_SIZE) as u64;
let end = start + chunk.len() as u64;
shardtree::LocatedTree::from_iter( shardtree::LocatedTree::from_iter(
start..end, start..end,
SAPLING_SHARD_HEIGHT.into(), SAPLING_SHARD_HEIGHT.into(),
chunk.iter_mut().map(|n| n.take().expect("always Some")), chunk.iter_mut().map(|n| n.take().expect("always Some")),
) )
}) })
.map(|res| (res.subtree, res.checkpoints)) .map(|res| (res.subtree, res.checkpoints))
.collect::<Vec<_>>(); .collect::<Vec<_>>();
// Update the Sapling note commitment tree with all newly read note commitments // Update the Sapling note commitment tree with all newly read note commitments
let mut subtrees = subtrees.into_iter(); let mut sapling_subtrees = sapling_subtrees.into_iter();
wdb.with_sapling_tree_mut::<_, _, Self::Error>(move |sapling_tree| { wdb.with_sapling_tree_mut::<_, _, Self::Error>(move |sapling_tree| {
for (tree, checkpoints) in &mut subtrees { for (tree, checkpoints) in &mut sapling_subtrees {
sapling_tree.insert_tree(tree, checkpoints)?; sapling_tree.insert_tree(tree, checkpoints)?;
} }
Ok(()) Ok(())
})?; })?;
}
// Create subtrees from the note commitments in parallel.
#[cfg(feature = "orchard")]
{
let orchard_subtrees = orchard_commitments
.par_chunks_mut(CHUNK_SIZE)
.enumerate()
.filter_map(|(i, chunk)| {
let start = start_position + (i * CHUNK_SIZE) as u64;
let end = start + chunk.len() as u64;
shardtree::LocatedTree::from_iter(
start..end,
ORCHARD_SHARD_HEIGHT.into(),
chunk.iter_mut().map(|n| n.take().expect("always Some")),
)
})
.map(|res| (res.subtree, res.checkpoints))
.collect::<Vec<_>>();
// Update the Sapling note commitment tree with all newly read note commitments
let mut orchard_subtrees = orchard_subtrees.into_iter();
wdb.with_orchard_tree_mut::<_, _, Self::Error>(move |orchard_tree| {
for (tree, checkpoints) in &mut orchard_subtrees {
orchard_tree.insert_tree(tree, checkpoints)?;
}
Ok(())
})?;
}
// Update now-expired transactions that didn't get mined. // Update now-expired transactions that didn't get mined.
wallet::update_expired_notes(wdb.conn.0, last_scanned_height)?; wallet::update_expired_notes(wdb.conn.0, last_scanned_height)?;

View File

@ -1834,6 +1834,7 @@ pub(crate) fn get_account_ids(
} }
/// Inserts information about a scanned block into the database. /// Inserts information about a scanned block into the database.
#[allow(clippy::too_many_arguments)]
pub(crate) fn put_block( pub(crate) fn put_block(
conn: &rusqlite::Transaction<'_>, conn: &rusqlite::Transaction<'_>,
block_height: BlockHeight, block_height: BlockHeight,
@ -1841,6 +1842,8 @@ pub(crate) fn put_block(
block_time: u32, block_time: u32,
sapling_commitment_tree_size: u32, sapling_commitment_tree_size: u32,
sapling_output_count: u32, sapling_output_count: u32,
#[cfg(feature = "orchard")] orchard_commitment_tree_size: u32,
#[cfg(feature = "orchard")] orchard_action_count: u32,
) -> Result<(), SqliteClientError> { ) -> Result<(), SqliteClientError> {
let block_hash_data = conn let block_hash_data = conn
.query_row( .query_row(
@ -1871,7 +1874,9 @@ pub(crate) fn put_block(
time, time,
sapling_commitment_tree_size, sapling_commitment_tree_size,
sapling_output_count, sapling_output_count,
sapling_tree sapling_tree,
orchard_commitment_tree_size,
orchard_action_count
) )
VALUES ( VALUES (
:height, :height,
@ -1879,21 +1884,32 @@ pub(crate) fn put_block(
:block_time, :block_time,
:sapling_commitment_tree_size, :sapling_commitment_tree_size,
:sapling_output_count, :sapling_output_count,
x'00' x'00',
:orchard_commitment_tree_size,
:orchard_action_count
) )
ON CONFLICT (height) DO UPDATE ON CONFLICT (height) DO UPDATE
SET hash = :hash, SET hash = :hash,
time = :block_time, time = :block_time,
sapling_commitment_tree_size = :sapling_commitment_tree_size, sapling_commitment_tree_size = :sapling_commitment_tree_size,
sapling_output_count = :sapling_output_count", sapling_output_count = :sapling_output_count,
orchard_commitment_tree_size = :orchard_commitment_tree_size,
orchard_action_count = :orchard_action_count",
)?; )?;
#[cfg(not(feature = "orchard"))]
let orchard_commitment_tree_size: Option<u32> = None;
#[cfg(not(feature = "orchard"))]
let orchard_action_count: Option<u32> = None;
stmt_upsert_block.execute(named_params![ stmt_upsert_block.execute(named_params![
":height": u32::from(block_height), ":height": u32::from(block_height),
":hash": &block_hash.0[..], ":hash": &block_hash.0[..],
":block_time": block_time, ":block_time": block_time,
":sapling_commitment_tree_size": sapling_commitment_tree_size, ":sapling_commitment_tree_size": sapling_commitment_tree_size,
":sapling_output_count": sapling_output_count, ":sapling_output_count": sapling_output_count,
":orchard_commitment_tree_size": orchard_commitment_tree_size,
":orchard_action_count": orchard_action_count,
])?; ])?;
Ok(()) Ok(())

View File

@ -652,6 +652,10 @@ mod tests {
block.block_time(), block.block_time(),
block.sapling().final_tree_size(), block.sapling().final_tree_size(),
block.sapling().commitments().len().try_into().unwrap(), block.sapling().commitments().len().try_into().unwrap(),
#[cfg(feature = "orchard")]
block.orchard().final_tree_size(),
#[cfg(feature = "orchard")]
block.orchard().commitments().len().try_into().unwrap(),
)?; )?;
for tx in block.transactions() { for tx in block.transactions() {