elliptic-curve: impl BatchInvert for NonZeroScalar (#1890)

This implements `BatchInvert` for `NonZeroScalar`.

To accomplish this, I did the following notable things:
- Remove all trait bounds on `BatchInvert` itself.
- Remove `CtOption` from `BatchInvert::batch_invert`s return type.
- Expose `invert_batch_internal()` internally, which now takes an
`invert` function instead of requiring `trait Invert`.
- Implement `MulAssign` for `NonZeroScalar`.

Things that could still be added:
- Mirror all `BatchInvert` implementations on `Scalar`, I left the ones
taking a reference to a slice out.
- Implement `MulAssign` for `NonZeroScalar` on more combinations.

It might be a tiny bit simpler if I implemented `Default` on
`NonZeroScalar`, but that seemed like a footgun to me.

Author: daxpedda

elliptic-curve/src/ops.rs

@@ -1,7 +1,7 @@
 //! Traits for arithmetic operations on elliptic curve field elements.
 
 use core::iter;
-pub use core::ops::{Add, AddAssign, Mul, Neg, Shr, ShrAssign, Sub, SubAssign};
+pub use core::ops::{Add, AddAssign, Mul, MulAssign, Neg, Shr, ShrAssign, Sub, SubAssign};
 pub use crypto_bigint::Invert;
 
 use crypto_bigint::Integer;
@@ -13,26 +13,24 @@ use alloc::{borrow::ToOwned, vec::Vec};
 
 /// Perform a batched inversion on a sequence of field elements (i.e. base field elements or scalars)
 /// at an amortized cost that should be practically as efficient as a single inversion.
-pub trait BatchInvert<FieldElements: ?Sized>: Field + Sized {
+pub trait BatchInvert<FieldElements: ?Sized> {
     /// The output of batch inversion. A container of field elements.
-    type Output: AsRef<[Self]>;
+    type Output;
 
     /// Invert a batch of field elements.
-    fn batch_invert(
-        field_elements: FieldElements,
-    ) -> CtOption<<Self as BatchInvert<FieldElements>>::Output>;
+    fn batch_invert(field_elements: FieldElements) -> <Self as BatchInvert<FieldElements>>::Output;
 }
 
 impl<const N: usize, T> BatchInvert<[T; N]> for T
 where
     T: Field,
 {
-    type Output = [Self; N];
+    type Output = CtOption<[Self; N]>;
 
     fn batch_invert(mut field_elements: [Self; N]) -> CtOption<[Self; N]> {
         let mut field_elements_pad = [Self::default(); N];
         let inversion_succeeded =
-            invert_batch_internal(&mut field_elements, &mut field_elements_pad);
+            invert_batch_internal(&mut field_elements, &mut field_elements_pad, invert);
 
         CtOption::new(field_elements, inversion_succeeded)
     }
@@ -43,11 +41,12 @@ impl<'this, T> BatchInvert<&'this mut [Self]> for T
 where
     T: Field,
 {
-    type Output = &'this mut [Self];
+    type Output = CtOption<&'this mut [Self]>;
 
     fn batch_invert(field_elements: &'this mut [Self]) -> CtOption<&'this mut [Self]> {
         let mut field_elements_pad: Vec<Self> = vec![Self::default(); field_elements.len()];
-        let inversion_succeeded = invert_batch_internal(field_elements, &mut field_elements_pad);
+        let inversion_succeeded =
+            invert_batch_internal(field_elements, &mut field_elements_pad, invert);
 
         CtOption::new(field_elements, inversion_succeeded)
     }
@@ -58,13 +57,13 @@ impl<T> BatchInvert<&[Self]> for T
 where
     T: Field,
 {
-    type Output = Vec<Self>;
+    type Output = CtOption<Vec<Self>>;
 
     fn batch_invert(field_elements: &[Self]) -> CtOption<Vec<Self>> {
         let mut field_elements: Vec<Self> = field_elements.to_owned();
         let mut field_elements_pad: Vec<Self> = vec![Self::default(); field_elements.len()];
         let inversion_succeeded =
-            invert_batch_internal(&mut field_elements, &mut field_elements_pad);
+            invert_batch_internal(&mut field_elements, &mut field_elements_pad, invert);
 
         CtOption::new(field_elements, inversion_succeeded)
     }
@@ -75,26 +74,35 @@ impl<T> BatchInvert<Vec<Self>> for T
 where
     T: Field,
 {
-    type Output = Vec<Self>;
+    type Output = CtOption<Vec<Self>>;
 
     fn batch_invert(mut field_elements: Vec<Self>) -> CtOption<Vec<Self>> {
         let mut field_elements_pad: Vec<Self> = vec![Self::default(); field_elements.len()];
         let inversion_succeeded =
-            invert_batch_internal(&mut field_elements, &mut field_elements_pad);
+            invert_batch_internal(&mut field_elements, &mut field_elements_pad, invert);
 
         CtOption::new(field_elements, inversion_succeeded)
     }
 }
 
+fn invert<T: Field>(scalar: T) -> (T, Choice) {
+    let scalar = scalar.invert();
+    let choice = scalar.is_some();
+    let scalar = scalar.unwrap_or(T::default());
+
+    (scalar, 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.
-fn invert_batch_internal<T: Field>(
+pub(crate) fn invert_batch_internal<T: Copy + Mul<Output = T> + MulAssign>(
     field_elements: &mut [T],
     field_elements_pad: &mut [T],
+    invert: fn(T) -> (T, Choice),
 ) -> Choice {
     let batch_size = field_elements.len();
     if batch_size != field_elements_pad.len() {
@@ -117,32 +125,32 @@ fn invert_batch_internal<T: Field>(
         *field_element_pad = acc;
     }
 
-    acc.invert()
-        .map(|mut acc| {
-            // Shift the iterator by one element back. The one we are skipping is served in `acc`.
-            let field_elements_pad = field_elements_pad
-                .iter()
-                .rev()
-                .skip(1)
-                .map(Some)
-                .chain(iter::once(None));
-
-            for (field_element, field_element_pad) in
-                field_elements.iter_mut().rev().zip(field_elements_pad)
-            {
-                if let Some(field_element_pad) = field_element_pad {
-                    // Store in a temporary so we can overwrite `field_element`.
-                    // $ a_{n-1} = {a_n}^{-1}*x_n $
-                    let tmp = acc * *field_element;
-                    // $ {x_n}^{-1} = a_{n}^{-1}*a_{n-1} $
-                    *field_element = acc * *field_element_pad;
-                    acc = tmp;
-                } else {
-                    *field_element = acc;
-                }
-            }
-        })
-        .is_some()
+    let (mut acc, choice) = invert(acc);
+
+    // Shift the iterator by one element back. The one we are skipping is served in `acc`.
+    let field_elements_pad = field_elements_pad
+        .iter()
+        .rev()
+        .skip(1)
+        .map(Some)
+        .chain(iter::once(None));
+
+    for (field_element, field_element_pad) in
+        field_elements.iter_mut().rev().zip(field_elements_pad)
+    {
+        if let Some(field_element_pad) = field_element_pad {
+            // Store in a temporary so we can overwrite `field_element`.
+            // $ a_{n-1} = {a_n}^{-1}*x_n $
+            let tmp = acc * *field_element;
+            // $ {x_n}^{-1} = a_{n}^{-1}*a_{n-1} $
+            *field_element = acc * *field_element_pad;
+            acc = tmp;
+        } else {
+            *field_element = acc;
+        }
+    }
+
+    choice
 }
 
 /// Linear combination.

elliptic-curve/src/scalar/nonzero.rs

@@ -2,14 +2,14 @@
 
 use crate::{
     CurveArithmetic, Error, FieldBytes, PrimeCurve, Scalar, ScalarPrimitive, SecretKey,
-    ops::{Invert, Reduce, ReduceNonZero},
+    ops::{self, BatchInvert, Invert, Reduce, ReduceNonZero},
     point::NonIdentity,
     scalar::IsHigh,
 };
 use base16ct::HexDisplay;
 use core::{
     fmt,
-    ops::{Deref, Mul, Neg},
+    ops::{Deref, Mul, MulAssign, Neg},
     str,
 };
 use crypto_bigint::{ArrayEncoding, Integer};
@@ -18,6 +18,9 @@ use rand_core::{CryptoRng, TryCryptoRng};
 use subtle::{Choice, ConditionallySelectable, ConstantTimeEq, CtOption};
 use zeroize::Zeroize;
 
+#[cfg(feature = "alloc")]
+use alloc::vec::Vec;
+
 #[cfg(feature = "serde")]
 use serdect::serde::{Deserialize, Serialize, de, ser};
 
@@ -96,6 +99,47 @@ where
     }
 }
 
+impl<const N: usize, C> BatchInvert<[Self; N]> for NonZeroScalar<C>
+where
+    C: CurveArithmetic + PrimeCurve,
+{
+    type Output = [Self; N];
+
+    fn batch_invert(mut field_elements: [Self; N]) -> [Self; N] {
+        let mut field_elements_pad = [Self {
+            scalar: Scalar::<C>::ONE,
+        }; N];
+        ops::invert_batch_internal(&mut field_elements, &mut field_elements_pad, |scalar| {
+            (scalar.invert(), Choice::from(1))
+        });
+
+        field_elements
+    }
+}
+
+#[cfg(feature = "alloc")]
+impl<C> BatchInvert<Vec<Self>> for NonZeroScalar<C>
+where
+    C: CurveArithmetic + PrimeCurve,
+{
+    type Output = Vec<Self>;
+
+    fn batch_invert(mut field_elements: Vec<Self>) -> Vec<Self> {
+        let mut field_elements_pad: Vec<Self> = vec![
+            Self {
+                scalar: Scalar::<C>::ONE,
+            };
+            field_elements.len()
+        ];
+
+        ops::invert_batch_internal(&mut field_elements, &mut field_elements_pad, |scalar| {
+            (scalar.invert(), Choice::from(1))
+        });
+
+        field_elements
+    }
+}
+
 impl<C> ConditionallySelectable for NonZeroScalar<C>
 where
     C: CurveArithmetic,
@@ -307,6 +351,15 @@ where
     }
 }
 
+impl<C> MulAssign for NonZeroScalar<C>
+where
+    C: PrimeCurve + CurveArithmetic,
+{
+    fn mul_assign(&mut self, rhs: Self) {
+        *self = *self * rhs;
+    }
+}
+
 impl<C> PartialEq for NonZeroScalar<C>
 where
     C: CurveArithmetic,