Expand integer trait (#489)

See also: #436 and entropyxyz/crypto-primes#37

Author: xuganyu96

src/const_choice.rs

@@ -1,6 +1,6 @@
 use subtle::{Choice, CtOption};
 
-use crate::{modular::BernsteinYangInverter, NonZero, Odd, Uint, Word};
+use crate::{modular::BernsteinYangInverter, Limb, NonZero, Odd, Uint, Word};
 
 /// A boolean value returned by constant-time `const fn`s.
 // TODO: should be replaced by `subtle::Choice` or `CtOption`
@@ -319,6 +319,20 @@ impl<const LIMBS: usize> ConstCtOption<Odd<Uint<LIMBS>>> {
     }
 }
 
+impl ConstCtOption<NonZero<Limb>> {
+    /// Returns the contained value, consuming the `self` value.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the value is none with a custom panic message provided by
+    /// `msg`.
+    #[inline]
+    pub const fn expect(self, msg: &str) -> NonZero<Limb> {
+        assert!(self.is_some.is_true_vartime(), "{}", msg);
+        self.value
+    }
+}
+
 impl<const SAT_LIMBS: usize, const UNSAT_LIMBS: usize>
     ConstCtOption<BernsteinYangInverter<SAT_LIMBS, UNSAT_LIMBS>>
 {

src/modular/boxed_monty_form.rs

@@ -1,4 +1,3 @@
-//! Implements `BoxedMontyForm`s, supporting modular arithmetic with a modulus whose size and value
 //! is chosen at runtime.
 
 mod add;
@@ -9,6 +8,7 @@ mod pow;
 mod sub;
 
 use super::{
+    div_by_2,
     reduction::{montgomery_reduction_boxed, montgomery_reduction_boxed_mut},
     Retrieve,
 };
@@ -219,6 +219,18 @@ impl BoxedMontyForm {
         debug_assert!(self.montgomery_form < self.params.modulus);
         self.montgomery_form.clone()
     }
+
+    /// Performs the modular division by 2, that is for given `x` returns `y`
+    /// such that `y * 2 = x mod p`. This means:
+    /// - if `x` is even, returns `x / 2`,
+    /// - if `x` is odd, returns `(x + p) / 2`
+    ///   (since the modulus `p` in Montgomery form is always odd, this divides entirely).
+    pub fn div_by_2(&self) -> Self {
+        Self {
+            montgomery_form: div_by_2::div_by_2_boxed(&self.montgomery_form, &self.params.modulus),
+            params: self.params.clone(),
+        }
+    }
 }
 
 impl Retrieve for BoxedMontyForm {
@@ -258,3 +270,26 @@ fn convert_to_montgomery(integer: &mut BoxedUint, params: &BoxedMontyParams) {
     #[cfg(feature = "zeroize")]
     product.zeroize();
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{BoxedMontyForm, BoxedMontyParams, BoxedUint, Odd};
+
+    #[test]
+    fn new_params_with_valid_modulus() {
+        let modulus = Odd::new(BoxedUint::from(3u8)).unwrap();
+        BoxedMontyParams::new(modulus);
+    }
+
+    #[test]
+    fn div_by_2() {
+        let modulus = Odd::new(BoxedUint::from(9u8)).unwrap();
+        let params = BoxedMontyParams::new(modulus);
+        let zero = BoxedMontyForm::zero(params.clone());
+        let one = BoxedMontyForm::one(params.clone());
+        let two = one.add(&one);
+
+        assert_eq!(zero.div_by_2(), zero);
+        assert_eq!(one.div_by_2().mul(&two), one);
+    }
+}

src/modular/div_by_2.rs

@@ -1,4 +1,6 @@
 use crate::Uint;
+#[cfg(feature = "alloc")]
+use crate::{BoxedUint, ConstantTimeSelect};
 
 pub(crate) fn div_by_2<const LIMBS: usize>(a: &Uint<LIMBS>, modulus: &Uint<LIMBS>) -> Uint<LIMBS> {
     // We are looking for such `x` that `x * 2 = y mod modulus`,
@@ -28,3 +30,19 @@ pub(crate) fn div_by_2<const LIMBS: usize>(a: &Uint<LIMBS>, modulus: &Uint<LIMBS
 
     Uint::<LIMBS>::select(&if_even, &if_odd, is_odd)
 }
+
+#[cfg(feature = "alloc")]
+pub(crate) fn div_by_2_boxed(a: &BoxedUint, modulus: &BoxedUint) -> BoxedUint {
+    debug_assert_eq!(a.bits_precision(), modulus.bits_precision());
+
+    let (mut half, is_odd) = a.shr1_with_carry();
+    let half_modulus = modulus.shr1();
+
+    let if_odd = half
+        .wrapping_add(&half_modulus)
+        .wrapping_add(&BoxedUint::one_with_precision(a.bits_precision()));
+
+    half.ct_assign(&if_odd, is_odd);
+
+    half
+}

src/odd.rs

@@ -143,3 +143,27 @@ impl Odd<BoxedUint> {
         Odd(ret)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    #[cfg(feature = "alloc")]
+    use super::BoxedUint;
+    use super::{Odd, Uint};
+
+    #[test]
+    fn not_odd_numbers() {
+        let zero = Odd::new(Uint::<4>::ZERO);
+        assert!(bool::from(zero.is_none()));
+        let two = Odd::new(Uint::<4>::from(2u8));
+        assert!(bool::from(two.is_none()));
+    }
+
+    #[cfg(feature = "alloc")]
+    #[test]
+    fn not_odd_numbers_boxed() {
+        let zero = Odd::new(BoxedUint::zero());
+        assert!(bool::from(zero.is_none()));
+        let two = Odd::new(BoxedUint::from(2u8));
+        assert!(bool::from(two.is_none()));
+    }
+}

src/uint/boxed.rs

@@ -10,6 +10,7 @@ mod bits;
 mod cmp;
 mod ct;
 mod div;
+mod div_limb;
 pub(crate) mod encoding;
 mod from;
 mod inv_mod;
@@ -19,6 +20,7 @@ mod neg;
 mod neg_mod;
 mod shl;
 mod shr;
+mod sqrt;
 mod sub;
 mod sub_mod;
 
@@ -92,6 +94,14 @@ impl BoxedUint {
             .fold(Choice::from(1), |acc, limb| acc & limb.is_zero())
     }
 
+    /// Is this [`BoxedUint`] not equal to zero?
+    pub fn is_nonzero(&self) -> Choice {
+        // TODO: why not just !self.is_zero()?
+        self.limbs
+            .iter()
+            .fold(Choice::from(0), |acc, limb| acc | limb.is_nonzero().into())
+    }
+
     /// Is this [`BoxedUint`] equal to one?
     pub fn is_one(&self) -> Choice {
         let mut iter = self.limbs.iter();

src/uint/boxed/bits.rs

@@ -1,6 +1,6 @@
 //! Bit manipulation functions.
 
-use crate::{BoxedUint, Limb, Zero};
+use crate::{BoxedUint, ConstChoice, Limb, Zero};
 use subtle::{Choice, ConditionallySelectable, ConstantTimeEq};
 
 impl BoxedUint {
@@ -24,6 +24,11 @@ impl BoxedUint {
         Limb::BITS * n - leading_zeros
     }
 
+    /// `floor(log2(self.bits_precision()))`.
+    pub(crate) fn log2_bits(&self) -> u32 {
+        u32::BITS - self.bits_precision().leading_zeros() - 1
+    }
+
     /// Calculate the number of bits needed to represent this number in variable-time with respect
     /// to `self`.
     pub fn bits_vartime(&self) -> u32 {
@@ -36,6 +41,18 @@ impl BoxedUint {
         Limb::BITS * (i as u32 + 1) - limb.leading_zeros()
     }
 
+    /// Returns `true` if the bit at position `index` is set, `false` otherwise.
+    ///
+    /// # Remarks
+    /// This operation is variable time with respect to `index` only.
+    pub fn bit_vartime(&self, index: u32) -> bool {
+        if index >= self.bits_precision() {
+            false
+        } else {
+            (self.limbs[(index / Limb::BITS) as usize].0 >> (index % Limb::BITS)) & 1 == 1
+        }
+    }
+
     /// Get the precision of this [`BoxedUint`] in bits.
     pub fn bits_precision(&self) -> u32 {
         self.limbs.len() as u32 * Limb::BITS
@@ -55,6 +72,45 @@ impl BoxedUint {
         count
     }
 
+    /// Calculate the number of trailing ones in the binary representation of this number.
+    pub fn trailing_ones(&self) -> u32 {
+        let limbs = self.as_limbs();
+
+        let mut count = 0;
+        let mut i = 0;
+        let mut nonmax_limb_not_encountered = ConstChoice::TRUE;
+        while i < limbs.len() {
+            let l = limbs[i];
+            let z = l.trailing_ones();
+            count += nonmax_limb_not_encountered.if_true_u32(z);
+            nonmax_limb_not_encountered =
+                nonmax_limb_not_encountered.and(ConstChoice::from_word_eq(l.0, Limb::MAX.0));
+            i += 1;
+        }
+
+        count
+    }
+
+    /// Calculate the number of trailing ones in the binary representation of this number,
+    /// variable time in `self`.
+    pub fn trailing_ones_vartime(&self) -> u32 {
+        let limbs = self.as_limbs();
+
+        let mut count = 0;
+        let mut i = 0;
+        while i < limbs.len() {
+            let l = limbs[i];
+            let z = l.trailing_ones();
+            count += z;
+            if z != Limb::BITS {
+                break;
+            }
+            i += 1;
+        }
+
+        count
+    }
+
     /// Sets the bit at `index` to 0 or 1 depending on the value of `bit_value`.
     pub(crate) fn set_bit(&mut self, index: u32, bit_value: Choice) {
         let limb_num = (index / Limb::BITS) as usize;
@@ -89,6 +145,18 @@ mod tests {
         result
     }
 
+    #[test]
+    fn bit_vartime() {
+        let u = uint_with_bits_at(&[16, 48, 112, 127, 255]);
+        assert!(!u.bit_vartime(0));
+        assert!(!u.bit_vartime(1));
+        assert!(u.bit_vartime(16));
+        assert!(u.bit_vartime(127));
+        assert!(u.bit_vartime(255));
+        assert!(!u.bit_vartime(256));
+        assert!(!u.bit_vartime(260));
+    }
+
     #[test]
     fn bits() {
         assert_eq!(0, BoxedUint::zero().bits());
@@ -119,4 +187,40 @@ mod tests {
         u.set_bit(150, Choice::from(0));
         assert_eq!(u, uint_with_bits_at(&[16, 79]));
     }
+
+    #[test]
+    fn trailing_ones() {
+        let u = !uint_with_bits_at(&[16, 79, 150]);
+        assert_eq!(u.trailing_ones(), 16);
+
+        let u = !uint_with_bits_at(&[79, 150]);
+        assert_eq!(u.trailing_ones(), 79);
+
+        let u = !uint_with_bits_at(&[150, 207]);
+        assert_eq!(u.trailing_ones(), 150);
+
+        let u = !uint_with_bits_at(&[0, 150, 207]);
+        assert_eq!(u.trailing_ones(), 0);
+
+        let u = !BoxedUint::zero_with_precision(256);
+        assert_eq!(u.trailing_ones(), 256);
+    }
+
+    #[test]
+    fn trailing_ones_vartime() {
+        let u = !uint_with_bits_at(&[16, 79, 150]);
+        assert_eq!(u.trailing_ones_vartime(), 16);
+
+        let u = !uint_with_bits_at(&[79, 150]);
+        assert_eq!(u.trailing_ones_vartime(), 79);
+
+        let u = !uint_with_bits_at(&[150, 207]);
+        assert_eq!(u.trailing_ones_vartime(), 150);
+
+        let u = !uint_with_bits_at(&[0, 150, 207]);
+        assert_eq!(u.trailing_ones_vartime(), 0);
+
+        let u = !BoxedUint::zero_with_precision(256);
+        assert_eq!(u.trailing_ones_vartime(), 256);
+    }
 }

src/uint/boxed/cmp.rs

@@ -10,6 +10,30 @@ use subtle::{
     Choice, ConditionallySelectable, ConstantTimeEq, ConstantTimeGreater, ConstantTimeLess,
 };
 
+impl BoxedUint {
+    /// Returns the Ordering between `self` and `rhs` in variable time.
+    pub fn cmp_vartime(&self, rhs: &Self) -> Ordering {
+        debug_assert_eq!(self.limbs.len(), rhs.limbs.len());
+        let mut i = self.limbs.len() - 1;
+        loop {
+            // TODO: investigate if directly comparing limbs is faster than performing a
+            // subtraction between limbs
+            let (val, borrow) = self.limbs[i].sbb(rhs.limbs[i], Limb::ZERO);
+            if val.0 != 0 {
+                return if borrow.0 != 0 {
+                    Ordering::Less
+                } else {
+                    Ordering::Greater
+                };
+            }
+            if i == 0 {
+                return Ordering::Equal;
+            }
+            i -= 1;
+        }
+    }
+}
+
 impl ConstantTimeEq for BoxedUint {
     #[inline]
     fn ct_eq(&self, other: &Self) -> Choice {

src/uint/boxed/div.rs

@@ -1,10 +1,23 @@
 //! [`BoxedUint`] division operations.
 
-use crate::{BoxedUint, CheckedDiv, ConstantTimeSelect, Limb, NonZero, Wrapping};
+use crate::{
+    uint::boxed, BoxedUint, CheckedDiv, ConstantTimeSelect, Limb, NonZero, Reciprocal, Wrapping,
+};
 use core::ops::{Div, DivAssign, Rem, RemAssign};
 use subtle::{Choice, ConstantTimeEq, ConstantTimeLess, CtOption};
 
 impl BoxedUint {
+    /// Computes `self` / `rhs` using a pre-made reciprocal,
+    /// returns the quotient (q) and remainder (r).
+    pub fn div_rem_limb_with_reciprocal(&self, reciprocal: &Reciprocal) -> (Self, Limb) {
+        boxed::div_limb::div_rem_limb_with_reciprocal(self, reciprocal)
+    }
+
+    /// Computes `self` / `rhs`, returns the quotient (q) and remainder (r).
+    pub fn div_rem_limb(&self, rhs: NonZero<Limb>) -> (Self, Limb) {
+        boxed::div_limb::div_rem_limb_with_reciprocal(self, &Reciprocal::new(rhs))
+    }
+
     /// Computes self / rhs, returns the quotient, remainder.
     pub fn div_rem(&self, rhs: &NonZero<Self>) -> (Self, Self) {
         // Since `rhs` is nonzero, this should always hold.
@@ -61,6 +74,14 @@ impl BoxedUint {
         self.div_rem(rhs).0
     }
 
+    /// Wrapped division is just normal division i.e. `self` / `rhs`
+    ///
+    /// There’s no way wrapping could ever happen.
+    /// This function exists, so that all operations are accounted for in the wrapping operations
+    pub fn wrapping_div_vartime(&self, rhs: &NonZero<Self>) -> Self {
+        self.div_rem_vartime(rhs).0
+    }
+
     /// Perform checked division, returning a [`CtOption`] which `is_some`
     /// only if the rhs != 0
     pub fn checked_div(&self, rhs: &Self) -> CtOption<Self> {