use UintRef methods for Uint bit shifts

Signed-off-by: Andrew Whitehead <cywolf@gmail.com>

Author: andrewwhitehead

benches/uint.rs

@@ -740,6 +740,26 @@ fn bench_mod_symbols(c: &mut Criterion) {
 fn bench_shl(c: &mut Criterion) {
     let mut group = c.benchmark_group("left shift");
 
+    group.bench_function("shl, U256", |b| {
+        b.iter_batched(|| U256::ONE, |x| x.shl(30), BatchSize::SmallInput);
+    });
+
+    group.bench_function("unbounded_shl, U256", |b| {
+        b.iter_batched(|| U256::ONE, |x| x.unbounded_shl(30), BatchSize::SmallInput);
+    });
+
+    group.bench_function("shl, U2048", |b| {
+        b.iter_batched(|| U2048::ONE, |x| x.shl(1024 + 10), BatchSize::SmallInput);
+    });
+
+    group.bench_function("unbounded_shl, U2048", |b| {
+        b.iter_batched(
+            || U2048::ONE,
+            |x| x.unbounded_shl(1024 + 10),
+            BatchSize::SmallInput,
+        );
+    });
+
     group.bench_function("shl_vartime, small, U2048", |b| {
         b.iter_batched(|| U2048::ONE, |x| x.shl_vartime(10), BatchSize::SmallInput);
     });
@@ -760,16 +780,32 @@ fn bench_shl(c: &mut Criterion) {
         );
     });
 
-    group.bench_function("shl, U2048", |b| {
-        b.iter_batched(|| U2048::ONE, |x| x.shl(1024 + 10), BatchSize::SmallInput);
-    });
-
     group.finish();
 }
 
 fn bench_shr(c: &mut Criterion) {
     let mut group = c.benchmark_group("right shift");
 
+    group.bench_function("shr, U256", |b| {
+        b.iter_batched(|| U256::ONE, |x| x.shr(30), BatchSize::SmallInput);
+    });
+
+    group.bench_function("unbounded_shr, U256", |b| {
+        b.iter_batched(|| U256::ONE, |x| x.unbounded_shr(30), BatchSize::SmallInput);
+    });
+
+    group.bench_function("shr, U2048", |b| {
+        b.iter_batched(|| U2048::ONE, |x| x.shr(1024 + 10), BatchSize::SmallInput);
+    });
+
+    group.bench_function("unbounded_shr, U2048", |b| {
+        b.iter_batched(
+            || U2048::ONE,
+            |x| x.unbounded_shr(1024 + 10),
+            BatchSize::SmallInput,
+        );
+    });
+
     group.bench_function("shr_vartime, small, U2048", |b| {
         b.iter_batched(|| U2048::ONE, |x| x.shr_vartime(10), BatchSize::SmallInput);
     });
@@ -790,10 +826,6 @@ fn bench_shr(c: &mut Criterion) {
         );
     });
 
-    group.bench_function("shr, U2048", |b| {
-        b.iter_batched(|| U2048::ONE, |x| x.shr(1024 + 10), BatchSize::SmallInput);
-    });
-
     group.finish();
 }
 

src/modular/bingcd/extension.rs

