Simplify BatchInvert trait

Companion PR to RustCrypto/traits#2455, which switches to in-place batch
inversion.

This massively simplifies the batch normalization implementation which
no longer requires a litany of generic parameters to abstract over
slices versus arrays because it accepts a temporary storage buffer the
caller can allocate however they choose.

It requires every curve add an `impl BatchInvert for FieldElement` but
also lets curves provide their own implementation, which is needed for
`k256` due to its lazy normalization.

Author: tarcieri

Cargo.lock

@@ -28,3 +28,5 @@ hash2curve = { path = "hash2curve" }
 primefield = { path = "primefield" }
 primeorder = { path = "primeorder" }
 wnaf = { path = "wnaf" }
+
+elliptic-curve = { git = "https://github.com/RustCrypto/traits", branch = "batchinvert-redux" }

Cargo.toml

@@ -17,6 +17,7 @@ use crate::U256;
 use elliptic_curve::{
     bigint::cpubits,
     ff::PrimeField,
+    ops::BatchInvert,
     subtle::{Choice, ConstantTimeEq, CtOption},
 };
 
@@ -92,6 +93,8 @@ primefield::fiat_monty_field_arithmetic! {
     selectnz: fiat_bignp256_selectznz
 }
 
+impl BatchInvert for FieldElement {}
+
 #[cfg(test)]
 mod tests {
     use super::{FieldElement, U256};

bignp256/src/arithmetic/field.rs

@@ -14,6 +14,7 @@ use crate::U256;
 use elliptic_curve::{
     bigint::cpubits,
     ff::PrimeField,
+    ops::BatchInvert,
     subtle::{Choice, ConstantTimeEq, CtOption},
 };
 
@@ -94,6 +95,8 @@ primefield::fiat_monty_field_arithmetic! {
     selectnz: fiat_bp256_selectznz
 }
 
+impl BatchInvert for FieldElement {}
+
 #[cfg(test)]
 mod tests {
     use super::{FieldElement, U256};

bp256/src/arithmetic/field.rs

@@ -14,6 +14,7 @@ use crate::U384;
 use elliptic_curve::{
     bigint::cpubits,
     ff::PrimeField,
+    ops::BatchInvert,
     subtle::{Choice, ConstantTimeEq, CtOption},
 };
 
@@ -94,6 +95,8 @@ primefield::fiat_monty_field_arithmetic! {
     selectnz: fiat_bp384_selectznz
 }
 
+impl BatchInvert for FieldElement {}
+
 #[cfg(test)]
 mod tests {
     use super::{FieldElement, U384};

bp384/src/arithmetic/field.rs

@@ -322,8 +322,9 @@ impl CurveGroup for EdwardsPoint {
     #[inline]
     fn batch_normalize(projective: &[Self], affine: &mut [Self::Affine]) {
         assert_eq!(projective.len(), affine.len());
-        let mut zs = alloc::vec![FieldElement::ONE; projective.len()];
-        batch_normalize_generic(projective, zs.as_mut_slice(), affine);
+        let mut zs = vec![FieldElement::ZERO; projective.len()];
+        let mut scratch = vec![FieldElement::ZERO; projective.len()];
+        batch_normalize_generic(projective, &mut zs, &mut scratch, affine);
     }
 }
 
@@ -768,9 +769,10 @@ impl<const N: usize> BatchNormalize<[EdwardsPoint; N]> for EdwardsPoint {
 
     #[inline]
     fn batch_normalize(points: &[Self; N]) -> [<Self as CurveGroup>::Affine; N] {
-        let zs = [FieldElement::ONE; N];
+        let mut zs = [FieldElement::ZERO; N];
+        let mut scratch = [FieldElement::ZERO; N];
         let mut affine_points = [AffinePoint::IDENTITY; N];
-        batch_normalize_generic(points, zs, &mut affine_points);
+        batch_normalize_generic(points, &mut zs, &mut scratch, &mut affine_points);
         affine_points
     }
 }
@@ -783,43 +785,36 @@ impl BatchNormalize<[EdwardsPoint]> for EdwardsPoint {
     fn batch_normalize(points: &[Self]) -> Vec<<Self as CurveGroup>::Affine> {
         use alloc::vec;
 
-        let mut zs = vec![FieldElement::ONE; points.len()];
+        let mut zs = vec![FieldElement::ZERO; points.len()];
+        let mut scratch = vec![FieldElement::ZERO; points.len()];
         let mut affine_points = vec![AffinePoint::IDENTITY; points.len()];
-        batch_normalize_generic(points, zs.as_mut_slice(), &mut affine_points);
+        batch_normalize_generic(points, &mut zs, &mut scratch, &mut affine_points);
         affine_points
     }
 }
 
-/// Generic implementation of batch normalization.
-fn batch_normalize_generic<P, Z, I, O>(points: &P, mut zs: Z, out: &mut O)
-where
-    FieldElement: BatchInvert<Z, Output = CtOption<I>>,
-    P: AsRef<[EdwardsPoint]> + ?Sized,
-    Z: AsMut<[FieldElement]>,
-    I: AsRef<[FieldElement]>,
-    O: AsMut<[AffinePoint]> + ?Sized,
-{
-    let points = points.as_ref();
-    let out = out.as_mut();
+fn batch_normalize_generic(
+    points: &[EdwardsPoint],
+    zs: &mut [FieldElement],
+    scratch: &mut [FieldElement],
+    out: &mut [AffinePoint],
+) {
+    debug_assert_eq!(points.len(), zs.len());
+    debug_assert_eq!(points.len(), scratch.len());
+    debug_assert_eq!(points.len(), out.len());
 
-    for (i, point) in points.iter().enumerate() {
-        // Even a single zero value will fail inversion for the entire batch.
-        // Put a dummy value (above `FieldElement::ONE`) so inversion succeeds
-        // and treat that case specially later-on.
-        zs.as_mut()[i].conditional_assign(&point.Z, !point.Z.ct_eq(&FieldElement::ZERO));
+    for (z, point) in zs.iter_mut().zip(points) {
+        *z = point.Z;
     }
 
-    // This is safe to unwrap since we assured that all elements are non-zero
-    let zs_inverses = <FieldElement as BatchInvert<Z>>::batch_invert(zs)
-        .expect("all elements should be non-zero");
+    // Zero `zs` (identity) are handled explicitly below, so the `Choice` here is informational only
+    let _ = FieldElement::batch_invert_in_place(zs, scratch);
 
     for i in 0..out.len() {
-        // If the `z` coordinate is non-zero, we can use it to invert;
-        // otherwise it defaults to the `IDENTITY` value.
         out[i] = AffinePoint::conditional_select(
             &AffinePoint {
-                x: points[i].X * zs_inverses.as_ref()[i],
-                y: points[i].Y * zs_inverses.as_ref()[i],
+                x: points[i].X * zs[i],
+                y: points[i].Y * zs[i],
             },
             &AffinePoint::IDENTITY,
             points[i].Z.ct_eq(&FieldElement::ZERO),

ed448-goldilocks/src/edwards/extended.rs

@@ -7,7 +7,6 @@ use crate::{
     AffinePoint, Decaf448, DecafPoint, Ed448, EdwardsPoint,
     curve::twedwards::extended::ExtendedPoint as TwistedExtendedPoint,
 };
-use elliptic_curve::ops::Reduce;
 use elliptic_curve::{
     Field,
     array::Array,
@@ -16,6 +15,7 @@ use elliptic_curve::{
         consts::{U28, U56, U84, U88},
         modular::ConstMontyParams,
     },
+    ops::{BatchInvert, Reduce},
     zeroize::DefaultIsZeroes,
 };
 use hash2curve::MapToCurve;
@@ -52,6 +52,8 @@ impl UpperHex for FieldElement {
     }
 }
 
+impl BatchInvert for FieldElement {}
+
 impl ConstantTimeEq for FieldElement {
     fn ct_eq(&self, other: &Self) -> Choice {
         self.0.ct_eq(&other.0)

ed448-goldilocks/src/field/element.rs

@@ -40,6 +40,7 @@
 //! Please see type-specific documentation for more information.
 
 #[cfg(feature = "alloc")]
+#[macro_use]
 extern crate alloc;
 #[cfg(feature = "std")]
 extern crate std;

ed448-goldilocks/src/lib.rs

@@ -30,7 +30,7 @@ use elliptic_curve::{
     Generate,
     bigint::{Odd, U256, modular::Retrieve},
     ff::{self, Field, PrimeField},
-    ops::Invert,
+    ops::{BatchInvert, Invert},
     rand_core::{TryCryptoRng, TryRng},
     subtle::{Choice, ConditionallySelectable, ConstantTimeEq, CtOption},
     zeroize::DefaultIsZeroes,
@@ -230,6 +230,47 @@ impl FieldElement {
     }
 }
 
+impl BatchInvert for FieldElement {
+    fn batch_invert_in_place(elements: &mut [Self], scratch: &mut [Self]) -> Choice {
+        // Implements "Montgomery's trick", a trick for computing many modular inverses at once.
+        //
+        // "Montgomery's trick" works by reducing the problem of computing `n` inverses
+        // to computing a single inversion, plus some storage and `O(n)` extra multiplications.
+        //
+        // See: https://iacr.org/archive/pkc2004/29470042/29470042.pdf section 2.2.
+        assert_eq!(elements.len(), scratch.len());
+
+        // TODO(tarcieri): share implementation with `elliptic-curve` somehow
+        // The only difference from the default impl: we need to normalize the elements first
+        for e in elements.iter_mut() {
+            *e = e.normalize();
+        }
+
+        let mut acc = Self::ONE;
+        let mut all_nonzero = Choice::from(1u8);
+
+        for (tmp, e) in scratch.iter_mut().zip(elements.iter()) {
+            // $ a_n = a_{n-1}*x_n $
+            *tmp = acc;
+            let is_zero = e.ct_eq(&Self::ZERO);
+            all_nonzero &= !is_zero;
+            acc = Self::conditional_select(&(acc * e), &acc, is_zero);
+        }
+
+        // `acc` is the product of every nonzero element, so this can't fail.
+        acc = acc.invert().unwrap_or(Self::ONE);
+
+        for (e, tmp) in elements.iter_mut().zip(scratch.iter()).rev() {
+            let is_zero = e.ct_eq(&Self::ZERO);
+            let new_acc = Self::conditional_select(&(acc * *e), &acc, is_zero);
+            *e = Self::conditional_select(&(acc * *tmp), e, is_zero);
+            acc = new_acc;
+        }
+
+        all_nonzero
+    }
+}
+
 impl Invert for FieldElement {
     type Output = CtOption<Self>;
 
@@ -698,26 +739,14 @@ mod tests {
 
     #[test]
     #[cfg(feature = "getrandom")]
-    fn batch_invert_array() {
-        let k: FieldElement = FieldElement::generate();
-        let l: FieldElement = FieldElement::generate();
-
-        let expected = [k.invert().unwrap(), l.invert().unwrap()];
-        let actual = <FieldElement as BatchInvert<_>>::batch_invert([k, l]).unwrap();
-
-        assert_eq!(expected[0], actual[0].normalize());
-        assert_eq!(expected[1], actual[1].normalize());
-    }
-
-    #[test]
-    #[cfg(all(feature = "alloc", feature = "getrandom"))]
     fn batch_invert() {
         let k: FieldElement = FieldElement::generate();
         let l: FieldElement = FieldElement::generate();
 
-        let expected = vec![k.invert().unwrap(), l.invert().unwrap()];
-        let field_elements = vec![k, l]; // to test impl of `BatchInvert` for `Vec`
-        let actual = <FieldElement as BatchInvert<_>>::batch_invert(field_elements).unwrap();
+        let expected = [k.invert().unwrap(), l.invert().unwrap()];
+        let mut actual = [k, l];
+        let mut scratch = [FieldElement::ZERO; 2];
+        FieldElement::batch_invert_in_place(&mut actual, &mut scratch);
 
         assert_eq!(expected[0], actual[0].normalize());
         assert_eq!(expected[1], actual[1].normalize());

k256/src/arithmetic/field.rs

@@ -257,9 +257,10 @@ impl<const N: usize> BatchNormalize<[ProjectivePoint; N]> for ProjectivePoint {
 
     #[inline]
     fn batch_normalize(points: &[Self; N]) -> [<Self as CurveGroup>::Affine; N] {
-        let zs = [FieldElement::ONE; N];
+        let mut zs = [FieldElement::ZERO; N];
+        let mut scratch = [FieldElement::ZERO; N];
         let mut affine_points = [AffinePoint::IDENTITY; N];
-        batch_normalize_generic(points, zs, &mut affine_points);
+        batch_normalize(points, &mut zs, &mut scratch, &mut affine_points);
         affine_points
     }
 }
@@ -270,40 +271,34 @@ impl BatchNormalize<[ProjectivePoint]> for ProjectivePoint {
 
     #[inline]
     fn batch_normalize(points: &[Self]) -> Vec<<Self as CurveGroup>::Affine> {
-        let zs = vec![FieldElement::ONE; points.len()];
+        let mut zs = vec![FieldElement::ZERO; points.len()];
+        let mut scratch = vec![FieldElement::ZERO; points.len()];
         let mut affine_points = vec![AffinePoint::IDENTITY; points.len()];
-        batch_normalize_generic(points, zs, &mut affine_points);
+        batch_normalize(points, &mut zs, &mut scratch, &mut affine_points);
         affine_points
     }
 }
 
-fn batch_normalize_generic<P, Z, I, O>(points: &P, mut zs: Z, out: &mut O)
-where
-    FieldElement: BatchInvert<Z, Output = CtOption<I>>,
-    P: AsRef<[ProjectivePoint]> + ?Sized,
-    Z: AsMut<[FieldElement]>,
-    I: AsRef<[FieldElement]>,
-    O: AsMut<[AffinePoint]> + ?Sized,
-{
-    let points = points.as_ref();
-    let out = out.as_mut();
+fn batch_normalize(
+    points: &[ProjectivePoint],
+    zs: &mut [FieldElement],
+    scratch: &mut [FieldElement],
+    out: &mut [AffinePoint],
+) {
+    debug_assert_eq!(points.len(), zs.len());
+    debug_assert_eq!(points.len(), scratch.len());
+    debug_assert_eq!(points.len(), out.len());
 
-    for i in 0..points.len() {
-        // Even a single zero value will fail inversion for the entire batch.
-        // Put a dummy value (above `FieldElement::ONE`) so inversion succeeds
-        // and treat that case specially later-on.
-        zs.as_mut()[i].conditional_assign(&points[i].z, !points[i].z.normalizes_to_zero());
+    for (z, point) in zs.iter_mut().zip(points) {
+        *z = point.z;
     }
 
-    // This is safe to unwrap since we assured that all elements are non-zero
-    let zs_inverses = <FieldElement as BatchInvert<Z>>::batch_invert(zs)
-        .expect("all elements should be non-zero");
+    // Zero `zs` (identity) are handled explicitly below, so the `Choice` here is informational only
+    let _ = FieldElement::batch_invert_in_place(zs, scratch);
 
     for i in 0..out.len() {
-        // If the `z` coordinate is non-zero, we can use it to invert;
-        // otherwise it defaults to the `IDENTITY` value.
         out[i] = AffinePoint::conditional_select(
-            &points[i].to_affine_internal(zs_inverses.as_ref()[i]),
+            &points[i].to_affine_internal(zs[i]),
             &AffinePoint::IDENTITY,
             points[i].z.normalizes_to_zero(),
         );
@@ -455,8 +450,9 @@ impl CurveGroup for ProjectivePoint {
     #[inline]
     fn batch_normalize(projective: &[Self], affine: &mut [Self::Affine]) {
         assert_eq!(projective.len(), affine.len());
-        let mut zs = vec![FieldElement::ONE; projective.len()];
-        batch_normalize_generic(projective, zs.as_mut_slice(), affine);
+        let mut zs = vec![FieldElement::ZERO; projective.len()];
+        let mut scratch = vec![FieldElement::ZERO; projective.len()];
+        batch_normalize(projective, &mut zs, &mut scratch, affine);
     }
 }