From 8f03208439db1613739bb5bcf84e6791445088c5 Mon Sep 17 00:00:00 2001 From: Jack Grigg Date: Tue, 13 Sep 2022 21:22:46 +0000 Subject: [PATCH 1/2] zcash_client_backend: Add tags to IVKs in the batch scanner This removes the dependency on `SaplingIvk::to_repr()`, and enables us to alter the type of `D::IncomingViewingKey` to improve the performance of batch scanning. For the welding rig, we already annotate the viewing keys with `AccountId`, so we use `(AccountId, Scope)` as the tag. --- zcash_client_backend/src/data_api/chain.rs | 10 ++-- zcash_client_backend/src/scan.rs | 58 ++++++++++++-------- zcash_client_backend/src/welding_rig.rs | 64 ++++++++++++++-------- zcash_primitives/src/sapling/keys.rs | 2 +- 4 files changed, 83 insertions(+), 51 deletions(-) diff --git a/zcash_client_backend/src/data_api/chain.rs b/zcash_client_backend/src/data_api/chain.rs index b5d9567a2..3aa10ce20 100644 --- a/zcash_client_backend/src/data_api/chain.rs +++ b/zcash_client_backend/src/data_api/chain.rs @@ -232,10 +232,12 @@ where let mut batch_runner = BatchRunner::new( 100, - dfvks - .iter() - .flat_map(|(_, dfvk)| [dfvk.to_ivk(Scope::External), dfvk.to_ivk(Scope::Internal)]) - .collect(), + dfvks.iter().flat_map(|(account, dfvk)| { + [ + ((**account, Scope::External), dfvk.to_ivk(Scope::External)), + ((**account, Scope::Internal), dfvk.to_ivk(Scope::Internal)), + ] + }), ); cache.with_blocks(last_height, limit, |block: CompactBlock| { diff --git a/zcash_client_backend/src/scan.rs b/zcash_client_backend/src/scan.rs index ecebb3358..dc56833bb 100644 --- a/zcash_client_backend/src/scan.rs +++ b/zcash_client_backend/src/scan.rs @@ -7,17 +7,18 @@ use zcash_note_encryption::{batch, BatchDomain, Domain, ShieldedOutput, COMPACT_ use zcash_primitives::{block::BlockHash, transaction::TxId}; /// A decrypted note. -pub(crate) struct DecryptedNote { - /// The incoming viewing key used to decrypt the note. - pub(crate) ivk: D::IncomingViewingKey, +pub(crate) struct DecryptedNote { + /// The tag corresponding to the incoming viewing key used to decrypt the note. + pub(crate) ivk_tag: A, /// The recipient of the note. pub(crate) recipient: D::Recipient, /// The note! pub(crate) note: D::Note, } -impl fmt::Debug for DecryptedNote +impl fmt::Debug for DecryptedNote where + A: fmt::Debug, D::IncomingViewingKey: fmt::Debug, D::Recipient: fmt::Debug, D::Note: fmt::Debug, @@ -25,7 +26,7 @@ where { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("DecryptedNote") - .field("ivk", &self.ivk) + .field("ivk_tag", &self.ivk_tag) .field("recipient", &self.recipient) .field("note", &self.note) .finish() @@ -40,10 +41,11 @@ struct OutputIndex { value: V, } -type OutputReplier = OutputIndex>>>>; +type OutputReplier = OutputIndex>>>>; /// A batch of outputs to trial decrypt. -struct Batch> { +struct Batch> { + tags: Vec, ivks: Vec, /// We currently store outputs and repliers as parallel vectors, because /// [`batch::try_note_decryption`] accepts a slice of domain/output pairs @@ -53,18 +55,20 @@ struct Batch> { /// all be part of the same struct, which would also track the output index /// (that is captured in the outer `OutputIndex` of each `OutputReplier`). outputs: Vec<(D, Output)>, - repliers: Vec>, + repliers: Vec>, } -impl Batch +impl Batch where + A: Clone, D: BatchDomain, Output: ShieldedOutput, - D::IncomingViewingKey: Clone, { /// Constructs a new batch. - fn new(ivks: Vec) -> Self { + fn new(tags: Vec, ivks: Vec) -> Self { + assert_eq!(tags.len(), ivks.len()); Self { + tags, ivks, outputs: vec![], repliers: vec![], @@ -86,7 +90,7 @@ where let result = OutputIndex { output_index: replier.output_index, value: decryption_result.map(|((note, recipient), ivk_idx)| DecryptedNote { - ivk: self.ivks[ivk_idx].clone(), + ivk_tag: self.tags[ivk_idx].clone(), recipient, note, }), @@ -100,7 +104,7 @@ where } } -impl + Clone> Batch { +impl + Clone> Batch { /// Adds the given outputs to this batch. /// /// `replier` will be called with the result of every output. @@ -108,7 +112,7 @@ impl + Clone> Batch &mut self, domain: impl Fn() -> D, outputs: &[Output], - replier: channel::Sender>>>, + replier: channel::Sender>>>, ) { self.outputs .extend(outputs.iter().cloned().map(|output| (domain(), output))); @@ -123,30 +127,36 @@ impl + Clone> Batch type ResultKey = (BlockHash, TxId); /// Logic to run batches of trial decryptions on the global threadpool. -pub(crate) struct BatchRunner> { +pub(crate) struct BatchRunner> { batch_size_threshold: usize, - acc: Batch, - pending_results: HashMap>>>>, + acc: Batch, + pending_results: + HashMap>>>>, } -impl BatchRunner +impl BatchRunner where + A: Clone, D: BatchDomain, Output: ShieldedOutput, - D::IncomingViewingKey: Clone, { /// Constructs a new batch runner for the given incoming viewing keys. - pub(crate) fn new(batch_size_threshold: usize, ivks: Vec) -> Self { + pub(crate) fn new( + batch_size_threshold: usize, + ivks: impl Iterator, + ) -> Self { + let (tags, ivks) = ivks.unzip(); Self { batch_size_threshold, - acc: Batch::new(ivks), + acc: Batch::new(tags, ivks), pending_results: HashMap::default(), } } } -impl BatchRunner +impl BatchRunner where + A: Clone + Send + 'static, D: BatchDomain + Send + 'static, D::IncomingViewingKey: Clone + Send, D::Memo: Send, @@ -184,7 +194,7 @@ where /// Subsequent calls to `Self::add_outputs` will be accumulated into a new batch. pub(crate) fn flush(&mut self) { if !self.acc.is_empty() { - let mut batch = Batch::new(self.acc.ivks.clone()); + let mut batch = Batch::new(self.acc.tags.clone(), self.acc.ivks.clone()); mem::swap(&mut batch, &mut self.acc); rayon::spawn_fifo(|| batch.run()); } @@ -199,7 +209,7 @@ where &mut self, block_tag: BlockHash, txid: TxId, - ) -> HashMap<(TxId, usize), DecryptedNote> { + ) -> HashMap<(TxId, usize), DecryptedNote> { self.pending_results .remove(&(block_tag, txid)) // We won't have a pending result if the transaction didn't have outputs of diff --git a/zcash_client_backend/src/welding_rig.rs b/zcash_client_backend/src/welding_rig.rs index 764169040..036a68cb1 100644 --- a/zcash_client_backend/src/welding_rig.rs +++ b/zcash_client_backend/src/welding_rig.rs @@ -38,10 +38,13 @@ use crate::{ /// [`CompactSaplingOutput`]: crate::proto::compact_formats::CompactSaplingOutput /// [`scan_block`]: crate::welding_rig::scan_block pub trait ScanningKey { + /// The type representing the scope of the scanning key. + type Scope: Clone + Eq + std::hash::Hash + Send + 'static; + /// The type of key that is used to decrypt Sapling outputs; type SaplingNk: Clone; - type SaplingKeys: IntoIterator; + type SaplingKeys: IntoIterator; /// The type of nullifier extracted when a note is successfully /// obtained by trial decryption. @@ -63,14 +66,23 @@ pub trait ScanningKey { } impl ScanningKey for DiversifiableFullViewingKey { + type Scope = Scope; type SaplingNk = NullifierDerivingKey; - type SaplingKeys = [(SaplingIvk, Self::SaplingNk); 2]; + type SaplingKeys = [(Self::Scope, SaplingIvk, Self::SaplingNk); 2]; type Nf = sapling::Nullifier; fn to_sapling_keys(&self) -> Self::SaplingKeys { [ - (self.to_ivk(Scope::External), self.to_nk(Scope::External)), - (self.to_ivk(Scope::Internal), self.to_nk(Scope::Internal)), + ( + Scope::External, + self.to_ivk(Scope::External), + self.to_nk(Scope::External), + ), + ( + Scope::Internal, + self.to_ivk(Scope::Internal), + self.to_nk(Scope::Internal), + ), ] } @@ -88,12 +100,13 @@ impl ScanningKey for DiversifiableFullViewingKey { /// /// [`ExtendedFullViewingKey`]: zcash_primitives::zip32::ExtendedFullViewingKey impl ScanningKey for ExtendedFullViewingKey { + type Scope = Scope; type SaplingNk = NullifierDerivingKey; - type SaplingKeys = [(SaplingIvk, Self::SaplingNk); 1]; + type SaplingKeys = [(Self::Scope, SaplingIvk, Self::SaplingNk); 1]; type Nf = sapling::Nullifier; fn to_sapling_keys(&self) -> Self::SaplingKeys { - [(self.fvk.vk.ivk(), self.fvk.vk.nk)] + [(Scope::External, self.fvk.vk.ivk(), self.fvk.vk.nk)] } fn sapling_nf( @@ -110,12 +123,13 @@ impl ScanningKey for ExtendedFullViewingKey { /// /// [`SaplingIvk`]: zcash_primitives::sapling::SaplingIvk impl ScanningKey for SaplingIvk { + type Scope = (); type SaplingNk = (); - type SaplingKeys = [(SaplingIvk, Self::SaplingNk); 1]; + type SaplingKeys = [(Self::Scope, SaplingIvk, Self::SaplingNk); 1]; type Nf = (); fn to_sapling_keys(&self) -> Self::SaplingKeys { - [(self.clone(), ())] + [((), self.clone(), ())] } fn sapling_nf(_key: &Self::SaplingNk, _note: &Note, _witness: &IncrementalWitness) {} @@ -163,11 +177,14 @@ pub fn scan_block( ) } -pub(crate) fn add_block_to_runner( +pub(crate) fn add_block_to_runner( params: &P, block: CompactBlock, - batch_runner: &mut BatchRunner, CompactOutputDescription>, -) { + batch_runner: &mut BatchRunner<(AccountId, S), SaplingDomain

, CompactOutputDescription>, +) where + P: consensus::Parameters + Send + 'static, + S: Clone + Send + 'static, +{ let block_hash = block.hash(); let block_height = block.height(); @@ -198,7 +215,9 @@ pub(crate) fn scan_block_with_runner, existing_witnesses: &mut [&mut IncrementalWitness], - mut batch_runner: Option<&mut BatchRunner, CompactOutputDescription>>, + mut batch_runner: Option< + &mut BatchRunner<(AccountId, K::Scope), SaplingDomain

, CompactOutputDescription>, + >, ) -> Vec> { let mut wtxs: Vec> = vec![]; let block_height = block.height(); @@ -276,7 +295,7 @@ pub(crate) fn scan_block_with_runner>(); @@ -284,11 +303,12 @@ pub(crate) fn scan_block_with_runner>(); @@ -516,6 +536,7 @@ mod tests { #[test] fn scan_block_with_my_tx() { fn go(scan_multithreaded: bool) { + let account = AccountId::from(0); let extsk = ExtendedSpendingKey::master(&[]); let extfvk = ExtendedFullViewingKey::from(&extsk); @@ -535,8 +556,7 @@ mod tests { extfvk .to_sapling_keys() .iter() - .map(|(k, _)| k.clone()) - .collect(), + .map(|(scope, ivk, _)| ((account, *scope), ivk.clone())), ); add_block_to_runner(&Network::TestNetwork, cb.clone(), &mut runner); @@ -550,7 +570,7 @@ mod tests { let txs = scan_block_with_runner( &Network::TestNetwork, cb, - &[(&AccountId::from(0), &extfvk)], + &[(&account, &extfvk)], &[], &mut tree, &mut [], @@ -565,7 +585,7 @@ mod tests { assert_eq!(tx.shielded_spends.len(), 0); assert_eq!(tx.shielded_outputs.len(), 1); assert_eq!(tx.shielded_outputs[0].index, 0); - assert_eq!(tx.shielded_outputs[0].account, AccountId::from(0)); + assert_eq!(tx.shielded_outputs[0].account, account); assert_eq!(tx.shielded_outputs[0].note.value, 5); // Check that the witness root matches @@ -579,6 +599,7 @@ mod tests { #[test] fn scan_block_with_txs_after_my_tx() { fn go(scan_multithreaded: bool) { + let account = AccountId::from(0); let extsk = ExtendedSpendingKey::master(&[]); let extfvk = ExtendedFullViewingKey::from(&extsk); @@ -598,8 +619,7 @@ mod tests { extfvk .to_sapling_keys() .iter() - .map(|(k, _)| k.clone()) - .collect(), + .map(|(scope, ivk, _)| ((account, *scope), ivk.clone())), ); add_block_to_runner(&Network::TestNetwork, cb.clone(), &mut runner); diff --git a/zcash_primitives/src/sapling/keys.rs b/zcash_primitives/src/sapling/keys.rs index 06e24cb6b..1cf5407f0 100644 --- a/zcash_primitives/src/sapling/keys.rs +++ b/zcash_primitives/src/sapling/keys.rs @@ -191,7 +191,7 @@ impl FullViewingKey { /// to other people. For example, a user can give an external [`SaplingIvk`] to a merchant /// terminal, enabling it to only detect "real" transactions from customers and not /// internal transactions from the wallet. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] pub enum Scope { /// A scope used for wallet-external operations, namely deriving addresses to give to /// other users in order to receive funds. From f7b776005137d1812c660cce9eac61ca957008af Mon Sep 17 00:00:00 2001 From: Jack Grigg Date: Tue, 13 Sep 2022 21:58:32 +0000 Subject: [PATCH 2/2] zcash_client_backend: Add some typedefs for complex types --- zcash_client_backend/src/scan.rs | 4 ++-- zcash_client_backend/src/welding_rig.rs | 9 +++++---- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/zcash_client_backend/src/scan.rs b/zcash_client_backend/src/scan.rs index dc56833bb..e628850cf 100644 --- a/zcash_client_backend/src/scan.rs +++ b/zcash_client_backend/src/scan.rs @@ -42,6 +42,7 @@ struct OutputIndex { } type OutputReplier = OutputIndex>>>>; +type OutputResult = channel::Receiver>>>; /// A batch of outputs to trial decrypt. struct Batch> { @@ -130,8 +131,7 @@ type ResultKey = (BlockHash, TxId); pub(crate) struct BatchRunner> { batch_size_threshold: usize, acc: Batch, - pending_results: - HashMap>>>>, + pending_results: HashMap>, } impl BatchRunner diff --git a/zcash_client_backend/src/welding_rig.rs b/zcash_client_backend/src/welding_rig.rs index 036a68cb1..f514cd414 100644 --- a/zcash_client_backend/src/welding_rig.rs +++ b/zcash_client_backend/src/welding_rig.rs @@ -177,10 +177,13 @@ pub fn scan_block( ) } +type TaggedBatchRunner = + BatchRunner<(AccountId, S), SaplingDomain

, CompactOutputDescription>; + pub(crate) fn add_block_to_runner( params: &P, block: CompactBlock, - batch_runner: &mut BatchRunner<(AccountId, S), SaplingDomain

, CompactOutputDescription>, + batch_runner: &mut TaggedBatchRunner, ) where P: consensus::Parameters + Send + 'static, S: Clone + Send + 'static, @@ -215,9 +218,7 @@ pub(crate) fn scan_block_with_runner, existing_witnesses: &mut [&mut IncrementalWitness], - mut batch_runner: Option< - &mut BatchRunner<(AccountId, K::Scope), SaplingDomain

, CompactOutputDescription>, - >, + mut batch_runner: Option<&mut TaggedBatchRunner>, ) -> Vec> { let mut wtxs: Vec> = vec![]; let block_height = block.height();