Factor out witness updates from scan_output

This commit is contained in:
Kris Nuttycombe 2022-07-25 19:29:56 -06:00
parent 5873950648
commit 878a428c3c
1 changed files with 21 additions and 27 deletions

View File

@ -18,7 +18,7 @@ use zcash_primitives::{
zip32::{AccountId, ExtendedFullViewingKey},
};
use crate::proto::compact_formats::{CompactBlock, CompactSaplingOutput};
use crate::proto::compact_formats::CompactBlock;
use crate::wallet::{WalletShieldedOutput, WalletShieldedSpend, WalletTx};
/// Scans a [`CompactSaplingOutput`] with a set of [`ScanningKey`]s.
@ -35,29 +35,11 @@ fn scan_output<P: consensus::Parameters, K: ScanningKey>(
params: &P,
height: BlockHeight,
index: usize,
output: CompactSaplingOutput,
output: CompactOutputDescription,
vks: &[(&AccountId, &K)],
spent_from_accounts: &HashSet<AccountId>,
tree: &mut CommitmentTree<Node>,
existing_witnesses: &mut [&mut IncrementalWitness<Node>],
block_witnesses: &mut [&mut IncrementalWitness<Node>],
new_witnesses: &mut [&mut IncrementalWitness<Node>],
) -> Option<WalletShieldedOutput<K::Nf>> {
let output = CompactOutputDescription::try_from(output).ok()?;
// Increment tree and witnesses
let node = Node::new(output.cmu.to_repr());
for witness in existing_witnesses {
witness.append(node).unwrap();
}
for witness in block_witnesses {
witness.append(node).unwrap();
}
for witness in new_witnesses {
witness.append(node).unwrap();
}
tree.append(node).unwrap();
for (account, vk) in vks.iter() {
let (note, to) = match vk.try_decryption(params, height, &output) {
Some(ret) => ret,
@ -283,26 +265,38 @@ pub fn scan_block<P: consensus::Parameters, K: ScanningKey>(
})
.collect();
for (idx, c_out) in tx.outputs.into_iter().enumerate() {
for (index, output) in tx.outputs.into_iter().enumerate() {
// Grab mutable references to new witnesses from previous outputs
// in this transaction so that we can update them. Scoped so we
// don't hold mutable references to shielded_outputs for too long.
let mut new_witnesses: Vec<_> = shielded_outputs
let new_witnesses: Vec<_> = shielded_outputs
.iter_mut()
.map(|output| &mut output.witness)
.collect();
let output = CompactOutputDescription::try_from(output).ok().unwrap();
// Increment tree and witnesses
let node = Node::new(output.cmu.to_repr());
for witness in &mut *existing_witnesses {
witness.append(node).unwrap();
}
for witness in &mut block_witnesses {
witness.append(node).unwrap();
}
for witness in new_witnesses {
witness.append(node).unwrap();
}
tree.append(node).unwrap();
if let Some(output) = scan_output(
params,
block_height,
idx,
c_out,
index,
output,
vks,
&spent_from_accounts,
tree,
existing_witnesses,
&mut block_witnesses,
&mut new_witnesses,
) {
shielded_outputs.push(output);
}