Have pow_montgomery_form_amm reduce its output (#1133)

Extracts an `almost_montgomery_reduce` function which accepts a `&mut
UintRef` as input and conditionally subtracts the modulus twice, as
explained in extensive comments.

This required porting `conditional_borrowing_sub_assign` to `UintRef`,
as well as adding a `CtLt` impl to `UintRef`. `Uint::lt` and
`BoxedUint::conditional_borrowing_sub_assign` now both delegate these
functions to `UintRef`.

This also adds a `AsRef<GenericMontyParams<Self::Integer>>` bound to
`Monty::Params` which can be used to obtain the underlying struct, which
is used to query the params for the modulus in `pow_montgomery_form_amm`

Author: tarcieri

src/modular/boxed_monty_form.rs

@@ -127,6 +127,12 @@ impl BoxedMontyParams {
     }
 }
 
+impl AsRef<GenericMontyParams<BoxedUint>> for BoxedMontyParams {
+    fn as_ref(&self) -> &GenericMontyParams<BoxedUint> {
+        &self.0
+    }
+}
+
 impl Debug for BoxedMontyParams {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         self.0.debug_struct(f.debug_struct("BoxedMontyParams"))

src/modular/boxed_monty_form/mul.rs

@@ -7,7 +7,7 @@
 
 use super::{BoxedMontyForm, BoxedMontyParams};
 use crate::{
-    AmmMultiplier, BoxedUint, CtLt, Limb, MontyMultiplier, Mul, MulAssign, Square, SquareAssign,
+    AmmMultiplier, BoxedUint, Limb, MontyMultiplier, Mul, MulAssign, Square, SquareAssign,
     modular::mul::montgomery_multiply_inner, word,
 };
 
@@ -49,32 +49,6 @@ impl BoxedMontyForm {
             params: self.params.clone(),
         }
     }
-
-    /// Instantiate [`BoxedMontyForm`] from the result of an "Almost Montgomery Multiplication".
-    pub(crate) fn from_amm(mut z: BoxedUint, params: BoxedMontyParams) -> Self {
-        // Ensure the output is properly reduced.
-        //
-        // Using the properties of `almost_montgomery_mul()` (see its documentation):
-        // - We have an incoming `x` which is fully reduced (`floor(x / modulus) = 0`).
-        // - We build an array of `powers` which are produced by multiplying the previous power by
-        //   `x`, so for each power `floor(power / modulus) <= 1`.
-        // - Then we take turns squaring the accumulator `z` (bringing `floor(z / modulus)` to 1
-        //   regardless of the previous reduction level) and multiplying by a power of `x`
-        //   (bringing `floor(z / modulus)` to at most 2).
-        // - Then we either exit the loop, or square again, which brings `floor(z / modulus)` back
-        //   to 1.
-        //
-        // We now need to reduce `z` at most twice to bring it within `[0, modulus)`.
-        let modulus = params.modulus();
-        z.conditional_borrowing_sub_assign(modulus, !z.ct_lt(modulus));
-        z.conditional_borrowing_sub_assign(modulus, !z.ct_lt(modulus));
-        debug_assert!(&z < modulus);
-
-        Self {
-            montgomery_form: z,
-            params,
-        }
-    }
 }
 
 impl Mul<&BoxedMontyForm> for &BoxedMontyForm {

src/modular/boxed_monty_form/pow.rs

@@ -17,10 +17,15 @@ impl BoxedMontyForm {
     /// NOTE: `exponent_bits` may be leaked in the time pattern.
     #[must_use]
     pub fn pow_bounded_exp(&self, exponent: &BoxedUint, exponent_bits: u32) -> Self {
-        let z =
-            pow_montgomery_form_amm(&self.montgomery_form, exponent, exponent_bits, &self.params);
-
-        Self::from_amm(z, self.params.clone())
+        Self {
+            montgomery_form: pow_montgomery_form_amm(
+                &self.montgomery_form,
+                exponent,
+                exponent_bits,
+                &self.params,
+            ),
+            params: self.params.clone(),
+        }
     }
 }
 

src/modular/monty_params.rs

@@ -70,6 +70,12 @@ impl<U: Unsigned> GenericMontyParams<U> {
     }
 }
 
