zcash_client_backend: Add Orchard support to batch scanning.

Cargo.toml

@@ -115,3 +115,6 @@ zip32 = "0.1"
 lto = true
 panic = 'abort'
 codegen-units = 1
+
+[patch.crates-io]
+orchard = { git = "https://github.com/nuttycom/orchard", rev = "9729cd8d266a6121bd8c7f3b43053440b787d413" }

zcash_client_backend/src/data_api.rs

@@ -693,7 +693,7 @@ pub struct ScannedBlockCommitments {
     /// The ordered vector of note commitments for Orchard outputs of the block.
     /// Present only when the `orchard` feature is enabled.
     #[cfg(feature = "orchard")]
-    pub orchard: Vec<(orchard::note::NoteCommitment, Retention<BlockHeight>)>,
+    pub orchard: Vec<(orchard::tree::MerkleHashOrchard, Retention<BlockHeight>)>,
 }
 
 /// The subset of information that is relevant to this wallet that has been
@@ -707,7 +707,7 @@ pub struct ScannedBlock {
     transactions: Vec<WalletTx>,
     sapling: ScannedBundles<sapling::Node, sapling::Nullifier>,
     #[cfg(feature = "orchard")]
-    orchard: ScannedBundles<orchard::note::NoteCommitment, orchard::note::Nullifier>,
+    orchard: ScannedBundles<orchard::tree::MerkleHashOrchard, orchard::note::Nullifier>,
 }
 
 impl ScannedBlock {
@@ -719,7 +719,7 @@ impl ScannedBlock {
         transactions: Vec<WalletTx>,
         sapling: ScannedBundles<sapling::Node, sapling::Nullifier>,
         #[cfg(feature = "orchard")] orchard: ScannedBundles<
-            orchard::note::NoteCommitment,
+            orchard::tree::MerkleHashOrchard,
             orchard::note::Nullifier,
         >,
     ) -> Self {
@@ -763,7 +763,7 @@ impl ScannedBlock {
     #[cfg(feature = "orchard")]
     pub fn orchard(
         &self,
-    ) -> &ScannedBundles<orchard::note::NoteCommitment, orchard::note::Nullifier> {
+    ) -> &ScannedBundles<orchard::tree::MerkleHashOrchard, orchard::note::Nullifier> {
         &self.orchard
     }
 

zcash_client_backend/src/data_api/chain.rs

@@ -150,8 +150,7 @@ use zcash_primitives::consensus::{self, BlockHeight};
 use crate::{
     data_api::{NullifierQuery, WalletWrite},
     proto::compact_formats::CompactBlock,
-    scan::BatchRunner,
-    scanning::{add_block_to_runner, scan_block_with_runner, ScanningKeys},
+    scanning::{scan_block_with_runners, BatchRunners, ScanningKeys},
 };
 
 pub mod error;
@@ -213,22 +212,26 @@ pub trait BlockSource {
 /// blocks.
 #[derive(Clone, Debug)]
 pub struct ScanSummary {
-    scanned_range: Range<BlockHeight>,
-    spent_sapling_note_count: usize,
-    received_sapling_note_count: usize,
+    pub(crate) scanned_range: Range<BlockHeight>,
+    pub(crate) spent_sapling_note_count: usize,
+    pub(crate) received_sapling_note_count: usize,
+    #[cfg(feature = "orchard")]
+    pub(crate) spent_orchard_note_count: usize,
+    #[cfg(feature = "orchard")]
+    pub(crate) received_orchard_note_count: usize,
 }
 
 impl ScanSummary {
-    /// Constructs a new [`ScanSummary`] from its constituent parts.
-    pub fn from_parts(
-        scanned_range: Range<BlockHeight>,
-        spent_sapling_note_count: usize,
-        received_sapling_note_count: usize,
-    ) -> Self {
+    /// Constructs a new [`ScanSummary`] for the provided block range.
+    pub fn for_range(scanned_range: Range<BlockHeight>) -> Self {
         Self {
             scanned_range,
-            spent_sapling_note_count,
-            received_sapling_note_count,
+            spent_sapling_note_count: 0,
+            received_sapling_note_count: 0,
+            #[cfg(feature = "orchard")]
+            spent_orchard_note_count: 0,
+            #[cfg(feature = "orchard")]
+            received_orchard_note_count: 0,
         }
     }
 
@@ -252,6 +255,16 @@ impl ScanSummary {
     pub fn received_sapling_note_count(&self) -> usize {
         self.received_sapling_note_count
     }
+
+    #[cfg(feature = "orchard")]
+    pub fn spent_orchard_note_count(&self) -> usize {
+        self.spent_orchard_note_count
+    }
+
+    #[cfg(feature = "orchard")]
+    pub fn received_orchard_note_count(&self) -> usize {
+        self.received_orchard_note_count
+    }
 }
 
 /// Scans at most `limit` blocks from the provided block source for in order to find transactions
@@ -275,31 +288,23 @@ where
     DbT: WalletWrite,
 {
     // Fetch the UnifiedFullViewingKeys we are tracking
-    let mut scanning_keys = ScanningKeys::from_account_ufvks(
-        data_db
-            .get_unified_full_viewing_keys()
-            .map_err(Error::Wallet)?,
-    );
-
-    let mut sapling_runner = BatchRunner::<_, _, _, _, ()>::new(
-        100,
-        scanning_keys
-            .sapling_keys()
-            .iter()
-            .map(|(id, key)| (*id, key.prepare())),
-    );
+    let account_ufvks = data_db
+        .get_unified_full_viewing_keys()
+        .map_err(Error::Wallet)?;
+    let mut scanning_keys = ScanningKeys::from_account_ufvks(account_ufvks);
+    let mut runners = BatchRunners::<_, (), ()>::for_keys(100, &scanning_keys);
 
     block_source.with_blocks::<_, DbT::Error>(
         Some(from_height),
         Some(limit),
         |block: CompactBlock| {
-            add_block_to_runner(params, block, &mut sapling_runner);
+            runners.add_block(params, block);
 
             Ok(())
         },
     )?;
 
-    sapling_runner.flush();
+    runners.flush();
 
     let mut prior_block_metadata = if from_height > BlockHeight::from(0) {
         data_db
@@ -317,42 +322,37 @@ where
     );
 
     let mut scanned_blocks = vec![];
-    let mut scan_end_height = from_height;
-    let mut received_note_count = 0;
-    let mut spent_note_count = 0;
+    let mut scan_summary = ScanSummary::for_range(from_height..from_height);
     block_source.with_blocks::<_, DbT::Error>(
         Some(from_height),
         Some(limit),
         |block: CompactBlock| {
-            scan_end_height = block.height() + 1;
-            let scanned_block = scan_block_with_runner(
+            scan_summary.scanned_range.end = block.height() + 1;
+            let scanned_block = scan_block_with_runners::<_, _, (), ()>(
                 params,
                 block,
                 &scanning_keys,
                 prior_block_metadata.as_ref(),
-                Some(&mut sapling_runner),
+                Some(&mut runners),
             )
             .map_err(Error::Scan)?;
 
-            let (s, r) = scanned_block
-                .transactions
-                .iter()
-                .fold((0, 0), |(s, r), wtx| {
-                    (
-                        s + wtx.sapling_spends().len(),
-                        r + wtx.sapling_outputs().len(),
-                    )
-                });
-            spent_note_count += s;
-            received_note_count += r;
-
-            let spent_nf: Vec<&sapling::Nullifier> = scanned_block
+            for wtx in &scanned_block.transactions {
+                scan_summary.spent_sapling_note_count += wtx.sapling_spends().len();
+                scan_summary.received_sapling_note_count += wtx.sapling_outputs().len();
+                #[cfg(feature = "orchard")]
+                {
+                    scan_summary.spent_orchard_note_count += wtx.orchard_spends().len();
+                    scan_summary.received_orchard_note_count += wtx.orchard_outputs().len();
+                }
+            }
+
+            let sapling_spent_nf: Vec<&sapling::Nullifier> = scanned_block
                 .transactions
                 .iter()
                 .flat_map(|tx| tx.sapling_spends().iter().map(|spend| spend.nf()))
                 .collect();
-
-            scanning_keys.retain_sapling_nullifiers(|(_, nf)| !spent_nf.contains(&nf));
+            scanning_keys.retain_sapling_nullifiers(|(_, nf)| !sapling_spent_nf.contains(&nf));
             scanning_keys.extend_sapling_nullifiers(scanned_block.transactions.iter().flat_map(
                 |tx| {
                     tx.sapling_outputs().iter().flat_map(|out| {
@@ -364,6 +364,27 @@ where
                 },
             ));
 
+            #[cfg(feature = "orchard")]
+            {
+                let orchard_spent_nf: Vec<&orchard::note::Nullifier> = scanned_block
+                    .transactions
+                    .iter()
+                    .flat_map(|tx| tx.orchard_spends().iter().map(|spend| spend.nf()))
+                    .collect();
+
+                scanning_keys.retain_orchard_nullifiers(|(_, nf)| !orchard_spent_nf.contains(&nf));
+                scanning_keys.extend_orchard_nullifiers(
+                    scanned_block.transactions.iter().flat_map(|tx| {
+                        tx.orchard_outputs().iter().flat_map(|out| {
+                            out.key_source()
+                                .account()
+                                .zip(out.nf().copied())
+                                .into_iter()
+                        })
+                    }),
+                );
+            }
+
             prior_block_metadata = Some(scanned_block.to_block_metadata());
             scanned_blocks.push(scanned_block);
 
@@ -372,11 +393,7 @@ where
     )?;
 
     data_db.put_blocks(scanned_blocks).map_err(Error::Wallet)?;
-    Ok(ScanSummary::from_parts(
-        from_height..scan_end_height,
-        spent_note_count,
-        received_note_count,
-    ))
+    Ok(scan_summary)
 }
 
 #[cfg(feature = "test-dependencies")]

zcash_client_backend/src/proto.rs

@@ -10,7 +10,7 @@ use std::{
 use incrementalmerkletree::frontier::CommitmentTree;
 
 use nonempty::NonEmpty;
-use sapling::{self, note::ExtractedNoteCommitment, Node, Nullifier, NOTE_COMMITMENT_TREE_DEPTH};
+use sapling::{self, note::ExtractedNoteCommitment, Node, NOTE_COMMITMENT_TREE_DEPTH};
 use zcash_primitives::{
     block::{BlockHash, BlockHeader},
     consensus::{self, BlockHeight, Parameters},
@@ -169,8 +169,49 @@ impl TryFrom<&compact_formats::CompactSaplingOutput>
 }
 
 impl compact_formats::CompactSaplingSpend {
-    pub fn nf(&self) -> Result<Nullifier, ()> {
-        Nullifier::from_slice(&self.nf).map_err(|_| ())
+    pub fn nf(&self) -> Result<sapling::Nullifier, ()> {
+        sapling::Nullifier::from_slice(&self.nf).map_err(|_| ())
+    }
+}
+
+#[cfg(feature = "orchard")]
+impl TryFrom<&compact_formats::CompactOrchardAction> for orchard::note_encryption::CompactAction {
+    type Error = ();
+
+    fn try_from(value: &compact_formats::CompactOrchardAction) -> Result<Self, Self::Error> {
+        Ok(orchard::note_encryption::CompactAction::from_parts(
+            value.nf()?,
+            value.cmx()?,
+            value.ephemeral_key()?,
+            value.ciphertext[..].try_into().map_err(|_| ())?,
+        ))
+    }
+}
+
+#[cfg(feature = "orchard")]
+impl compact_formats::CompactOrchardAction {
+    pub fn cmx(&self) -> Result<orchard::note::ExtractedNoteCommitment, ()> {
+        Option::from(orchard::note::ExtractedNoteCommitment::from_bytes(
+            &self.cmx[..].try_into().map_err(|_| ())?,
+        ))
+        .ok_or(())
+    }
+
+    pub fn nf(&self) -> Result<orchard::note::Nullifier, ()> {
+        let nf_bytes: [u8; 32] = self.nullifier[..].try_into().map_err(|_| ())?;
+        Option::from(orchard::note::Nullifier::from_bytes(&nf_bytes)).ok_or(())
+    }
+
+    /// Returns the ephemeral public key for this output.
+    ///
+    /// A convenience method that parses [`CompactOutput.epk`].
+    ///
+    /// [`CompactOutput.epk`]: #structfield.epk
+    pub fn ephemeral_key(&self) -> Result<EphemeralKeyBytes, ()> {
+        self.ephemeral_key[..]
+            .try_into()
+            .map(EphemeralKeyBytes)
+            .map_err(|_| ())
     }
 }
 

zcash_client_backend/src/scan.rs

@@ -404,12 +404,16 @@ where
     /// `replier` will be called with the result of every output.
     fn add_outputs(
         &mut self,
-        domain: impl Fn() -> D,
+        domain: impl Fn(&Output) -> D,
         outputs: &[Output],
         replier: channel::Sender<OutputItem<KeyId, D, Dec::Memo>>,
     ) {
-        self.outputs
-            .extend(outputs.iter().cloned().map(|output| (domain(), output)));
+        self.outputs.extend(
+            outputs
+                .iter()
+                .cloned()
+                .map(|output| (domain(&output), output)),
+        );
         self.repliers.extend((0..outputs.len()).map(|output_index| {
             OutputReplier(OutputIndex {
                 output_index,
@@ -539,7 +543,7 @@ where
         &mut self,
         block_tag: BlockHash,
         txid: TxId,
-        domain: impl Fn() -> D,
+        domain: impl Fn(&Output) -> D,
         outputs: &[Output],
     ) {
         let (tx, rx) = channel::unbounded();