Merge pull request #178 from zcash/batch-note-decryption

Speed up batched note decryption

Cargo.toml

@@ -84,5 +84,5 @@ debug = true
 
 [patch.crates-io]
 halo2 = { git = "https://github.com/zcash/halo2.git", rev = "27c4187673a9c6ade13fbdbd4f20955530c22d7f" }
-zcash_note_encryption = { git = "https://github.com/zcash/librustzcash.git", rev = "cc533a9da4f6a7209a7be05f82b12a03969152c9" }
+zcash_note_encryption = { git = "https://github.com/zcash/librustzcash.git", rev = "13b023387bafdc7b5712c933dc0e16ee94b96a6a" }
 incrementalmerkletree = { git = "https://github.com/zcash/incrementalmerkletree.git", rev = "b7bd6246122a6e9ace8edb51553fbf5228906cbb" }

benches/note_decryption.rs

@@ -1,4 +1,6 @@
-use criterion::{criterion_group, criterion_main, Criterion, Throughput};
+use std::array;
+
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
 use orchard::{
     builder::Builder,
     bundle::Flags,
@@ -9,7 +11,7 @@ use orchard::{
     Anchor, Bundle,
 };
 use rand::rngs::OsRng;
-use zcash_note_encryption::{try_compact_note_decryption, try_note_decryption};
+use zcash_note_encryption::{batch, try_compact_note_decryption, try_note_decryption};
 
 #[cfg(unix)]
 use pprof::criterion::{Output, PProfProfiler};
@@ -63,6 +65,7 @@ fn bench_note_decryption(c: &mut Criterion) {
 
     let compact = {
         let mut group = c.benchmark_group("note-decryption");
+        group.throughput(Throughput::Elements(1));
 
         group.bench_function("valid", |b| {
             b.iter(|| try_note_decryption(&domain, &valid_ivk, action).unwrap())
@@ -93,6 +96,47 @@ fn bench_note_decryption(c: &mut Criterion) {
             })
         });
     }
+
+    {
+        // Benchmark with 2 IVKs to emulate a wallet with two pools of funds.
+        let ivks = 2;
+        let valid_ivks = vec![valid_ivk; ivks];
+        let actions: Vec<_> = (0..100)
+            .map(|_| (OrchardDomain::for_action(action), action.clone()))
+            .collect();
+        let compact: Vec<_> = (0..100)
+            .map(|_| {
+                (
+                    OrchardDomain::for_action(action),
+                    CompactAction::from(action),
+                )
+            })
+            .collect();
+
+        let mut group = c.benchmark_group("batch-note-decryption");
+
+        for size in array::IntoIter::new([10, 50, 100]) {
+            group.throughput(Throughput::Elements((ivks * size) as u64));
+
+            group.bench_function(BenchmarkId::new("valid", size), |b| {
+                b.iter(|| batch::try_note_decryption(&valid_ivks, &actions[..size]))
+            });
+
+            group.bench_function(BenchmarkId::new("invalid", size), |b| {
+                b.iter(|| batch::try_note_decryption(&invalid_ivks[..ivks], &actions[..size]))
+            });
+
+            group.bench_function(BenchmarkId::new("compact-valid", size), |b| {
+                b.iter(|| batch::try_compact_note_decryption(&valid_ivks, &compact[..size]))
+            });
+
+            group.bench_function(BenchmarkId::new("compact-invalid", size), |b| {
+                b.iter(|| {
+                    batch::try_compact_note_decryption(&invalid_ivks[..ivks], &compact[..size])
+                })
+            });
+        }
+    }
 }
 
 #[cfg(unix)]

src/keys.rs

@@ -6,7 +6,7 @@ use std::mem;
 use aes::Aes256;
 use blake2b_simd::{Hash as Blake2bHash, Params};
 use fpe::ff1::{BinaryNumeralString, FF1};
-use group::GroupEncoding;
+use group::{prime::PrimeCurveAffine, Curve, GroupEncoding};
 use halo2::arithmetic::FieldExt;
 use pasta_curves::pallas;
 use rand::RngCore;
