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![];
#[cfg(feature = "orchard")]
let mut orchard_commitments = vec![];
let mut last_scanned_height = None;
let mut note_positions = vec![];
for block in blocks.into_iter() {
@ -541,6 +543,10 @@ impl<P: consensus::Parameters> WalletWrite for WalletDb<rusqlite::Connection, P>
block.block_time(),
block.sapling().final_tree_size(),
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() {
@ -635,6 +641,8 @@ impl<P: consensus::Parameters> WalletWrite for WalletDb<rusqlite::Connection, P>
last_scanned_height = Some(block.height());
let block_commitments = block.into_commitments();
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.
@ -652,31 +660,63 @@ impl<P: consensus::Parameters> WalletWrite for WalletDb<rusqlite::Connection, P>
{
// Create subtrees from the note commitments in parallel.
const CHUNK_SIZE: usize = 1024;
let subtrees = sapling_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;
{
let sapling_subtrees = sapling_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,
SAPLING_SHARD_HEIGHT.into(),
chunk.iter_mut().map(|n| n.take().expect("always Some")),
)
})
.map(|res| (res.subtree, res.checkpoints))
.collect::<Vec<_>>();
shardtree::LocatedTree::from_iter(
start..end,
SAPLING_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 subtrees = subtrees.into_iter();
wdb.with_sapling_tree_mut::<_, _, Self::Error>(move |sapling_tree| {
for (tree, checkpoints) in &mut subtrees {
sapling_tree.insert_tree(tree, checkpoints)?;
}
// Update the Sapling note commitment tree with all newly read note commitments
let mut sapling_subtrees = sapling_subtrees.into_iter();
wdb.with_sapling_tree_mut::<_, _, Self::Error>(move |sapling_tree| {
for (tree, checkpoints) in &mut sapling_subtrees {
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.
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.
#[allow(clippy::too_many_arguments)]
pub(crate) fn put_block(
conn: &rusqlite::Transaction<'_>,
block_height: BlockHeight,
@ -1841,6 +1842,8 @@ pub(crate) fn put_block(
block_time: u32,
sapling_commitment_tree_size: u32,
sapling_output_count: u32,
#[cfg(feature = "orchard")] orchard_commitment_tree_size: u32,
#[cfg(feature = "orchard")] orchard_action_count: u32,
) -> Result<(), SqliteClientError> {
let block_hash_data = conn
.query_row(
@ -1871,7 +1874,9 @@ pub(crate) fn put_block(
time,
sapling_commitment_tree_size,
sapling_output_count,
sapling_tree
sapling_tree,
orchard_commitment_tree_size,
orchard_action_count
)
VALUES (
:height,
@ -1879,21 +1884,32 @@ pub(crate) fn put_block(
:block_time,
:sapling_commitment_tree_size,
:sapling_output_count,
x'00'
x'00',
:orchard_commitment_tree_size,
:orchard_action_count
)
ON CONFLICT (height) DO UPDATE
SET hash = :hash,
time = :block_time,
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![
":height": u32::from(block_height),
":hash": &block_hash.0[..],
":block_time": block_time,
":sapling_commitment_tree_size": sapling_commitment_tree_size,
":sapling_output_count": sapling_output_count,
":orchard_commitment_tree_size": orchard_commitment_tree_size,
":orchard_action_count": orchard_action_count,
])?;
Ok(())

View File

@ -652,6 +652,10 @@ mod tests {
block.block_time(),
block.sapling().final_tree_size(),
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() {