From 8e13986101d9a1574b17b870e93faaca85a4d834 Mon Sep 17 00:00:00 2001
From: Jack Grigg <jack@electriccoin.co>
Date: Wed, 11 Aug 2021 01:19:58 +0100
Subject: [PATCH] Implement `Domain::batch_epk` for note decryption

Improves throughput of batched trial decryption by around 10%.
---
 src/keys.rs            | 23 +++++++++++++++++++++--
 src/note_encryption.rs | 18 ++++++++++++++++++
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/src/keys.rs b/src/keys.rs
index 5b34656b..4ed263a9 100644
--- a/src/keys.rs
+++ b/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;
@@ -567,15 +567,34 @@ impl SharedSecret {
         self.0.to_bytes()
     }
 
+    /// Only for use in batched note encryption.
+    pub(crate) fn batch_to_affine(shared_secrets: &[Option<Self>]) -> Vec<pallas::Affine> {
+        let secrets: Vec<_> = shared_secrets
+            .iter()
+            .filter_map(|s| s.as_ref().map(|s| *(s.0)))
+            .collect();
+        let mut secrets_affine = vec![pallas::Affine::identity(); shared_secrets.len()];
+        group::Curve::batch_normalize(&secrets, &mut secrets_affine);
+        secrets_affine
+    }
+
     /// 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()
     }
diff --git a/src/note_encryption.rs b/src/note_encryption.rs
index eed1b409..8836f4e3 100644
--- a/src/note_encryption.rs
+++ b/src/note_encryption.rs
@@ -141,6 +141,24 @@ 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();
+
+        let secrets_affine = SharedSecret::batch_to_affine(&shared_secrets);
+
+        let mut secrets_affine = secrets_affine.into_iter();
+        shared_secrets
+            .into_iter()
+            .map(|s| s.and_then(|_| secrets_affine.next()))
+            .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,