Add BatchDomain::batch_ka_agree_dec, a per-viewing-key batched agreement hook

CHANGELOG.md

@@ -9,6 +9,13 @@ and this library adheres to Rust's notion of
 ### Added
 - A blanket `impl<D, O, const CIPHERTEXT_SIZE: usize> ShieldedOutput<D, CIPHERTEXT_SIZE> for &O`
     where `D: Domain, O: ShieldedOutput<D, CIPHERTEXT_SIZE>`
+- `BatchDomain::batch_ka_agree_dec`, a provided method that computes
+  `Domain::ka_agree_dec` for a batch of prepared ephemeral keys against a single
+  incoming viewing key. Domains for which same-scalar multiplications can share
+  work can override it; the default implementation (and therefore the behavior
+  of existing `BatchDomain` implementations) is the per-item computation. The
+  batch decryption functions in the `batch` module now derive their shared
+  secrets through this method, one call per viewing key.
 
 ### Changed
 - **Breaking change:** removed the constants `COMPACT_NOTE_SIZE`,

src/batch.rs

@@ -54,16 +54,36 @@ where
     // Fetch the ephemeral keys for each output, and batch-parse and prepare them.
     let ephemeral_keys = D::batch_epk(outputs.iter().map(|(_, output)| output.ephemeral_key()));
 
-    // Derive the shared secrets for all combinations of (ivk, output).
-    // The scalar multiplications cannot benefit from batching.
-    let items = ephemeral_keys.iter().flat_map(|(epk, ephemeral_key)| {
-        ivks.iter().map(move |ivk| {
-            (
-                epk.as_ref().map(|epk| D::ka_agree_dec(ivk, epk)),
-                ephemeral_key,
-            )
+    // Derive the shared secrets for all combinations of (ivk, output), one batched
+    // same-key agreement per ivk: domains for which same-scalar multiplications can
+    // share work accelerate here, and the default `batch_ka_agree_dec` implementation
+    // is exactly the previous per-item computation.
+    // Reassembly below is in the (output-major, ivk-minor) order the batch-KDF
+    // expects, moving values out of the per-ivk columns (`SharedSecret` need not be
+    // `Clone`).
+    let mut columns: Vec<_> = ivks
+        .iter()
+        .map(|ivk| {
+            D::batch_ka_agree_dec(ivk, ephemeral_keys.iter().map(|(epk, _)| epk.as_ref()))
+                .into_iter()
         })
-    });
+        .collect();
+    let mut secrets: Vec<Option<D::SharedSecret>> =
+        Vec::with_capacity(ephemeral_keys.len() * ivks.len());
+    for _ in 0..ephemeral_keys.len() {
+        for column in columns.iter_mut() {
+            secrets.push(
+                column
+                    .next()
+                    .expect("all columns have one entry per output"),
+            );
+        }
+    }
+    let items = secrets.into_iter().zip(
+        ephemeral_keys
+            .iter()
+            .flat_map(|(_, ephemeral_key)| core::iter::repeat(ephemeral_key).take(ivks.len())),
+    );
 
     // Run the batch-KDF to obtain the symmetric keys from the shared secrets.
     let keys = D::batch_kdf(items);

src/lib.rs

@@ -349,6 +349,29 @@ pub trait BatchDomain: Domain {
             })
             .collect()
     }
+
+    /// Computes `Self::ka_agree_dec` on a batch of prepared ephemeral keys against a
+    /// single incoming viewing key.
+    ///
+    /// For each item, if the prepared ephemeral key is `None` (i.e. its encoding could
+    /// not be parsed), this returns `None` at that position.
+    ///
+    /// Trial decryption multiplies many ephemeral keys by the same viewing key, so
+    /// domains for which same-scalar multiplications can share work (for example,
+    /// lockstep ladders over a shared batched field inversion) can override this to
+    /// reduce the cost of the scalar multiplications, which dominate batched trial
+    /// decryption. The default implementation performs the per-item computation.
+    fn batch_ka_agree_dec<'a>(
+        ivk: &Self::IncomingViewingKey,
+        epks: impl Iterator<Item = Option<&'a Self::PreparedEphemeralPublicKey>>,
+    ) -> Vec<Option<Self::SharedSecret>>
+    where
+        Self::PreparedEphemeralPublicKey: 'a,
+    {
+        // Default implementation: do the non-batched thing.
+        epks.map(|epk| epk.map(|epk| Self::ka_agree_dec(ivk, epk)))
+            .collect()
+    }
 }
 
 /// Trait that provides access to the components of an encrypted transaction output.