ConstChoice improvements (#440)

* Add `sqrt` benchmark
* Rename `ct_choice` module to `const_choice`
* Use the Hacker's Delight integer comparison for `Limb'`s subtle impls
* Add `ConstCtOption` struct

Author: fjarri

benches/uint.rs

@@ -160,8 +160,8 @@ fn bench_inv_mod(c: &mut Criterion) {
                 let m = U256::random(&mut OsRng) | U256::ONE;
                 loop {
                     let x = U256::random(&mut OsRng);
-                    let (_, is_some) = x.inv_odd_mod(&m);
-                    if is_some.into() {
+                    let inv_x = x.inv_odd_mod(&m);
+                    if inv_x.is_some().into() {
                         break (x, m);
                     }
                 }
@@ -177,8 +177,8 @@ fn bench_inv_mod(c: &mut Criterion) {
                 let m = U256::random(&mut OsRng) | U256::ONE;
                 loop {
                     let x = U256::random(&mut OsRng);
-                    let (_, is_some) = x.inv_odd_mod(&m);
-                    if is_some.into() {
+                    let inv_x = x.inv_odd_mod(&m);
+                    if inv_x.is_some().into() {
                         break (x, m);
                     }
                 }
@@ -194,8 +194,8 @@ fn bench_inv_mod(c: &mut Criterion) {
                 let m = U256::random(&mut OsRng);
                 loop {
                     let x = U256::random(&mut OsRng);
-                    let (_, is_some) = black_box(x.inv_mod(&m));
-                    if is_some.into() {
+                    let inv_x = x.inv_mod(&m);
+                    if inv_x.is_some().into() {
                         break (x, m);
                     }
                 }
@@ -208,6 +208,25 @@ fn bench_inv_mod(c: &mut Criterion) {
     group.finish();
 }
 
-criterion_group!(benches, bench_shl, bench_shr, bench_division, bench_inv_mod);
+fn bench_sqrt(c: &mut Criterion) {
+    let mut group = c.benchmark_group("sqrt");
+
+    group.bench_function("sqrt, U256", |b| {
+        b.iter_batched(
+            || U256::random(&mut OsRng),
+            |x| x.sqrt(),
+            BatchSize::SmallInput,
+        )
+    });
+}
+
+criterion_group!(
+    benches,
+    bench_shl,
+    bench_shr,
+    bench_division,
+    bench_inv_mod,
+    bench_sqrt
+);
 
 criterion_main!(benches);

src/ct_choice.rs renamed to src/const_choice.rs

@@ -1,6 +1,6 @@
-use subtle::Choice;
+use subtle::{Choice, CtOption};
 
-use crate::Word;
+use crate::{NonZero, Uint, Word};
 
 /// A boolean value returned by constant-time `const fn`s.
 // TODO: should be replaced by `subtle::Choice` or `CtOption`
@@ -71,10 +71,19 @@ impl ConstChoice {
     /// Returns the truthy value if `x < y`, and the falsy value otherwise.
     #[inline]
     pub(crate) const fn from_word_lt(x: Word, y: Word) -> Self {
+        // See "Hacker's Delight" 2nd ed, section 2-12 (Comparison predicates)
         let bit = (((!x) & y) | (((!x) | y) & (x.wrapping_sub(y)))) >> (Word::BITS - 1);
         Self::from_word_lsb(bit)
     }
 
+    /// Returns the truthy value if `x > y`, and the falsy value otherwise.
+    #[inline]
+    pub(crate) const fn from_word_gt(x: Word, y: Word) -> Self {
+        // See "Hacker's Delight" 2nd ed, section 2-12 (Comparison predicates)
+        let bit = (((!y) & x) | (((!y) | x) & (y.wrapping_sub(x)))) >> (Word::BITS - 1);
+        Self::from_word_lsb(bit)
+    }
+
     /// Returns the truthy value if `x < y`, and the falsy value otherwise.
     #[inline]
     pub(crate) const fn from_u32_lt(x: u32, y: u32) -> Self {
@@ -147,6 +156,7 @@ impl ConstChoice {
 }
 
 impl From<ConstChoice> for Choice {
+    #[inline]
     fn from(choice: ConstChoice) -> Self {
         Choice::from(choice.to_u8())
     }
@@ -164,11 +174,144 @@ impl PartialEq for ConstChoice {
     }
 }
 
+/// An equivalent of `subtle::CtOption` usable in a `const fn` context.
+#[derive(Debug, Clone)]
+pub struct ConstCtOption<T> {
+    value: T,
+    is_some: ConstChoice,
+}
+
+impl<T> ConstCtOption<T> {
+    #[inline]
+    pub(crate) const fn new(value: T, is_some: ConstChoice) -> Self {
+        Self { value, is_some }
+    }
+
+    #[inline]
+    pub(crate) const fn some(value: T) -> Self {
+        Self {
+            value,
+            is_some: ConstChoice::TRUE,
+        }
+    }
+
+    #[inline]
+    pub(crate) const fn none(dummy_value: T) -> Self {
+        Self {
+            value: dummy_value,
+            is_some: ConstChoice::FALSE,
+        }
+    }
+
+    /// Returns a reference to the contents of this structure.
+    ///
+    /// **Note:** if the second element is `None`, the first value may take any value.
+    #[inline]
+    pub(crate) const fn components_ref(&self) -> (&T, ConstChoice) {
+        // Since Rust is not smart enough to tell that we would be moving the value,
+        // and hence no destructors will be called, we have to return a reference instead.
+        // See https://github.com/rust-lang/rust/issues/66753
+        (&self.value, self.is_some)
+    }
+
+    /// Returns a true [`ConstChoice`] if this value is `Some`.
+    #[inline]
+    pub const fn is_some(&self) -> ConstChoice {
+        self.is_some
+    }
+
+    /// Returns a true [`ConstChoice`] if this value is `None`.
+    #[inline]
+    pub const fn is_none(&self) -> ConstChoice {
+        self.is_some.not()
+    }
+
+    /// This returns the underlying value but panics if it is not `Some`.
+    #[inline]
+    pub fn unwrap(self) -> T {
+        assert!(self.is_some.is_true_vartime());
+        self.value
+    }
+}
+
+impl<T> From<ConstCtOption<T>> for CtOption<T> {
+    #[inline]
+    fn from(value: ConstCtOption<T>) -> Self {
+        CtOption::new(value.value, value.is_some.into())
+    }
+}
+
+// Need specific implementations to work around the
+// "destructors cannot be evaluated at compile-time" error
+// See https://github.com/rust-lang/rust/issues/66753
+
+impl<const LIMBS: usize> ConstCtOption<Uint<LIMBS>> {
+    /// This returns the underlying value if it is `Some` or the provided value otherwise.
+    #[inline]
+    pub const fn unwrap_or(self, def: Uint<LIMBS>) -> Uint<LIMBS> {
+        Uint::select(&def, &self.value, self.is_some)
+    }
+
+    /// 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) -> Uint<LIMBS> {
+        assert!(self.is_some.is_true_vartime(), "{}", msg);
+        self.value
+    }
+}
+
+impl<const LIMBS: usize> ConstCtOption<(Uint<LIMBS>, Uint<LIMBS>)> {
+    /// 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) -> (Uint<LIMBS>, Uint<LIMBS>) {
+        assert!(self.is_some.is_true_vartime(), "{}", msg);
+        self.value
+    }
+}
+
+impl<const LIMBS: usize> ConstCtOption<NonZero<Uint<LIMBS>>> {
+    /// 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<Uint<LIMBS>> {
+        assert!(self.is_some.is_true_vartime(), "{}", msg);
+        self.value
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::ConstChoice;
     use crate::Word;
 
+    #[test]
+    fn from_word_lt() {
+        assert_eq!(ConstChoice::from_word_lt(4, 5), ConstChoice::TRUE);
+        assert_eq!(ConstChoice::from_word_lt(5, 5), ConstChoice::FALSE);
+        assert_eq!(ConstChoice::from_word_lt(6, 5), ConstChoice::FALSE);
+    }
+
+    #[test]
+    fn from_word_gt() {
+        assert_eq!(ConstChoice::from_word_gt(4, 5), ConstChoice::FALSE);
+        assert_eq!(ConstChoice::from_word_gt(5, 5), ConstChoice::FALSE);
+        assert_eq!(ConstChoice::from_word_gt(6, 5), ConstChoice::TRUE);
+    }
+
     #[test]
     fn select() {
         let a: Word = 1;

src/lib.rs

@@ -167,7 +167,7 @@ pub mod modular;
 #[cfg(feature = "generic-array")]
 mod array;
 mod checked;
-mod ct_choice;
+mod const_choice;
 mod limb;
 mod non_zero;
 mod traits;
@@ -176,7 +176,7 @@ mod wrapping;
 
 pub use crate::{
     checked::Checked,
-    ct_choice::ConstChoice,
+    const_choice::{ConstChoice, ConstCtOption},
     limb::{Limb, WideWord, Word},
     non_zero::NonZero,
     traits::*,

src/limb/cmp.rs

@@ -54,16 +54,14 @@ impl ConstantTimeEq for Limb {
 impl ConstantTimeGreater for Limb {
     #[inline]
     fn ct_gt(&self, other: &Self) -> Choice {
-        let borrow = other.sbb(*self, Limb::ZERO).1;
-        Choice::from(borrow.0 as u8 & 1)
+        ConstChoice::from_word_gt(self.0, other.0).into()
     }
 }
 
 impl ConstantTimeLess for Limb {
     #[inline]
     fn ct_lt(&self, other: &Self) -> Choice {
-        let borrow = self.sbb(*other, Limb::ZERO).1;
-        Choice::from(borrow.0 as u8 & 1)
+        ConstChoice::from_word_lt(self.0, other.0).into()
     }
 }
 

src/modular/boxed_residue/mul.rs

@@ -301,7 +301,7 @@ fn sub_vv(z: &mut [Limb], x: &[Limb], y: &[Limb]) -> Limb {
     for (i, (&xi, &yi)) in x.iter().zip(y.iter()).enumerate().take(z.len()) {
         let zi = xi.wrapping_sub(yi).wrapping_sub(c);
         z[i] = zi;
-        // see "Hacker's Delight", section 2-12 (overflow detection)
+        // See "Hacker's Delight" 2nd ed, section 2-13 (Overflow detection)
         c = ((yi & !xi) | ((yi | !xi) & zi)) >> (Word::BITS - 1)
     }
 

src/modular/dyn_residue.rs

@@ -50,8 +50,9 @@ impl<const LIMBS: usize> DynResidueParams<LIMBS> {
         let r = Uint::MAX.rem(&nz_modulus).wrapping_add(&Uint::ONE);
         let r2 = Uint::rem_wide(r.square_wide(), &nz_modulus);
 
+        let maybe_inverse = modulus.inv_mod2k_vartime(Word::BITS);
         // If the inverse exists, it means the modulus is odd.
-        let (inv_mod_limb, modulus_is_odd) = modulus.inv_mod2k_vartime(Word::BITS);
+        let (inv_mod_limb, modulus_is_odd) = maybe_inverse.components_ref();
         let mod_neg_inv = Limb(Word::MIN.wrapping_sub(inv_mod_limb.limbs[0].0));
 
         let r3 = montgomery_reduction(&r2.square_wide(), modulus, mod_neg_inv);

src/modular/dyn_residue/inv.rs

@@ -4,7 +4,7 @@ use super::{DynResidue, DynResidueParams};
 use crate::{
     modular::{inv::inv_montgomery_form, BernsteinYangInverter},
     traits::Invert,
-    ConstChoice, Inverter, PrecomputeInverter, PrecomputeInverterWithAdjuster, Uint,
+    ConstCtOption, Inverter, PrecomputeInverter, PrecomputeInverterWithAdjuster, Uint,
 };
 use core::fmt;
 use subtle::CtOption;
@@ -14,28 +14,28 @@ impl<const LIMBS: usize> DynResidue<LIMBS> {
     /// I.e. `self * self^-1 = 1`.
     /// If the number was invertible, the second element of the tuple is the truthy value,
     /// otherwise it is the falsy value (in which case the first element's value is unspecified).
-    pub const fn invert(&self) -> (Self, ConstChoice) {
-        let (montgomery_form, is_some) = inv_montgomery_form(
+    pub const fn invert(&self) -> ConstCtOption<Self> {
+        let maybe_inverse = inv_montgomery_form(
             &self.montgomery_form,
             &self.residue_params.modulus,
             &self.residue_params.r3,
             self.residue_params.mod_neg_inv,
         );
+        let (montgomery_form, is_some) = maybe_inverse.components_ref();
 
         let value = Self {
-            montgomery_form,
+            montgomery_form: *montgomery_form,
             residue_params: self.residue_params,
         };
 
-        (value, is_some)
+        ConstCtOption::new(value, is_some)
     }
 }
 
 impl<const LIMBS: usize> Invert for DynResidue<LIMBS> {
     type Output = CtOption<Self>;
     fn invert(&self) -> Self::Output {
-        let (value, is_some) = self.invert();
-        CtOption::new(value, is_some.into())
+        self.invert().into()
     }
 }
 
@@ -114,7 +114,7 @@ mod tests {
             U256::from_be_hex("77117F1273373C26C700D076B3F780074D03339F56DD0EFB60E7F58441FD3685");
         let x_mod = DynResidue::new(&x, params);
 
-        let (inv, _is_some) = x_mod.invert();
+        let inv = x_mod.invert().unwrap();
         let res = x_mod * inv;
 
         assert_eq!(res.retrieve(), U256::ONE);

src/modular/inv.rs

@@ -1,13 +1,14 @@
-use crate::{modular::reduction::montgomery_reduction, ConstChoice, Limb, Uint};
+use crate::{modular::reduction::montgomery_reduction, ConstCtOption, Limb, Uint};
 
 pub const fn inv_montgomery_form<const LIMBS: usize>(
     x: &Uint<LIMBS>,
     modulus: &Uint<LIMBS>,
     r3: &Uint<LIMBS>,
     mod_neg_inv: Limb,
-) -> (Uint<LIMBS>, ConstChoice) {
-    let (inverse, is_some) = x.inv_odd_mod(modulus);
-    (
+) -> ConstCtOption<Uint<LIMBS>> {
+    let maybe_inverse = x.inv_odd_mod(modulus);
+    let (inverse, is_some) = maybe_inverse.components_ref();
+    ConstCtOption::new(
         montgomery_reduction(&inverse.split_mul(r3), modulus, mod_neg_inv),
         is_some,
     )