@@ -95,19 +95,17 @@ impl<const LIMBS: usize, const EXTRA: usize> ExtendedUint<LIMBS, EXTRA> {
     pub const fn shr_vartime(&self, shift: u32) -> Self {
         debug_assert!(shift <= Uint::<EXTRA>::BITS);
 
-        let shift_is_zero = Choice::from_u32_eq(shift, 0);
-        let left_shift = shift_is_zero.select_u32(Uint::<EXTRA>::BITS - shift, 0);
-
         let hi = self.1.shr_vartime(shift);
-        // TODO: replace with carrying_shl
-        let carry = Uint::select(&self.1, &Uint::ZERO, shift_is_zero).shl_vartime(left_shift);
+        let carry = self.1.unbounded_shl_vartime(Uint::<EXTRA>::BITS - shift);
         let mut lo = self.0.shr_vartime(shift);
 
         // Apply carry
-        let limb_diff = LIMBS.wrapping_sub(EXTRA) as u32;
+        let limb_diff = LIMBS.saturating_sub(EXTRA) as u32;
         // safe to vartime; shr_vartime is variable in the value of shift only. Since this shift
         // is a public constant, the constant time property of this algorithm is not impacted.
-        let carry = carry.resize::<LIMBS>().shl_vartime(limb_diff * Limb::BITS);
+        let carry = carry
+            .resize::<LIMBS>()
+            .unbounded_shl_by_limbs_vartime(limb_diff);
         lo = lo.bitxor(&carry);
 
         Self(lo, hi)

src/uint/boxed/shl.rs

@@ -165,14 +165,9 @@ impl BoxedUint {
     /// Computes `self << 1` in constant-time.
     pub(crate) fn shl1(&self) -> (Self, Limb) {
         let mut ret = self.clone();
-        let carry = Limb::select(Limb::ZERO, Limb::ONE, ret.shl1_assign());
+        let carry = ret.as_mut_uint_ref().shl1_assign();
         (ret, carry)
     }
-
-    /// Computes `self <<= 1` in constant-time.
-    pub(crate) fn shl1_assign(&mut self) -> Choice {
-        self.as_mut_uint_ref().shl1_assign()
-    }
 }
 
 macro_rules! impl_shl {
@@ -246,10 +241,9 @@ mod tests {
 
     #[test]
     fn shl1_assign() {
-        let mut n = BoxedUint::from(0x3c442b21f19185fe433f0a65af902b8fu128);
+        let n = BoxedUint::from(0x3c442b21f19185fe433f0a65af902b8fu128);
         let n_shl1 = BoxedUint::from(0x78885643e3230bfc867e14cb5f20571eu128);
-        n.shl1_assign();
-        assert_eq!(n, n_shl1);
+        assert_eq!(n.shl1().0, n_shl1);
     }
 
     #[test]

src/uint/boxed/shr.rs

@@ -178,7 +178,7 @@ impl BoxedUint {
 
     /// Computes `self >>= 1` in-place in constant-time.
     pub(crate) fn shr1_assign(&mut self) -> Choice {
-        self.as_mut_uint_ref().shr1_assign()
+        self.as_mut_uint_ref().shr1_assign().lsb_to_choice()
     }
 }
 

src/uint/ref_type/shl.rs

@@ -3,7 +3,7 @@
 use core::num::NonZeroU32;
 
 use super::UintRef;
-use crate::{Choice, Limb, primitives::u32_bits, word};
+use crate::{Choice, Limb, primitives::u32_bits};
 
 #[cfg(feature = "alloc")]
 use crate::primitives::u32_rem;
@@ -50,15 +50,16 @@ impl UintRef {
     pub(crate) const fn bounded_shl_assign(&mut self, shift: u32, shift_upper_bound: u32) {
         assert!(shift < shift_upper_bound, "`shift` exceeds upper bound");
 
-        if shift_upper_bound <= Limb::BITS {
-            self.shl_assign_limb(shift);
+        let bit_shift = if shift_upper_bound <= Limb::BITS {
+            shift
         } else {
             self.bounded_shl_by_limbs_assign(
                 shift >> Limb::LOG2_BITS,
                 shift_upper_bound.div_ceil(Limb::BITS),
             );
-            self.shl_assign_limb(shift & (Limb::BITS - 1));
-        }
+            shift & (Limb::BITS - 1)
+        };
+        self.shl_assign_limb_with_carry(bit_shift, Limb::ZERO);
     }
 
     /// Left-shifts by `shift * Limb::BITS` bits where `shift < shift_upper_bound`.
@@ -107,7 +108,6 @@ impl UintRef {
     /// 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`.
-    #[cfg(feature = "alloc")]
     #[inline(always)]
     pub(crate) const fn unbounded_shl_assign_by_limbs_vartime(&mut self, shift: u32) {
         let shift = shift as usize;
@@ -128,7 +128,6 @@ impl UintRef {
     /// 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`.
-    #[cfg(feature = "alloc")]
     #[inline(always)]
     pub const fn unbounded_shl_assign_vartime(&mut self, shift: u32) {
         let shift_limbs = shift / Limb::BITS;
@@ -170,7 +169,7 @@ impl UintRef {
     /// Left-shifts by a single bit in constant-time, returning [`Choice::TRUE`]
     /// if the least significant bit was set, and [`Choice::FALSE`] otherwise.
     #[inline(always)]
-    pub const fn shl1_assign(&mut self) -> Choice {
+    pub const fn shl1_assign(&mut self) -> Limb {
         let mut carry = Limb::ZERO;
         let mut i = 0;
         while i < self.nlimbs() {
@@ -179,25 +178,25 @@ impl UintRef {
             carry = new_carry;
             i += 1;
         }
-        word::choice_from_lsb(carry.0)
+        carry
     }
 
     /// Conditionally left-shifts by `shift` bits where `0 < shift < Limb::BITS`, returning
     /// the carry.
     ///
     /// # Panics
     /// - if `shift >= Limb::BITS`.
-    #[inline]
+    #[inline(always)]
     pub(crate) const fn conditional_shl_assign_limb_nonzero(
         &mut self,
         shift: NonZeroU32,
+        mut carry: Limb,
         choice: Choice,
     ) -> Limb {
         assert!(shift.get() < Limb::BITS);
 
         let lshift = shift.get();
         let rshift = Limb::BITS - lshift;
-        let mut carry = Limb::ZERO;
 
         let mut i = 0;
         while i < self.nlimbs() {
@@ -221,9 +220,19 @@ impl UintRef {
     /// - if `shift >= Limb::BITS`.
     #[inline(always)]
     pub const fn shl_assign_limb(&mut self, shift: u32) -> Limb {
+        self.shl_assign_limb_with_carry(shift, Limb::ZERO)
+    }
+
+    /// Left-shifts by `shift` bits where `0 < shift < Limb::BITS`, returning the carry.
+    ///
+    /// # Panics
+    /// - if `shift >= Limb::BITS`.
+    #[inline(always)]
+    pub const fn shl_assign_limb_with_carry(&mut self, shift: u32, carry: Limb) -> Limb {
         let nz = Choice::from_u32_nz(shift);
         self.conditional_shl_assign_limb_nonzero(
             NonZeroU32::new(nz.select_u32(1, shift)).expect("ensured non-zero"),
+            carry,
             nz,
         )
     }
@@ -263,9 +272,7 @@ impl UintRef {
 
 #[cfg(test)]
 mod tests {
-    #[cfg(feature = "alloc")]
-    use crate::Uint;
-    use crate::{Limb, U256};
+    use crate::{Limb, U256, Uint};
 
     const N: U256 =
         U256::from_be_hex("FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141");
@@ -278,17 +285,17 @@ mod tests {
         let mut n = N;
         let carry = n.as_mut_uint_ref().shl1_assign();
         assert_eq!(n, TWO_N);
-        assert!(carry.to_bool());
+        assert_eq!(carry, Limb::ONE);
 
         let mut m = U256::MAX;
         let carry = m.as_mut_uint_ref().shl1_assign();
         assert_eq!(m, U256::MAX.shl_vartime(1));
-        assert!(carry.to_bool());
+        assert_eq!(carry, Limb::ONE);
 
         let mut z = U256::ZERO;
         let carry = z.as_mut_uint_ref().shl1_assign();
         assert_eq!(z, U256::ZERO);
-        assert!(!carry.to_bool());
+        assert_eq!(carry, Limb::ZERO);
     }
 
     #[test]
@@ -324,7 +331,6 @@ mod tests {
         assert_eq!(carry, N.limbs[U256::LIMBS - 1].shr(1));
     }
 
-    #[cfg(feature = "alloc")]
     #[test]
     fn unbounded_shl_by_limbs_vartime() {
         let refval = Uint::<2>::from_words([1, 99]);