Improve BoxedUint::shl/shr performance by performing operations inplace when possible

Author: fjarri

src/uint/boxed.rs

@@ -253,6 +253,11 @@ impl BoxedUint {
 
         limbs.into()
     }
+
+    /// Set the value of `self` to zero in-place.
+    pub(crate) fn set_to_zero(&mut self) {
+        self.limbs.as_mut().fill(Limb::ZERO)
+    }
 }
 
 impl NonZero<BoxedUint> {

src/uint/boxed/shl.rs

@@ -14,15 +14,16 @@ impl BoxedUint {
         let overflow = !shift.ct_lt(&self.bits_precision());
         let shift = shift % self.bits_precision();
         let mut result = self.clone();
+        let mut temp = self.clone();
 
         for i in 0..shift_bits {
             let bit = Choice::from(((shift >> i) & 1) as u8);
-            result = Self::conditional_select(
-                &result,
-                // Will not overflow by construction
-                &result.shl_vartime(1 << i).expect("shift within range"),
-                bit,
-            );
+            temp.set_to_zero();
+            // Will not overflow by construction
+            result
+                .shl_vartime_into(&mut temp, 1 << i)
+                .expect("shift within range");
+            result.conditional_assign(&temp, bit);
         }
 
         (
@@ -35,45 +36,55 @@ impl BoxedUint {
         )
     }
 
-    /// Computes `self << shift`.
+    /// Computes `self << shift` and writes the result into `dest`.
     /// Returns `None` if `shift >= Self::BITS`.
     ///
+    /// WARNING: for performance reasons, `dest` is assumed to be pre-zeroized.
+    ///
     /// NOTE: this operation is variable time with respect to `shift` *ONLY*.
     ///
     /// When used with a fixed `shift`, this function is constant-time with respect to `self`.
     #[inline(always)]
-    pub fn shl_vartime(&self, shift: u32) -> Option<Self> {
+    fn shl_vartime_into(&self, dest: &mut Self, shift: u32) -> Option<()> {
         if shift >= self.bits_precision() {
             return None;
         }
 
         let nlimbs = self.nlimbs();
-        let mut limbs = vec![Limb::ZERO; nlimbs].into_boxed_slice();
-
         let shift_num = (shift / Limb::BITS) as usize;
         let rem = shift % Limb::BITS;
 
-        let mut i = nlimbs;
-        while i > shift_num {
-            i -= 1;
-            limbs[i] = self.limbs[i - shift_num];
+        for i in shift_num..nlimbs {
+            dest.limbs[i] = self.limbs[i - shift_num];
         }
 
         if rem == 0 {
-            return Some(Self { limbs });
+            return Some(());
         }
 
         let mut carry = Limb::ZERO;
 
-        while i < nlimbs {
-            let shifted = limbs[i].shl(rem);
-            let new_carry = limbs[i].shr(Limb::BITS - rem);
-            limbs[i] = shifted.bitor(carry);
+        for i in shift_num..nlimbs {
+            let shifted = dest.limbs[i].shl(rem);
+            let new_carry = dest.limbs[i].shr(Limb::BITS - rem);
+            dest.limbs[i] = shifted.bitor(carry);
             carry = new_carry;
-            i += 1;
         }
 
-        Some(Self { limbs })
+        Some(())
+    }
+
+    /// Computes `self << shift`.
+    /// Returns `None` if `shift >= Self::BITS`.
+    ///
+    /// NOTE: this operation is variable time with respect to `shift` *ONLY*.
+    ///
+    /// When used with a fixed `shift`, this function is constant-time with respect to `self`.
+    #[inline(always)]
+    pub fn shl_vartime(&self, shift: u32) -> Option<Self> {
+        let mut result = Self::zero_with_precision(self.bits_precision());
+        let success = self.shl_vartime_into(&mut result, shift);
+        success.map(|_| result)
     }
 
     /// Computes `self >> 1` in constant-time.
@@ -118,6 +129,18 @@ impl ShlAssign<u32> for BoxedUint {
 mod tests {
     use super::BoxedUint;
 
+    #[test]
+    fn shl() {
+        let one = BoxedUint::one_with_precision(128);
+
+        assert_eq!(BoxedUint::from(2u8), one.shl(1).0);
+        assert_eq!(BoxedUint::from(4u8), one.shl(2).0);
+        assert_eq!(
+            BoxedUint::from(0x80000000000000000u128),
+            one.shl_vartime(67).unwrap()
+        );
+    }
+
     #[test]
     fn shl_vartime() {
         let one = BoxedUint::one_with_precision(128);

src/uint/boxed/shr.rs

@@ -14,15 +14,16 @@ impl BoxedUint {
         let overflow = !shift.ct_lt(&self.bits_precision());
         let shift = shift % self.bits_precision();
         let mut result = self.clone();
+        let mut temp = self.clone();
 
         for i in 0..shift_bits {
             let bit = Choice::from(((shift >> i) & 1) as u8);
-            result = Self::conditional_select(
-                &result,
-                // Will not overflow by construction
-                &result.shr_vartime(1 << i).expect("shift within range"),
-                bit,
-            );
+            temp.set_to_zero();
+            // Will not overflow by construction
+            result
+                .shr_vartime_into(&mut temp, 1 << i)
+                .expect("shift within range");
+            result.conditional_assign(&temp, bit);
         }
 
         (
@@ -38,42 +39,50 @@ impl BoxedUint {
     /// Computes `self >> shift`.
     /// Returns `None` if `shift >= Self::BITS`.
     ///
+    /// WARNING: for performance reasons, `dest` is assumed to be pre-zeroized.
+    ///
     /// NOTE: this operation is variable time with respect to `shift` *ONLY*.
     ///
     /// When used with a fixed `shift`, this function is constant-time with respect to `self`.
     #[inline(always)]
-    pub fn shr_vartime(&self, shift: u32) -> Option<Self> {
+    fn shr_vartime_into(&self, dest: &mut Self, shift: u32) -> Option<()> {
         if shift >= self.bits_precision() {
             return None;
         }
 
         let nlimbs = self.nlimbs();
-        let mut limbs = vec![Limb::ZERO; nlimbs].into_boxed_slice();
-
         let shift_num = (shift / Limb::BITS) as usize;
         let rem = shift % Limb::BITS;
 
-        let mut i = 0;
-        while i < nlimbs - shift_num {
-            limbs[i] = self.limbs[i + shift_num];
-            i += 1;
+        for i in 0..nlimbs - shift_num {
+            dest.limbs[i] = self.limbs[i + shift_num];
         }
 
         if rem == 0 {
-            return Some(Self { limbs });
+            return Some(());
         }
 
-        let mut carry = Limb::ZERO;
-
-        while i > 0 {
-            i -= 1;
-            let shifted = limbs[i].shr(rem);
-            let new_carry = limbs[i].shl(Limb::BITS - rem);
-            limbs[i] = shifted.bitor(carry);
-            carry = new_carry;
+        for i in 0..nlimbs - shift_num - 1 {
+            let shifted = dest.limbs[i].shr(rem);
+            let carry = dest.limbs[i + 1].shl(Limb::BITS - rem);
+            dest.limbs[i] = shifted.bitor(carry);
         }
+        dest.limbs[nlimbs - shift_num - 1] = dest.limbs[nlimbs - shift_num - 1].shr(rem);
 
-        Some(Self { limbs })
+        Some(())
+    }
+
+    /// Computes `self >> shift`.
+    /// Returns `None` if `shift >= Self::BITS`.
+    ///
+    /// NOTE: this operation is variable time with respect to `shift` *ONLY*.
+    ///
+    /// When used with a fixed `shift`, this function is constant-time with respect to `self`.
+    #[inline(always)]
+    pub fn shr_vartime(&self, shift: u32) -> Option<Self> {
+        let mut result = Self::zero_with_precision(self.bits_precision());
+        let success = self.shr_vartime_into(&mut result, shift);
+        success.map(|_| result)
     }
 
     /// Computes `self >> 1` in constant-time, returning a true [`Choice`] if the overflowing bit
@@ -142,6 +151,15 @@ mod tests {
         assert_eq!(n, n_shr1);
     }
 
+    #[test]
+    fn shr() {
+        let n = BoxedUint::from(0x80000000000000000u128);
+        assert_eq!(BoxedUint::zero(), n.shr(68).0);
+        assert_eq!(BoxedUint::one(), n.shr(67).0);
+        assert_eq!(BoxedUint::from(2u8), n.shr(66).0);
+        assert_eq!(BoxedUint::from(4u8), n.shr(65).0);
+    }
+
     #[test]
     fn shr_vartime() {
         let n = BoxedUint::from(0x80000000000000000u128);

tests/boxed_uint_proptests.rs

@@ -5,6 +5,7 @@
 use core::cmp::Ordering;
 use crypto_bigint::{BoxedUint, CheckedAdd, Limb, NonZero};
 use num_bigint::{BigUint, ModInverse};
+use num_traits::identities::One;
 use proptest::prelude::*;
 
 fn to_biguint(uint: &BoxedUint) -> BigUint {
@@ -212,4 +213,21 @@ proptest! {
             prop_assert_eq!(expected, to_biguint(&actual));
         }
     }
+
+    #[test]
+    fn boxed_shl(a in uint(), shift in any::<u16>()) {
+        let a_bi = to_biguint(&a);
+
+        // Add a 50% probability of overflow.
+        let shift = u32::from(shift) % (a.bits_precision() * 2);
+
+        let expected = to_uint((a_bi << shift as usize) & ((BigUint::one() << a.bits_precision() as usize) - BigUint::one()));
+        let (actual, overflow) = a.shl(shift);
+
+        assert_eq!(expected, actual);
+        if shift >= a.bits_precision() {
+            assert_eq!(actual, BoxedUint::zero());
+            assert!(bool::from(overflow));
+        }
+    }
 }