zcash_note_encryption: Add support for batch-parsing `ephemeral_key`

This is useful when the underlying curve requires an inversion to parse
an encoded point (such as for Jubjub).

components/zcash_note_encryption/src/batch.rs

@@ -35,18 +35,15 @@ fn batch_note_decryption<D: Domain, Output: ShieldedOutput<D>, F, FR>(
 where
     F: Fn(&D, &D::IncomingViewingKey, &EphemeralKeyBytes, &Output, D::SymmetricKey) -> Option<FR>,
 {
-    // Fetch the ephemeral keys for each output.
-    let ephemeral_keys: Vec<_> = outputs
-        .iter()
-        .map(|(_, output)| output.ephemeral_key())
-        .collect();
+    // Fetch the ephemeral keys for each output and batch-parse them.
+    let ephemeral_keys = D::batch_epk(outputs.iter().map(|(_, output)| output.ephemeral_key()));
 
     // Derive the shared secrets for all combinations of (ivk, output).
-    // None of this work can benefit from batching.
+    // The scalar multiplications cannot benefit from batching.
     let items = ivks.iter().flat_map(|ivk| {
-        ephemeral_keys.iter().map(move |ephemeral_key| {
+        ephemeral_keys.iter().map(move |(epk, ephemeral_key)| {
             (
-                D::epk(ephemeral_key).map(|epk| D::ka_agree_dec(ivk, &epk)),
+                epk.as_ref().map(|epk| D::ka_agree_dec(ivk, epk)),
                 ephemeral_key,
             )
         })
@@ -64,7 +61,7 @@ where
                 .zip(outputs.iter())
         })
         .zip(keys)
-        .map(|(((ivk, ephemeral_key), (domain, output)), key)| {
+        .map(|(((ivk, (_, ephemeral_key)), (domain, output)), key)| {
             // The `and_then` propagates any potential rejection from `D::epk`.
             key.and_then(|key| decrypt_inner(domain, ivk, ephemeral_key, output, key))
         })

components/zcash_note_encryption/src/lib.rs

@@ -139,6 +139,22 @@ pub trait Domain {
 
     fn epk(ephemeral_key: &EphemeralKeyBytes) -> Option<Self::EphemeralPublicKey>;
 
+    /// Computes `Self::epk` on a batch of ephemeral keys.
+    ///
+    /// This is useful for protocols where the underlying curve requires an inversion to
+    /// parse an encoded point.
+    ///
+    /// For usability, this returns tuples of the ephemeral keys and the result of parsing
+    /// them.
+    fn batch_epk(
+        ephemeral_keys: impl Iterator<Item = EphemeralKeyBytes>,
+    ) -> Vec<(Option<Self::EphemeralPublicKey>, EphemeralKeyBytes)> {
+        // Default implementation: do the non-batched thing.
+        ephemeral_keys
+            .map(|ephemeral_key| (Self::epk(&ephemeral_key), ephemeral_key))
+            .collect()
+    }
+
     fn check_epk_bytes<F: Fn(&Self::EphemeralSecretKey) -> NoteValidity>(
         note: &Self::Note,
         check: F,