@@ -263,7 +263,7 @@ impl FullViewingKey {
 /// Defined in [Zcash Protocol Spec § 4.2.3: Orchard Key Components][orchardkeycomponents].
 ///
 /// [orchardkeycomponents]: https://zips.z.cash/protocol/nu5.pdf#orchardkeycomponents
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub struct DiversifierKey([u8; 32]);
 
 impl From<&FullViewingKey> for DiversifierKey {
@@ -345,7 +345,7 @@ impl Diversifier {
 /// decryption of notes). When we actually want to serialize ivk, we're guaranteed to get
 /// a valid base field element encoding, because we always construct ivk from an integer
 /// in the correct range.
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 struct KeyAgreementPrivateKey(NonZeroPallasScalar);
 
 impl From<&FullViewingKey> for KeyAgreementPrivateKey {
@@ -382,7 +382,7 @@ impl KeyAgreementPrivateKey {
 /// Defined in [Zcash Protocol Spec § 5.6.4.3: Orchard Raw Incoming Viewing Keys][orchardinviewingkeyencoding].
 ///
 /// [orchardinviewingkeyencoding]: https://zips.z.cash/protocol/nu5.pdf#orchardinviewingkeyencoding
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub struct IncomingViewingKey {
     dk: DiversifierKey,
     ivk: KeyAgreementPrivateKey,
@@ -567,15 +567,44 @@ impl SharedSecret {
         self.0.to_bytes()
     }
 
+    /// Only for use in batched note encryption.
+    pub(crate) fn batch_to_affine(
+        shared_secrets: Vec<Option<Self>>,
+    ) -> impl Iterator<Item = Option<pallas::Affine>> {
+        // Filter out the positions for which ephemeral_key was not a valid encoding.
+        let secrets: Vec<_> = shared_secrets
+            .iter()
+            .filter_map(|s| s.as_ref().map(|s| *(s.0)))
+            .collect();
+
+        // Batch-normalize the shared secrets.
+        let mut secrets_affine = vec![pallas::Affine::identity(); secrets.len()];
+        group::Curve::batch_normalize(&secrets, &mut secrets_affine);
+
+        // Re-insert the invalid ephemeral_key positions.
+        let mut secrets_affine = secrets_affine.into_iter();
+        shared_secrets
+            .into_iter()
+            .map(move |s| s.and_then(|_| secrets_affine.next()))
+    }
+
     /// Defined in [Zcash Protocol Spec § 5.4.5.6: Orchard Key Agreement][concreteorchardkdf].
     ///
     /// [concreteorchardkdf]: https://zips.z.cash/protocol/nu5.pdf#concreteorchardkdf
     pub(crate) fn kdf_orchard(self, ephemeral_key: &EphemeralKeyBytes) -> Blake2bHash {
+        Self::kdf_orchard_inner(self.0.to_affine(), ephemeral_key)
+    }
+
+    /// Only for direct use in batched note encryption.
+    pub(crate) fn kdf_orchard_inner(
+        secret: pallas::Affine,
+        ephemeral_key: &EphemeralKeyBytes,
+    ) -> Blake2bHash {
         Params::new()
             .hash_length(32)
             .personal(KDF_ORCHARD_PERSONALIZATION)
             .to_state()
-            .update(&self.0.to_bytes())
+            .update(&secret.to_bytes())
             .update(&ephemeral_key.0)
             .finalize()
     }

src/note_encryption.rs

@@ -141,6 +141,19 @@ impl Domain for OrchardDomain {
         secret.kdf_orchard(ephemeral_key)
     }
 
+    fn batch_kdf<'a>(
+        items: impl Iterator<Item = (Option<Self::SharedSecret>, &'a EphemeralKeyBytes)>,
+    ) -> Vec<Option<Self::SymmetricKey>> {
+        let (shared_secrets, ephemeral_keys): (Vec<_>, Vec<_>) = items.unzip();
+
+        SharedSecret::batch_to_affine(shared_secrets)
+            .zip(ephemeral_keys.into_iter())
+            .map(|(secret, ephemeral_key)| {
+                secret.map(|dhsecret| SharedSecret::kdf_orchard_inner(dhsecret, ephemeral_key))
+            })
+            .collect()
+    }
+
     fn note_plaintext_bytes(
         note: &Self::Note,
         _: &Self::Recipient,