Improve BatchInvert (#1864)

This PR does a couple of things:
- Changes the implementation to require two write stores instead of 3
write stores and one read.
- Add a `BatchInvert::batch_invert_mut()` methods that allows users to
pass one of the two stores themselves. With this a single store can be
allocated instead of 3.
- Reduce the bounds from `Invert + Mul + Copy + Default +
ConditionallySelectable` to `Field`.

This was largely inspired by [`curve25519-dalek`s
implementation](https://github.com/dalek-cryptography/curve25519-dalek/blob/dd5bd108d6985491acb3de25497fb082a29c9fb7/curve25519-dalek/src/scalar.rs#L788-L837).
I adjusted it a bit to save two unnecessary multiplications by one.

Companion PR: RustCrypto/elliptic-curves#1205.

Author: daxpedda

elliptic-curve/src/ops.rs

@@ -1,78 +1,77 @@
 //! 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 crypto_bigint::Invert;
 
 use crypto_bigint::Integer;
-use subtle::{Choice, ConditionallySelectable, CtOption};
+use ff::Field;
+use subtle::{Choice, CtOption};
 
 #[cfg(feature = "alloc")]
-use alloc::vec::Vec;
+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>: Invert + Sized {
+pub trait BatchInvert<FieldElements: ?Sized>: Field + Sized {
     /// The output of batch inversion. A container of field elements.
     type Output: AsRef<[Self]>;
 
     /// Invert a batch of field elements.
     fn batch_invert(
         field_elements: &FieldElements,
     ) -> CtOption<<Self as BatchInvert<FieldElements>>::Output>;
+
+    /// Invert a batch of field elements in-place.
+    ///
+    /// # ⚠️ Warning
+    ///
+    /// Even though `field_elements` is modified regardless of success, on failure it does not
+    /// contain correctly inverted scalars and should be discarded instead.
+    ///
+    /// Consider using [`Self::batch_invert()`] instead.
+    fn batch_invert_mut(field_elements: &mut FieldElements) -> Choice;
 }
 
 impl<const N: usize, T> BatchInvert<[T; N]> for T
 where
-    T: Invert<Output = CtOption<Self>>
-        + Mul<Self, Output = Self>
-        + Copy
-        + Default
-        + ConditionallySelectable,
+    T: Field,
 {
     type Output = [Self; N];
 
     fn batch_invert(field_elements: &[Self; N]) -> CtOption<[Self; N]> {
-        let mut field_elements_multiples = [Self::default(); N];
-        let mut field_elements_multiples_inverses = [Self::default(); N];
-        let mut field_elements_inverses = [Self::default(); N];
-
-        let inversion_succeeded = invert_batch_internal(
-            field_elements,
-            &mut field_elements_multiples,
-            &mut field_elements_multiples_inverses,
-            &mut field_elements_inverses,
-        );
+        let mut field_elements_inverses = *field_elements;
+        let inversion_succeeded = Self::batch_invert_mut(&mut field_elements_inverses);
 
         CtOption::new(field_elements_inverses, inversion_succeeded)
     }
+
+    fn batch_invert_mut(field_elements: &mut [T; N]) -> Choice {
+        let mut field_elements_pad = [Self::default(); N];
+
+        invert_batch_internal(field_elements, &mut field_elements_pad)
+    }
 }
 
 #[cfg(feature = "alloc")]
 impl<T> BatchInvert<[T]> for T
 where
-    T: Invert<Output = CtOption<Self>>
-        + Mul<Self, Output = Self>
-        + Copy
-        + Default
-        + ConditionallySelectable,
+    T: Field,
 {
     type Output = Vec<Self>;
 
     fn batch_invert(field_elements: &[Self]) -> CtOption<Vec<Self>> {
-        let mut field_elements_multiples: Vec<Self> = vec![Self::default(); field_elements.len()];
-        let mut field_elements_multiples_inverses: Vec<Self> =
-            vec![Self::default(); field_elements.len()];
-        let mut field_elements_inverses: Vec<Self> = vec![Self::default(); field_elements.len()];
-
-        let inversion_succeeded = invert_batch_internal(
-            field_elements,
-            field_elements_multiples.as_mut(),
-            field_elements_multiples_inverses.as_mut(),
-            field_elements_inverses.as_mut(),
-        );
+        let mut field_elements_inverses: Vec<Self> = field_elements.to_owned();
+        let inversion_succeeded = Self::batch_invert_mut(field_elements_inverses.as_mut_slice());
 
         CtOption::new(field_elements_inverses, inversion_succeeded)
     }
+
+    fn batch_invert_mut(field_elements: &mut [T]) -> Choice {
+        let mut field_elements_pad: Vec<Self> = vec![Self::default(); field_elements.len()];
+
+        invert_batch_internal(field_elements, field_elements_pad.as_mut())
+    }
 }
 
 /// Implements "Montgomery's trick", a trick for computing many modular inverses at once.
@@ -81,44 +80,51 @@ where
 /// 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: Invert<Output = CtOption<T>> + Mul<T, Output = T> + Default + ConditionallySelectable,
->(
-    field_elements: &[T],
-    field_elements_multiples: &mut [T],
-    field_elements_multiples_inverses: &mut [T],
-    field_elements_inverses: &mut [T],
+fn invert_batch_internal<T: Field>(
+    field_elements: &mut [T],
+    field_elements_pad: &mut [T],
 ) -> Choice {
     let batch_size = field_elements.len();
-    if batch_size == 0
-        || batch_size != field_elements_multiples.len()
-        || batch_size != field_elements_multiples_inverses.len()
-    {
+    if batch_size == 0 || batch_size != field_elements_pad.len() {
         return Choice::from(0);
     }
 
-    field_elements_multiples[0] = field_elements[0];
-    for i in 1..batch_size {
+    let mut acc = field_elements[0];
+    field_elements_pad[0] = acc;
+
+    for (field_element, field_element_pad) in field_elements
+        .iter_mut()
+        .zip(field_elements_pad.iter_mut())
+        .skip(1)
+    {
         // $ a_n = a_{n-1}*x_n $
-        field_elements_multiples[i] = field_elements_multiples[i - 1] * field_elements[i];
+        acc *= *field_element;
+        *field_element_pad = acc;
     }
 
-    field_elements_multiples[batch_size - 1]
-        .invert()
-        .map(|multiple_of_inverses_of_all_field_elements| {
-            field_elements_multiples_inverses[batch_size - 1] =
-                multiple_of_inverses_of_all_field_elements;
-            for i in (1..batch_size).rev() {
-                // $ a_{n-1} = {a_n}^{-1}*x_n $
-                field_elements_multiples_inverses[i - 1] =
-                    field_elements_multiples_inverses[i] * field_elements[i];
-            }
-
-            field_elements_inverses[0] = field_elements_multiples_inverses[0];
-            for i in 1..batch_size {
-                // $ {x_n}^{-1} = a_{n}^{-1}*a_{n-1} $
-                field_elements_inverses[i] =
-                    field_elements_multiples_inverses[i] * field_elements_multiples[i - 1];
+    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()