+impl<U: Unsigned> AsRef<Self> for GenericMontyParams<U> {
+    fn as_ref(&self) -> &Self {
+        self
+    }
+}
+
 impl<U: Unsigned> CtAssign for GenericMontyParams<U> {
     fn ct_assign(&mut self, other: &Self, choice: Choice) {
         self.modulus.ct_assign(&other.modulus, choice);

src/modular/mul.rs

@@ -1,5 +1,23 @@
 use super::reduction::montgomery_reduction;
-use crate::{Limb, Odd, Uint, WideWord, Word};
+use crate::{CtLt, Limb, Odd, Uint, UintRef, WideWord, Word};
+
+/// Ensure the output of an "almost" Montgomery multiplication is properly reduced.
+///
+/// Using the properties of `almost_montgomery_mul()` (see its documentation):
+/// - We have an incoming `x` which is fully reduced (`floor(x / modulus) = 0`).
+/// - We build an array of `powers` which are produced by multiplying the previous power by
+///   `x`, so for each power `floor(power / modulus) <= 1`.
+/// - Then we take turns squaring the accumulator `z` (bringing `floor(z / modulus)` to 1
+///   regardless of the previous reduction level) and multiplying by a power of `x`
+///   (bringing `floor(z / modulus)` to at most 2).
+/// - Then we either exit the loop, or square again, which brings `floor(z / modulus)` back
+///   to 1.
+///
+/// We now need to reduce `z` at most twice to bring it within `[0, modulus)`.
+pub(crate) fn almost_montgomery_reduce(z: &mut UintRef, modulus: &UintRef) {
+    z.conditional_borrowing_sub_assign(modulus, !z.ct_lt(modulus));
+    z.conditional_borrowing_sub_assign(modulus, !z.ct_lt(modulus));
+}
 
 /// Based on Algorithm 14.36 in Handbook of Applied Cryptography
 /// <https://cacr.uwaterloo.ca/hac/about/chap14.pdf>

src/modular/pow.rs

@@ -1,5 +1,5 @@
 use super::MontyParams;
-use super::mul::{mul_montgomery_form, square_montgomery_form};
+use super::mul::{almost_montgomery_reduce, mul_montgomery_form, square_montgomery_form};
 use crate::{AmmMultiplier, CtEq, Limb, Monty, Uint, Unsigned, Word, word};
 use core::{array, mem};
 
@@ -32,11 +32,7 @@ pub const fn pow_montgomery_form<
 
 /// Performs modular exponentiation using "Almost Montgomery Multiplication".
 ///
-/// NOTE: the resulting output will be reduced to the *bit length* of the modulus, but not fully
-/// reduced and may exceed the modulus.
-///
-/// A final reduction is required to ensure AMM results are fully reduced, and should not be exposed
-/// outside the internals of this crate.
+/// Returns a result which has been fully reduced by the modulus specified in `params`.
 ///
 /// `exponent_bits` represents the length of the exponent in bits.
 ///
@@ -121,6 +117,7 @@ where
         }
     }
 
+    almost_montgomery_reduce(z.as_mut(), params.as_ref().modulus().as_uint_ref());
     z
 }
 

src/traits.rs

@@ -10,7 +10,10 @@ pub use num_traits::{
     WrappingSub,
 };
 
-use crate::{Choice, CtOption, Limb, NonZero, Odd, Reciprocal, UintRef, modular::Retrieve};
+use crate::{
+    Choice, CtOption, Limb, NonZero, Odd, Reciprocal, UintRef,
+    modular::{GenericMontyParams, Retrieve},
+};
 use core::fmt::{self, Debug};
 
 #[cfg(feature = "rand_core")]
@@ -1137,7 +1140,14 @@ pub trait Monty:
     type Multiplier<'a>: Debug + Clone + MontyMultiplier<'a, Monty = Self>;
 
     /// The precomputed data needed for this representation.
-    type Params: 'static + Clone + Debug + Eq + Sized + Send + Sync;
+    type Params: 'static
+        + AsRef<GenericMontyParams<Self::Integer>>
+        + Clone
+        + Debug
+        + Eq
+        + Sized
+        + Send
+        + Sync;
 
     /// Create the precomputed data for Montgomery representation of integers modulo `modulus`,
     /// variable time in `modulus`.

