Add allocation-free EdwardsPoint::compress_batch() (#832)

Author: daxpedda

curve25519-dalek/benches/dalek_benchmarks.rs

@@ -34,7 +34,7 @@ mod edwards_benches {
                     let mut rng = OsRng.unwrap_err();
                     let points: Vec<EdwardsPoint> =
                         (0..size).map(|_| EdwardsPoint::random(&mut rng)).collect();
-                    b.iter(|| EdwardsPoint::compress_batch(&points));
+                    b.iter(|| EdwardsPoint::compress_batch_alloc(&points));
                 },
             );
         }

curve25519-dalek/src/edwards.rs

@@ -599,7 +599,7 @@ impl EdwardsPoint {
 
         // Compute the denominators in a batch
         let mut denominators = eds.iter().map(|p| &p.Z - &p.Y).collect::<Vec<_>>();
-        FieldElement::batch_invert(&mut denominators);
+        FieldElement::invert_batch_alloc(&mut denominators);
 
         // Now compute the Montgomery u coordinate for every point
         let mut ret = Vec::with_capacity(eds.len());
@@ -616,12 +616,24 @@ impl EdwardsPoint {
         self.to_affine().compress()
     }
 
+    /// Compress several `EdwardsPoint`s into `CompressedEdwardsY` format, using a batch inversion
+    /// for a significant speedup.
+    pub fn compress_batch<const N: usize>(inputs: &[EdwardsPoint; N]) -> [CompressedEdwardsY; N] {
+        let mut zs: [_; N] = core::array::from_fn(|i| inputs[i].Z);
+        FieldElement::invert_batch(&mut zs);
+
+        core::array::from_fn(|i| {
+            let x = &inputs[i].X * &zs[i];
+            let y = &inputs[i].Y * &zs[i];
+            AffinePoint { x, y }.compress()
+        })
+    }
     /// Compress several `EdwardsPoint`s into `CompressedEdwardsY` format, using a batch inversion
     /// for a significant speedup.
     #[cfg(feature = "alloc")]
-    pub fn compress_batch(inputs: &[EdwardsPoint]) -> Vec<CompressedEdwardsY> {
+    pub fn compress_batch_alloc(inputs: &[EdwardsPoint]) -> Vec<CompressedEdwardsY> {
         let mut zs = inputs.iter().map(|input| input.Z).collect::<Vec<_>>();
-        FieldElement::batch_invert(&mut zs);
+        FieldElement::invert_batch_alloc(&mut zs);
 
         inputs
             .iter()
@@ -2175,30 +2187,49 @@ mod test {
             CompressedEdwardsY::identity()
         );
 
+        assert_eq!(
+            EdwardsPoint::compress_batch(&[EdwardsPoint::identity()]),
+            [CompressedEdwardsY::identity()]
+        );
         #[cfg(feature = "alloc")]
-        {
-            let compressed = EdwardsPoint::compress_batch(&[EdwardsPoint::identity()]);
-            assert_eq!(&compressed, &[CompressedEdwardsY::identity()]);
-        }
+        assert_eq!(
+            &EdwardsPoint::compress_batch_alloc(&[EdwardsPoint::identity()]),
+            &[CompressedEdwardsY::identity()]
+        );
     }
 
-    #[cfg(all(feature = "alloc", feature = "rand_core"))]
+    #[cfg(feature = "rand_core")]
     #[test]
     fn compress_batch() {
         let mut rng = rand::rng();
 
         // TODO(tarcieri): proptests?
-        // Make some points deterministically then randomly
-        let mut points = (1u64..16)
-            .map(|n| constants::ED25519_BASEPOINT_POINT * Scalar::from(n))
-            .collect::<Vec<_>>();
-        points.extend(core::iter::repeat_with(|| EdwardsPoint::random(&mut rng)).take(100));
-        let compressed = EdwardsPoint::compress_batch(&points);
+
+        // Make some test points deterministically then randomly
+        const TEST_VEC_LEN: usize = 117;
+        let points: [EdwardsPoint; TEST_VEC_LEN] = core::array::from_fn(|i| {
+            if i < 17 {
+                // The first 17 are multiple of the basepoint
+                constants::ED25519_BASEPOINT_POINT * Scalar::from(i as u64)
+            } else {
+                // The rest are random
+                EdwardsPoint::random(&mut rng)
+            }
+        });
+
+        // Compress the points individually. This is our reference result
+        let expected_compressed = core::array::from_fn(|i| points[i].compress());
 
         // Check that the batch-compressed points match the individually compressed ones
-        for (point, compressed) in points.iter().zip(&compressed) {
-            assert_eq!(&point.compress(), compressed);
-        }
+        assert_eq!(EdwardsPoint::compress_batch(&points), expected_compressed);
+
+        // Check that the batch-compressed (with alloc) points match the individually compressed
+        // ones
+        #[cfg(feature = "alloc")]
+        assert_eq!(
+            EdwardsPoint::compress_batch_alloc(&points),
+            expected_compressed
+        );
     }
 
     #[test]

curve25519-dalek/src/field.rs

@@ -209,17 +209,39 @@ impl FieldElement {
         (t19, t3)
     }
 
+    /// Given a slice of pub(crate)lic `FieldElements`, replace each with its inverse.
+    ///
+    /// When an input `FieldElement` is zero, its value is unchanged.
+    pub(crate) fn invert_batch<const N: usize>(inputs: &mut [FieldElement; N]) {
+        let mut scratch = [FieldElement::ONE; N];
+
+        Self::internal_invert_batch(inputs, &mut scratch);
+    }
+
     /// Given a slice of pub(crate)lic `FieldElements`, replace each with its inverse.
     ///
     /// When an input `FieldElement` is zero, its value is unchanged.
     #[cfg(feature = "alloc")]
-    pub(crate) fn batch_invert(inputs: &mut [FieldElement]) {
+    pub(crate) fn invert_batch_alloc(inputs: &mut [FieldElement]) {
+        let n = inputs.len();
+        let mut scratch = vec![FieldElement::ONE; n];
+
+        Self::internal_invert_batch(inputs, &mut scratch);
+    }
+
+    /// Given a slice of pub(crate)lic `FieldElements`, replace each with its inverse. `scratch` can
+    /// contain anything, so long as its length is the same as `inputs`.
+    ///
+    /// When an input `FieldElement` is zero, its value is unchanged.
+    ///
+    /// # Panics
+    /// Panics when `scratch.len() != inputs.len()`
+    fn internal_invert_batch(inputs: &mut [FieldElement], scratch: &mut [FieldElement]) {
         // Montgomery’s Trick and Fast Implementation of Masked AES
         // Genelle, Prouff and Quisquater
         // Section 3.2
 
-        let n = inputs.len();
-        let mut scratch = vec![FieldElement::ONE; n];
+        debug_assert_eq!(inputs.len(), scratch.len());
 
         // Keep an accumulator of all of the previous products
         let mut acc = FieldElement::ONE;
@@ -240,12 +262,12 @@ impl FieldElement {
 
         // Pass through the vector backwards to compute the inverses
         // in place
-        for (input, scratch) in inputs.iter_mut().rev().zip(scratch.into_iter().rev()) {
+        for (input, scratch) in inputs.iter_mut().rev().zip(scratch.iter_mut().rev()) {
             let tmp = &acc * input;
             // input <- acc * scratch, then acc <- tmp
             // Again, we skip zeros in a constant-time way
             let nz = !input.is_zero();
-            input.conditional_assign(&(&acc * &scratch), nz);
+            input.conditional_assign(&(&acc * scratch), nz);
             acc.conditional_assign(&tmp, nz);
         }
     }
@@ -559,7 +581,7 @@ mod test {
 
     #[test]
     #[cfg(feature = "alloc")]
-    fn batch_invert_a_matches_nonbatched() {
+    fn invert_batch_a_matches_nonbatched() {
         let a = FieldElement::from_bytes(&A_BYTES);
         let ap58 = FieldElement::from_bytes(&AP58_BYTES);
         let asq = FieldElement::from_bytes(&ASQ_BYTES);
@@ -568,7 +590,7 @@ mod test {
         let a2 = &a + &a;
         let a_list = vec![a, ap58, asq, ainv, a0, a2];
         let mut ainv_list = a_list.clone();
-        FieldElement::batch_invert(&mut ainv_list[..]);
+        FieldElement::invert_batch_alloc(&mut ainv_list[..]);
         for i in 0..6 {
             assert_eq!(a_list[i].invert(), ainv_list[i]);
         }
@@ -677,8 +699,8 @@ mod test {
 
     #[test]
     #[cfg(feature = "alloc")]
-    fn batch_invert_empty() {
-        FieldElement::batch_invert(&mut []);
+    fn invert_batch_empty() {
+        FieldElement::invert_batch_alloc(&mut []);
     }
 
     // The following two consts were generated with the following sage script:

curve25519-dalek/src/ristretto.rs

@@ -606,7 +606,7 @@ impl RistrettoPoint {
 
         let mut invs: Vec<FieldElement> = states.iter().map(|state| state.efgh()).collect();
 
-        FieldElement::batch_invert(&mut invs[..]);
+        FieldElement::invert_batch_alloc(&mut invs[..]);
 
         states
             .iter()