BoxedResidue: avoid allocations in pow loop (#394)

Performs Montgomery multiplications and squarings in-place, avoiding
allocations, by constructing a reusable `MontgomeryMultiplier`.

Montgomery arithmetic/modpow, BoxedUint^BoxedUint
                        time:   [24.265 µs 24.274 µs 24.288 µs]
                        change: [-24.321% -24.194% -24.081%] (p = 0.00 < 0.05)
                        Performance has improved.

Author: tarcieri

src/modular/boxed_residue/mul.rs

@@ -1,34 +1,36 @@
 //! Multiplications between boxed residues.
 
-use super::{montgomery_reduction_boxed_mut, BoxedResidue};
-use crate::traits::Square;
-use crate::{BoxedUint, Limb};
-use core::ops::{Mul, MulAssign};
+use super::{montgomery_reduction_boxed_mut, BoxedResidue, BoxedResidueParams};
+use crate::{traits::Square, uint::mul::mul_limbs, BoxedUint, Limb};
+use core::{
+    borrow::Borrow,
+    ops::{Mul, MulAssign},
+};
+
+#[cfg(feature = "zeroize")]
+use zeroize::Zeroize;
 
 impl BoxedResidue {
     /// Multiplies by `rhs`.
     pub fn mul(&self, rhs: &Self) -> Self {
         debug_assert_eq!(&self.residue_params, &rhs.residue_params);
 
+        let montgomery_form = MontgomeryMultiplier::from(self.residue_params.borrow())
+            .mul(&self.montgomery_form, &rhs.montgomery_form);
+
         Self {
-            montgomery_form: mul_montgomery_form(
-                &self.montgomery_form,
-                &rhs.montgomery_form,
-                &self.residue_params.modulus,
-                self.residue_params.mod_neg_inv,
-            ),
+            montgomery_form,
             residue_params: self.residue_params.clone(),
         }
     }
 
     /// Computes the (reduced) square of a residue.
     pub fn square(&self) -> Self {
+        let montgomery_form =
+            MontgomeryMultiplier::from(self.residue_params.borrow()).square(&self.montgomery_form);
+
         Self {
-            montgomery_form: square_montgomery_form(
-                &self.montgomery_form,
-                &self.residue_params.modulus,
-                self.residue_params.mod_neg_inv,
-            ),
+            montgomery_form,
             residue_params: self.residue_params.clone(),
         }
     }
@@ -67,13 +69,8 @@ impl Mul<BoxedResidue> for BoxedResidue {
 impl MulAssign<&BoxedResidue> for BoxedResidue {
     fn mul_assign(&mut self, rhs: &BoxedResidue) {
         debug_assert_eq!(&self.residue_params, &rhs.residue_params);
-
-        self.montgomery_form = mul_montgomery_form(
-            &self.montgomery_form,
-            &rhs.montgomery_form,
-            &self.residue_params.modulus,
-            self.residue_params.mod_neg_inv,
-        );
+        MontgomeryMultiplier::from(self.residue_params.borrow())
+            .mul_assign(&mut self.montgomery_form, &rhs.montgomery_form);
     }
 }
 
@@ -89,55 +86,80 @@ impl Square for BoxedResidue {
     }
 }
 
-pub(super) fn mul_montgomery_form(
-    a: &BoxedUint,
-    b: &BoxedUint,
-    modulus: &BoxedUint,
-    mod_neg_inv: Limb,
-) -> BoxedUint {
-    let mut ret = a.clone();
-    mul_montgomery_form_assign(&mut ret, b, modulus, mod_neg_inv);
-    ret
+impl<'a> From<&'a BoxedResidueParams> for MontgomeryMultiplier<'a> {
+    fn from(residue_params: &'a BoxedResidueParams) -> MontgomeryMultiplier<'a> {
+        MontgomeryMultiplier::new(&residue_params.modulus, residue_params.mod_neg_inv)
+    }
 }
 
-pub(super) fn mul_montgomery_form_assign(
-    a: &mut BoxedUint,
-    b: &BoxedUint,
-    modulus: &BoxedUint,
+/// Montgomery multiplier with a pre-allocated internal buffer to avoid additional allocations.
+pub(super) struct MontgomeryMultiplier<'a> {
+    product: BoxedUint,
+    modulus: &'a BoxedUint,
     mod_neg_inv: Limb,
-) {
-    debug_assert_eq!(a.bits_precision(), modulus.bits_precision());
-    debug_assert_eq!(b.bits_precision(), modulus.bits_precision());
+}
 
-    let mut product = a.mul(b);
-    montgomery_reduction_boxed_mut(&mut product, modulus, mod_neg_inv, a);
+impl<'a> MontgomeryMultiplier<'a> {
+    /// Create a new Montgomery multiplier.
+    pub(super) fn new(modulus: &'a BoxedUint, mod_neg_inv: Limb) -> Self {
+        Self {
+            product: BoxedUint::zero_with_precision(modulus.bits_precision() * 2),
+            modulus,
+            mod_neg_inv,
+        }
+    }
 
-    #[cfg(feature = "zeroize")]
-    zeroize::Zeroize::zeroize(&mut product);
-}
+    /// Multiply two values in Montgomery form.
+    pub(super) fn mul(&mut self, a: &BoxedUint, b: &BoxedUint) -> BoxedUint {
+        let mut ret = a.clone();
+        self.mul_assign(&mut ret, b);
+        ret
+    }
 
-#[inline]
-pub(super) fn square_montgomery_form(
-    a: &BoxedUint,
-    modulus: &BoxedUint,
-    mod_neg_inv: Limb,
-) -> BoxedUint {
-    let mut ret = a.clone();
-    square_montgomery_form_assign(&mut ret, modulus, mod_neg_inv);
-    ret
-}
+    /// Multiply two values in Montgomery form, assigning the product to `a`.
+    pub(super) fn mul_assign(&mut self, a: &mut BoxedUint, b: &BoxedUint) {
+        debug_assert_eq!(a.bits_precision(), self.modulus.bits_precision());
+        debug_assert_eq!(b.bits_precision(), self.modulus.bits_precision());
 
-#[inline]
-pub(super) fn square_montgomery_form_assign(
-    a: &mut BoxedUint,
-    modulus: &BoxedUint,
-    mod_neg_inv: Limb,
-) {
-    debug_assert_eq!(a.bits_precision(), modulus.bits_precision());
+        self.clear_product();
+        mul_limbs(&a.limbs, &b.limbs, &mut self.product.limbs);
+        self.montgomery_reduction(a);
+    }
+
+    /// Square the given value in Montgomery form.
+    #[inline]
+    pub(super) fn square(&mut self, a: &BoxedUint) -> BoxedUint {
+        let mut ret = a.clone();
+        self.square_assign(&mut ret);
+        ret
+    }
+
+    /// Square the given value in Montgomery form, assigning the result to `a`.
+    #[inline]
+    pub(super) fn square_assign(&mut self, a: &mut BoxedUint) {
+        debug_assert_eq!(a.bits_precision(), self.modulus.bits_precision());
+        self.clear_product();
+        mul_limbs(&a.limbs, &a.limbs, &mut self.product.limbs);
+        self.montgomery_reduction(a);
+    }
+
+    /// Clear the internal product buffer.
+    fn clear_product(&mut self) {
+        self.product
+            .limbs
+            .iter_mut()
+            .for_each(|limb| *limb = Limb::ZERO);
+    }
 
-    let mut square = a.square();
-    montgomery_reduction_boxed_mut(&mut square, modulus, mod_neg_inv, a);
+    /// Perform Montgomery reduction on the internal product buffer.
+    fn montgomery_reduction(&mut self, out: &mut BoxedUint) {
+        montgomery_reduction_boxed_mut(&mut self.product, self.modulus, self.mod_neg_inv, out);
+    }
+}
 
-    #[cfg(feature = "zeroize")]
-    zeroize::Zeroize::zeroize(&mut square);
+#[cfg(feature = "zeroize")]
+impl Drop for MontgomeryMultiplier<'_> {
+    fn drop(&mut self) {
+        self.product.zeroize();
+    }
 }

src/modular/boxed_residue/pow.rs

@@ -1,10 +1,8 @@
 //! Modular exponentiation support for [`BoxedResidue`].
 
-use super::{
-    mul::{mul_montgomery_form, mul_montgomery_form_assign, square_montgomery_form_assign},
-    BoxedResidue,
-};
+use super::{mul::MontgomeryMultiplier, BoxedResidue};
 use crate::{BoxedUint, Limb, PowBoundedExp, Word};
+use alloc::vec::Vec;
 use subtle::ConstantTimeEq;
 
 impl BoxedResidue {
@@ -58,12 +56,15 @@ fn pow_montgomery_form(
     const WINDOW: u32 = 4;
     const WINDOW_MASK: Word = (1 << WINDOW) - 1;
 
+    let mut multiplier = MontgomeryMultiplier::new(modulus, mod_neg_inv);
+
     // powers[i] contains x^i
-    let mut powers = vec![r.clone(); 1 << WINDOW];
-    powers[1] = x.clone();
+    let mut powers = Vec::with_capacity(1 << WINDOW);
+    powers.push(r.clone()); // 1 in Montgomery form
+    powers.push(x.clone());
 
-    for i in 2..powers.len() {
-        powers[i] = mul_montgomery_form(&powers[i - 1], x, modulus, mod_neg_inv);
+    for i in 2..(1 << WINDOW) {
+        powers.push(multiplier.mul(&powers[i - 1], x));
     }
 
     let starting_limb = ((exponent_bits - 1) / Limb::BITS) as usize;
@@ -72,6 +73,7 @@ fn pow_montgomery_form(
     let starting_window_mask = (1 << (starting_bit_in_limb % WINDOW + 1)) - 1;
 
     let mut z = r.clone(); // 1 in Montgomery form
+    let mut power = powers[0].clone();
 
     for limb_num in (0..=starting_limb).rev() {
         let w = exponent.as_limbs()[limb_num].0;
@@ -91,17 +93,17 @@ fn pow_montgomery_form(
                 idx &= starting_window_mask;
             } else {
                 for _ in 1..=WINDOW {
-                    square_montgomery_form_assign(&mut z, modulus, mod_neg_inv);
+                    multiplier.square_assign(&mut z);
                 }
             }
 
             // Constant-time lookup in the array of powers
-            let mut power = powers[0].clone();
+            power.limbs.copy_from_slice(&powers[0].limbs);
             for i in 1..(1 << WINDOW) {
                 power.conditional_assign(&powers[i as usize], i.ct_eq(&idx));
             }
 
-            mul_montgomery_form_assign(&mut z, &power, modulus, mod_neg_inv);
+            multiplier.mul_assign(&mut z, &power);
         }
     }
 

src/uint.rs

@@ -19,7 +19,7 @@ pub(crate) mod div_limb;
 mod encoding;
 mod from;
 mod inv_mod;
-mod mul;
+pub(crate) mod mul;
 mod mul_mod;
 mod neg;
 mod neg_mod;

src/uint/boxed.rs

@@ -380,13 +380,19 @@ impl From<&[Limb]> for BoxedUint {
 
 impl From<Box<[Limb]>> for BoxedUint {
     fn from(limbs: Box<[Limb]>) -> BoxedUint {
-        Self { limbs }
+        Vec::from(limbs).into()
     }
 }
 
 impl From<Vec<Limb>> for BoxedUint {
-    fn from(limbs: Vec<Limb>) -> BoxedUint {
-        limbs.into_boxed_slice().into()
+    fn from(mut limbs: Vec<Limb>) -> BoxedUint {
+        if limbs.is_empty() {
+            limbs.push(Limb::ZERO);
+        }
+
+        Self {
+            limbs: limbs.into_boxed_slice(),
+        }
     }
 }
 

tests/boxed_residue_proptests.proptest-regressions

@@ -0,0 +1,7 @@
+# Seeds for failure cases proptest has generated in the past. It is
+# automatically read and these particular cases re-run before any
+# novel cases are generated.
+#
+# It is recommended to check this file in to source control so that
+# everyone who runs the test benefits from these saved cases.
+cc a2dbf0ee6304db81982e37ad9d9145a0f9de45730b9c41221dbf9fdfb9a246c5 # shrinks to a = BoxedUint(0x0000000000000002), b = BoxedUint(0x0000000000000100), n = BoxedResidueParams { modulus: BoxedUint(0x0000000000000003), r: BoxedUint(0x0000000000000001), r2: BoxedUint(0x0000000000000001), r3: BoxedUint(0x0000000000000001), mod_neg_inv: Limb(0x5555555555555555) }