zcash_client_backend: Add ScanningKeys type in preparation of suppl…

…ying Orchard keys.

zcash_client_backend/src/data_api/chain.rs

@@ -143,15 +143,15 @@
 //! # }
 //! ```
 
-use std::{collections::HashMap, ops::Range};
+use std::ops::Range;
 
 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, SaplingScanningKey, ScanningKey},
+    scanning::{add_block_to_runner, scan_block_with_runner, ScanningKey, ScanningKeys},
 };
 
 pub mod error;
@@ -275,29 +275,18 @@ where
     DbT: WalletWrite,
 {
     // Fetch the UnifiedFullViewingKeys we are tracking
-    let ufvks = data_db
-        .get_unified_full_viewing_keys()
-        .map_err(Error::Wallet)?;
-    // TODO: Change `scan_block` to also scan Orchard.
-    // https://github.com/zcash/librustzcash/issues/403
-    let sapling_ivks = ufvks
-        .iter()
-        .filter_map(|(account, ufvk)| ufvk.sapling().map(move |k| (account, k)))
-        .flat_map(|(account, dfvk)| {
-            SaplingScanningKey::from_account_dfvk(dfvk, Some(*account)).into_iter()
-        })
-        .map(|key| {
-            (
-                key.account_scope()
-                    .expect("Account ID is available for each key"),
-                key,
-            )
-        })
-        .collect::<HashMap<_, _>>();
+    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,
-        sapling_ivks.iter().map(|(id, key)| (*id, key.prepare())),
+        scanning_keys
+            .sapling_keys()
+            .iter()
+            .map(|(id, key)| (*id, key.prepare())),
     );
 
     block_source.with_blocks::<_, DbT::Error>(
@@ -321,9 +310,11 @@ where
     };
 
     // Get the nullifiers for the unspent notes we are tracking
-    let mut sapling_nullifiers = data_db
-        .get_sapling_nullifiers(NullifierQuery::Unspent)
-        .map_err(Error::Wallet)?;
+    scanning_keys.extend_sapling_nullifiers(
+        data_db
+            .get_sapling_nullifiers(NullifierQuery::Unspent)
+            .map_err(Error::Wallet)?,
+    );
 
     let mut scanned_blocks = vec![];
     let mut scan_end_height = from_height;
@@ -337,8 +328,7 @@ where
             let scanned_block = scan_block_with_runner(
                 params,
                 block,
-                &sapling_ivks,
-                &sapling_nullifiers,
+                &scanning_keys,
                 prior_block_metadata.as_ref(),
                 Some(&mut sapling_runner),
             )
@@ -362,15 +352,17 @@ where
                 .flat_map(|tx| tx.sapling_spends().iter().map(|spend| spend.nf()))
                 .collect();
 
-            sapling_nullifiers.retain(|(_, nf)| !spent_nf.contains(&nf));
-            sapling_nullifiers.extend(scanned_block.transactions.iter().flat_map(|tx| {
-                tx.sapling_outputs().iter().flat_map(|out| {
-                    out.recipient_key_meta()
-                        .zip(out.nf())
-                        .into_iter()
-                        .map(|((account_id, _), nf)| (account_id, *nf))
-                })
-            }));
+            scanning_keys.retain_sapling_nullifiers(|(_, nf)| !spent_nf.contains(&nf));
+            scanning_keys.extend_sapling_nullifiers(scanned_block.transactions.iter().flat_map(
+                |tx| {
+                    tx.sapling_outputs().iter().flat_map(|out| {
+                        out.recipient_key_meta()
+                            .zip(out.nf())
+                            .into_iter()
+                            .map(|((account_id, _), nf)| (account_id, *nf))
+                    })
+                },
+            ));
 
             prior_block_metadata = Some(scanned_block.to_block_metadata());
             scanned_blocks.push(scanned_block);