From 02af12b81a46ff165db14ff6da1ea0cc6dbfe0ba Mon Sep 17 00:00:00 2001
From: Sean Bowe <ewillbefull@gmail.com>
Date: Sat, 30 Jun 2018 16:50:07 -0600
Subject: [PATCH] Replace batch inversion for FieldElement with sequential
 variant of Montgomery's trick.

---
 src/field.rs | 73 +++++++++++++++++++---------------------------------
 1 file changed, 27 insertions(+), 46 deletions(-)

diff --git a/src/field.rs b/src/field.rs
index 938fa03..c9f4c49 100644
--- a/src/field.rs
+++ b/src/field.rs
@@ -142,58 +142,34 @@ impl FieldElement {
     /// Given a slice of public `FieldElements`, replace each with its inverse.
     ///
     /// All input `FieldElements` **MUST** be nonzero.
-    ///
-    /// This function is most efficient when the batch size (slice
-    /// length) is a power of 2.
     #[cfg(any(feature = "alloc", feature = "std"))]
     pub fn batch_invert(inputs: &mut [FieldElement]) {
-        // First, compute the product of all inputs using a product
-        // tree:
-        //
-        // Inputs: [x_0, x_1, x_2]
-        //
-        // Tree:
-        //
-        //                 x_0*x_1*x_2*1         tree[1]
-        //                   /       \
-        //               x_0*x_1     x_2*1       tree[2,3]
-        //                / \        / \
-        //              x_0  x_1   x_2  1        tree[4,5,6,7]
-        //
-        //  The leaves of the tree are the inputs.  We store the tree in
-        //  an array of length 2*n, similar to a binary heap.
-        //
-        //  To initialize the tree, set every node to 1, then fill in
-        //  the leaf nodes with the input variables.  Finally, set every
-        //  non-leaf node to be the product of its children.
+        // Montgomery’s Trick and Fast Implementation of Masked AES
+        // Genelle, Prouff and Quisquater
+        // Section 3.2
 
-        let n = inputs.len().next_power_of_two();
-        let mut tree = vec![FieldElement::one(); 2*n];
-        tree[n..n+inputs.len()].copy_from_slice(inputs);
-        for i in (1..n).rev() {
-            tree[i] = &tree[2*i] * &tree[2*i+1];
+        let n = inputs.len();
+        let mut scratch = vec![FieldElement::one(); n];
+
+        // Keep an accumulator of all of the previous products
+        let mut acc = FieldElement::one();
+
+        // Pass through the input vector, recording the previous
+        // products in the scratch space
+        for (input, scratch) in inputs.iter().zip(scratch.iter_mut()) {
+            *scratch = acc;
+            acc = &acc * input;
         }
 
-        // The root of the tree is the product of all inputs, and is
-        // stored at index 1.  Compute its inverse.
-        let allinv = tree[1].invert();
+        // Compute the inverse of all products
+        acc = acc.invert();
 
-        // To compute y_i = 1/x_i, start at the i-th leaf node of the
-        // tree, and walk up to the root of the tree, multiplying
-        // `allinv` by each sibling.  This computes
-        //
-        // y_i = y * (all x_j except x_i)
-        //
-        // using lg(n) multiplications for each y_i, taking n*lg(n) in
-        // total.
-        for i in 0..inputs.len() {
-            let mut inv = allinv;
-            let mut node = n + i;
-            while node > 1 {
-                inv *= &tree[node ^ 1];
-                node = node >> 1;
-            }
-            inputs[i] = inv;
+        // Pass through the vector backwards to compute the inverses
+        // in place
+        for (input, scratch) in inputs.iter_mut().rev().zip(scratch.into_iter().rev()) {
+            let tmp = &acc * input;
+            *input = &acc * &scratch;
+            acc = tmp;
         }
     }
 
@@ -496,4 +472,9 @@ mod test {
             assert_eq!(one_bytes[i], 0);
         }
     }
+
+    #[test]
+    fn batch_invert_empty() {
+        FieldElement::batch_invert(&mut []);
+    }
 }