Fix vartime division edge case

src/uint/boxed/div.rs

@@ -1,9 +1,9 @@
 //! [`BoxedUint`] division operations.
 
 use crate::{
-    uint::{boxed, div_limb::div2by1},
+    uint::{boxed, div_limb::div3by2},
     BoxedUint, CheckedDiv, ConstChoice, ConstantTimeSelect, DivRemLimb, Limb, NonZero, Reciprocal,
-    RemLimb, WideWord, Word, Wrapping,
+    RemLimb, Wrapping,
 };
 use core::ops::{Div, DivAssign, Rem, RemAssign};
 use subtle::{Choice, ConstantTimeLess, CtOption};
@@ -360,19 +360,7 @@ pub(crate) fn div_rem_vartime_in_place(x: &mut [Limb], y: &mut [Limb]) {
 
     for xi in (yc - 1..xc).rev() {
         // Divide high dividend words by the high divisor word to estimate the quotient word
-        let (mut quo, mut rem) = div2by1(x_hi.0, x[xi].0, &reciprocal);
-
-        for _ in 0..2 {
-            let qy = (quo as WideWord) * (y[yc - 2].0 as WideWord);
-            let rx = ((rem as WideWord) << Word::BITS) | (x[xi - 1].0 as WideWord);
-            // Constant-time check for q*y[-2] < r*x[-1], based on ConstChoice::from_word_lt
-            let diff = ConstChoice::from_word_lsb(
-                ((((!rx) & qy) | (((!rx) | qy) & (rx.wrapping_sub(qy)))) >> (WideWord::BITS - 1))
-                    as Word,
-            );
-            quo = diff.select_word(quo, quo.saturating_sub(1));
-            rem = diff.select_word(rem, rem.saturating_add(y[yc - 1].0));
-        }
+        let (mut quo, _) = div3by2(x_hi.0, x[xi].0, x[xi - 1].0, &reciprocal, y[yc - 2].0);
 
         // Subtract q*divisor from the dividend
         carry = Limb::ZERO;

src/uint/div.rs

@@ -1,12 +1,10 @@
 //! [`Uint`] division operations.
 
 use super::div_limb::{
-    div2by1, div_rem_limb_with_reciprocal, rem_limb_with_reciprocal, rem_limb_with_reciprocal_wide,
+    div3by2, div_rem_limb_with_reciprocal, rem_limb_with_reciprocal, rem_limb_with_reciprocal_wide,
     Reciprocal,
 };
-use crate::{
-    CheckedDiv, ConstChoice, DivRemLimb, Limb, NonZero, RemLimb, Uint, WideWord, Word, Wrapping,
-};
+use crate::{CheckedDiv, ConstChoice, DivRemLimb, Limb, NonZero, RemLimb, Uint, Word, Wrapping};
 use core::ops::{Div, DivAssign, Rem, RemAssign};
 use subtle::CtOption;
 
@@ -135,21 +133,7 @@ impl<const LIMBS: usize> Uint<LIMBS> {
 
         loop {
             // Divide high dividend words by the high divisor word to estimate the quotient word
-            let (mut quo, mut rem) = div2by1(x_hi.0, x[xi].0, &reciprocal);
-
-            i = 0;
-            while i < 2 {
-                let qy = (quo as WideWord) * (y[yc - 2].0 as WideWord);
-                let rx = ((rem as WideWord) << Word::BITS) | (x[xi - 1].0 as WideWord);
-                // Constant-time check for q*y[-2] < r*x[-1], based on ConstChoice::from_word_lt
-                let diff = ConstChoice::from_word_lsb(
-                    ((((!rx) & qy) | (((!rx) | qy) & (rx.wrapping_sub(qy))))
-                        >> (WideWord::BITS - 1)) as Word,
-                );
-                quo = diff.select_word(quo, quo.saturating_sub(1));
-                rem = diff.select_word(rem, rem.saturating_add(y[yc - 1].0));
-                i += 1;
-            }
+            let (mut quo, _) = div3by2(x_hi.0, x[xi].0, x[xi - 1].0, &reciprocal, y[yc - 2].0);
 
             // Subtract q*divisor from the dividend
             carry = Limb::ZERO;
@@ -283,21 +267,7 @@ impl<const LIMBS: usize> Uint<LIMBS> {
 
         loop {
             // Divide high dividend words by the high divisor word to estimate the quotient word
-            let (mut quo, mut rem) = div2by1(x_hi.0, x[xi].0, &reciprocal);
-
-            i = 0;
-            while i < 2 {
-                let qy = (quo as WideWord) * (y[yc - 2].0 as WideWord);
-                let rx = ((rem as WideWord) << Word::BITS) | (x[xi - 1].0 as WideWord);
-                // Constant-time check for q*y[-2] < r*x[-1], based on ConstChoice::from_word_lt
-                let diff = ConstChoice::from_word_lsb(
-                    ((((!rx) & qy) | (((!rx) | qy) & (rx.wrapping_sub(qy))))
-                        >> (WideWord::BITS - 1)) as Word,
-                );
-                quo = diff.select_word(quo, quo.saturating_sub(1));
-                rem = diff.select_word(rem, rem.saturating_add(y[yc - 1].0));
-                i += 1;
-            }
+            let (quo, _) = div3by2(x_hi.0, x[xi].0, x[xi - 1].0, &reciprocal, y[yc - 2].0);
 
             // Subtract q*divisor from the dividend
             carry = Limb::ZERO;
@@ -940,6 +910,27 @@ mod tests {
         assert_eq!(r, U256::ZERO);
     }
 
+    #[test]
+    fn div_edge() {
+        use crate::U128;
+        let lo = U128::from_be_hex("00000000000000000000000000000001");
+        let hi = U128::from_be_hex("00000000000000000000000000000001");
+        let y = U128::from_be_hex("00000000000000010000000000000001");
+        let x = U256::from((lo, hi));
+        let expect = (U256::from(Limb::MAX), U256::from(2u64));
+
+        let q_r = Uint::div_rem(&x, &NonZero::new(y.resize()).unwrap());
+        assert_eq!(q_r, expect);
+        let (q2, r2) = Uint::div_rem_vartime(&x, &NonZero::new(y).unwrap());
+        assert_eq!((q2, r2.resize()), expect);
+        let r3 = Uint::rem(&x, &NonZero::new(y.resize()).unwrap());
+        assert_eq!(r3, expect.1);
+        let r4 = Uint::rem_vartime(&x, &NonZero::new(y.resize()).unwrap());
+        assert_eq!(r4, expect.1);
+        let r5 = Uint::rem_wide_vartime((lo, hi), &NonZero::new(y).unwrap());
+        assert_eq!(r5.resize(), expect.1);
+    }
+
     #[test]
     fn reduce_one() {
         let r = U256::from(10u8).rem_vartime(&NonZero::new(U256::ONE).unwrap());

src/uint/div_limb.rs

@@ -5,7 +5,7 @@ use subtle::{Choice, ConditionallySelectable};
 
 use crate::{
     primitives::{addhilo, mulhilo},
-    ConstChoice, Limb, NonZero, Uint, Word,
+    ConstChoice, Limb, NonZero, Uint, WideWord, Word,
 };
 
 /// Calculates the reciprocal of the given 32-bit divisor with the highmost bit set.
@@ -145,6 +145,45 @@ pub(crate) const fn div2by1(u1: Word, u0: Word, reciprocal: &Reciprocal) -> (Wor
     (q1, r)
 }
 
+/// Calculate the quotient and the remainder of the division of a 3-word
+/// dividend `u`, with a precalculated leading-word reciprocal `v` and a second
+/// divisor word `v0`. The dividend and the lower divisor word must be left-shifted
+/// according to the `shift` value  of the reciprocal.
+#[inline(always)]
+pub(crate) const fn div3by2(
+    u2: Word,
+    u1: Word,
+    u0: Word,
+    reciprocal: &Reciprocal,
+    v0: Word,
+) -> (Word, Word) {
+    // This method corresponds to Algorithm Q:
+    // https://janmr.com/blog/2014/04/basic-multiple-precision-long-division/
+
+    let maxed = ConstChoice::from_word_eq(u2, reciprocal.divisor_normalized);
+    let (mut quo, mut rem) = div2by1(maxed.select_word(u2, 0), u1, reciprocal);
+    // When the leading dividend word equals the leading divisor word, cap the quotient
+    // at Word::MAX and set the remainder to the sum of the top dividend words.
+    quo = maxed.select_word(quo, Word::MAX);
+    rem = maxed.select_word(rem, u2.saturating_add(u1));
+
+    let mut i = 0;
+    while i < 2 {
+        let qy = (quo as WideWord) * (v0 as WideWord);
+        let rx = ((rem as WideWord) << Word::BITS) | (u0 as WideWord);
+        // Constant-time check for q*y[-2] < r*x[-1], based on ConstChoice::from_word_lt
+        let diff = ConstChoice::from_word_lsb(
+            ((((!rx) & qy) | (((!rx) | qy) & (rx.wrapping_sub(qy)))) >> (WideWord::BITS - 1))
+                as Word,
+        );
+        quo = diff.select_word(quo, quo.saturating_sub(1));
+        rem = diff.select_word(rem, rem.saturating_add(reciprocal.divisor_normalized));
+        i += 1;
+    }
+
+    (quo, rem)
+}
+
 /// A pre-calculated reciprocal for division by a single limb.
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
 pub struct Reciprocal {