Merge remote-tracking branch 'ebfull/sequential-montgomery-trick' into develop

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 []);
+    }
 }

src/scalar.rs

@@ -434,9 +434,6 @@ impl Scalar {
     /// *prove* that this is the case, you **SHOULD NOT USE THIS
     /// FUNCTION**.
     ///
-    /// This function is most efficient when the batch size (slice
-    /// length) is a power of 2.
-    ///
     /// # Example
     ///
     /// ```
@@ -474,38 +471,47 @@ impl Scalar {
         // Mark UnpackedScalars as zeroable.
         unsafe impl ZeroSafe for UnpackedScalar {}
 
-        let n = inputs.len().next_power_of_two();
+        let n = inputs.len();
         let one: UnpackedScalar = Scalar::one().unpack().to_montgomery();
 
-        // Wrap the tree storage in a ClearOnDrop to wipe it when we
-        // pass out of scope.
-        let tree_vec = vec![one; 2*n];
-        let mut tree = ClearOnDrop::new(tree_vec);
+        // Wrap the scratch storage in a ClearOnDrop to wipe it when
+        // we pass out of scope.
+        let scratch_vec = vec![one; n];
+        let mut scratch = ClearOnDrop::new(scratch_vec);
 
-        for i in 0..inputs.len() {
-            tree[n+i] = inputs[i].unpack().to_montgomery();
+        // Keep an accumulator of all of the previous products
+        let mut acc = Scalar::one().unpack().to_montgomery();
+
+        // Pass through the input vector, recording the previous
+        // products in the scratch space
+        for (input, scratch) in inputs.iter_mut().zip(scratch.iter_mut()) {
+            *scratch = acc;
+
+            // Avoid unnecessary Montgomery multiplication in second pass by
+            // keeping inputs in Montgomery form
+            let tmp = input.unpack().to_montgomery();
+            *input = tmp.pack();
+            acc = UnpackedScalar::montgomery_mul(&acc, &tmp);
         }
 
-        for i in (1..n).rev() {
-            tree[i] = UnpackedScalar::montgomery_mul(&tree[2*i], &tree[2*i+1]);
+        // acc is nonzero iff all inputs are nonzero
+        debug_assert!(acc.pack() != Scalar::zero());
+
+        // Compute the inverse of all products
+        acc = acc.montgomery_invert().from_montgomery();
+
+        // We need to return the product of all inverses later
+        let ret = acc.pack();
+
+        // 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 = UnpackedScalar::montgomery_mul(&acc, &input.unpack());
+            *input = UnpackedScalar::montgomery_mul(&acc, &scratch).pack();
+            acc = tmp;
         }
 
-        // tree[1] is zero iff any of the inputs are zero.
-        debug_assert!(tree[1].from_montgomery().pack() != Scalar::zero());
-
-        let allinv = tree[1].montgomery_invert();
-
-        for i in 0..inputs.len() {
-            let mut inv = allinv;
-            let mut node = n + i;
-            while node > 1 {
-                inv = UnpackedScalar::montgomery_mul(&inv, &tree[node ^1]);
-                node = node >> 1;
-            }
-            inputs[i] = inv.from_montgomery().pack();
-        }
-
-        allinv.from_montgomery().pack()
+        ret
     }
 
     /// Get the bits of the scalar.
@@ -1130,6 +1136,7 @@ mod test {
         assert_eq!(parsed, X);
     }
 
+    #[cfg(debug_assertions)]
     #[test]
     #[should_panic]
     fn batch_invert_with_a_zero_input_panics() {
@@ -1138,4 +1145,25 @@ mod test {
         // This should panic in debug mode.
         Scalar::batch_invert(&mut xs);
     }
+
+    #[test]
+    fn batch_invert_empty() {
+        assert_eq!(Scalar::one(), Scalar::batch_invert(&mut []));
+    }
+
+    #[test]
+    fn batch_invert_consistency() {
+        let mut x = Scalar::from_u64(1);
+        let mut v1: Vec<_> = (0..16).map(|_| {let tmp = x; x = x + x; tmp}).collect();
+        let v2 = v1.clone();
+
+        let expected: Scalar = v1.iter().product();
+        let expected = expected.invert();
+        let ret = Scalar::batch_invert(&mut v1);
+        assert_eq!(ret, expected);
+
+        for (a, b) in v1.iter().zip(v2.iter()) {
+            assert_eq!(a * b, Scalar::one());
+        }
+    }
 }