src/uint/boxed/sub.rs

@@ -1,8 +1,8 @@
 //! [`BoxedUint`] subtraction operations.
 
 use crate::{
-    BoxedUint, CheckedSub, Choice, CtOption, CtSelect, Limb, Sub, SubAssign, U64, U128, Uint,
-    Wrapping, WrappingSub,
+    BoxedUint, CheckedSub, Choice, CtOption, Limb, Sub, SubAssign, U64, U128, Uint, Wrapping,
+    WrappingSub,
 };
 
 impl BoxedUint {
@@ -61,18 +61,8 @@ impl BoxedUint {
         rhs: &Self,
         choice: Choice,
     ) -> Choice {
-        debug_assert!(self.bits_precision() <= rhs.bits_precision());
-        let mask = Limb::ct_select(&Limb::ZERO, &Limb::MAX, choice);
-        let mut borrow = Limb::ZERO;
-
-        for i in 0..self.nlimbs() {
-            let masked_rhs = *rhs.limbs.get(i).unwrap_or(&Limb::ZERO) & mask;
-            let (limb, b) = self.limbs[i].borrowing_sub(masked_rhs, borrow);
-            self.limbs[i] = limb;
-            borrow = b;
-        }
-
-        borrow.lsb_to_choice()
+        self.as_mut_uint_ref()
+            .conditional_borrowing_sub_assign(rhs.as_uint_ref(), choice)
     }
 }
 

src/uint/cmp.rs

@@ -3,7 +3,7 @@
 //! Constant-time unless explicitly noted otherwise.
 
 use super::Uint;
-use crate::{Choice, CtEq, Limb, word};
+use crate::{Choice, CtEq, Limb, UintRef, word};
 use core::cmp::Ordering;
 
 impl<const LIMBS: usize> Uint<LIMBS> {
@@ -46,11 +46,7 @@ impl<const LIMBS: usize> Uint<LIMBS> {
     /// Returns the truthy value if `self < rhs` and the falsy value otherwise.
     #[inline]
     pub(crate) const fn lt(lhs: &Self, rhs: &Self) -> Choice {
-        // We could use the same approach as in Limb::ct_lt(),
-        // but since we have to use Uint::wrapping_sub(), which calls `borrowing_sub()`,
-        // there are no savings compared to just calling `borrowing_sub()` directly.
-        let (_res, borrow) = lhs.borrowing_sub(rhs, Limb::ZERO);
-        word::choice_from_mask(borrow.0)
+        UintRef::lt(lhs.as_uint_ref(), rhs.as_uint_ref())
     }
 
     /// Returns the truthy value if `self <= rhs` and the falsy value otherwise.

src/uint/ref_type/cmp.rs

@@ -47,4 +47,41 @@ impl UintRef {
         }
         true
     }
+
+    /// Returns the truthy value if `self < rhs` and the falsy value otherwise.
+    ///
+    /// # Panics
+    /// If `lhs` and `rhs` are not the same size
+    #[inline]
+    pub(crate) const fn lt(lhs: &Self, rhs: &Self) -> Choice {
+        // TODO(tarcieri): support for differently sized inputs?
+        debug_assert!(
+            lhs.bits_precision() == rhs.bits_precision(),
+            "lhs and rhs sizes differ"
+        );
+
+        // NOTE: this is effectively a `borrowing_sub` that only computes the borrow
+        let mut i = 0;
+        let mut borrow = Limb::ZERO;
+
+        while i < lhs.nlimbs() {
+            borrow = lhs.limbs[i].borrowing_sub(rhs.limbs[i], borrow).1;
+            i += 1;
+        }
+
+        borrow.lsb_to_choice()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::UintRef;
+    use crate::Limb;
+
+    #[test]
+    fn lt() {
+        let lesser = UintRef::new(&[Limb::ZERO, Limb::ZERO, Limb::ZERO]);
+        let greater = UintRef::new(&[Limb::ZERO, Limb::ONE, Limb::ZERO]);
+        assert!(UintRef::lt(lesser, greater).to_bool());
+    }
 }