gh-156443: Keep PyLong loop carries as twodigits in shifts and division

Author: XiaohongGong

Lib/test/test_long.py

@@ -204,7 +204,56 @@ def test_division(self):
         self.check_division(710031681576388032, 26769404391308)
         self.check_division(1933622614268221, 30212853348836)
 
+    def test_divmod_full_limbs(self):
+        # Use saturated limbs and q near BASE to maximize |z| (~BASE*q)
+        # in the x_divrem inner loop: z = vk[i] + zhi - q*w0[i].
+        for n_div in (2, 3, 5, 8):
+            # n_div MASK limbs for the divisor (w0)
+            w = (1 << (n_div * SHIFT)) - 1
+            for n_num in (n_div, n_div + 1, n_div + 4):
+                # n_num MASK limbs for the dividend (vk)
+                v = (1 << (n_num * SHIFT)) - 1
+                with self.subTest(n_div=n_div, n_num=n_num):
+                    self.check_division(v, w)
+
+            # Known quotient and remainder: q includes values
+            # near BASE and r spans 0 .. w-1.
+            for q in (1, 2, MASK, BASE - 1, BASE, BASE + 1):
+                for r in (0, 1, MASK, w - 1):
+                    with self.subTest(n_div=n_div, q=q, r=r):
+                        v = q * w + r
+                        self.assertEqual(divmod(v, w), (q, r))
+                        self.check_division(v, w)
 
+    @support.requires_IEEE_754
+    def test_intradigit_shift(self):
+        # Unit tests for v_lshift and v_rshift in longobject.c.
+        # These two functions are not used by Python << and >>,
+        # so it is different with tests of test_xxx_l|rshift.
+        # We test them with other functions that use them.
+
+        # Full limb values.
+        one = (1 << SHIFT) - 1
+        two = (1 << (2 * SHIFT)) - 1
+        three = (1 << (3 * SHIFT)) - 1
+        four = (1 << (4 * SHIFT)) - 1
+        # Powers of 10.
+        ten_to_40 = 10**40
+        ten_to_20 = 10**20
+
+        # Test with "_PyLong_Frexp" (n -> float):
+        # - n.bit_length() <= 55  => v_lshift,
+        # - n.bit_length() > 55   => v_rshift.
+        self.check_float_conversion(one)
+        self.check_float_conversion(two)
+        self.check_float_conversion(ten_to_40)
+
+        # Test with "long_true_divide" (a / b):
+        # - (a.bit_length() - b.bit_length()) <= 55  => v_lshift,
+        # - (a.bit_length() - b.bit_length()) > 55   => v_rshift.
+        self.check_truediv(three, two)
+        self.check_truediv(four, two)
+        self.check_truediv(ten_to_40, ten_to_20)
 
     def test_karatsuba(self):
         digits = list(range(1, 5)) + list(range(KARATSUBA_CUTOFF,

Objects/longobject.c

@@ -1905,15 +1905,15 @@ static digit
 v_lshift(digit *z, digit *a, Py_ssize_t m, int d)
 {
     Py_ssize_t i;
-    digit carry = 0;
+    twodigits carry = 0;
 
     assert(0 <= d && d < PyLong_SHIFT);
     for (i=0; i < m; i++) {
         twodigits acc = (twodigits)a[i] << d | carry;
         z[i] = (digit)acc & PyLong_MASK;
-        carry = (digit)(acc >> PyLong_SHIFT);
+        carry = acc >> PyLong_SHIFT;
     }
-    return carry;
+    return (digit)carry;
 }
 
 /* Shift digit vector a[0:m] d bits right, with 0 <= d < PyLong_SHIFT.  Put
@@ -1923,16 +1923,16 @@ static digit
 v_rshift(digit *z, digit *a, Py_ssize_t m, int d)
 {
     Py_ssize_t i;
-    digit carry = 0;
-    digit mask = ((digit)1 << d) - 1U;
+    twodigits carry = 0;
+    twodigits mask = ((twodigits)1 << d) - 1U;
 
     assert(0 <= d && d < PyLong_SHIFT);
     for (i=m; i-- > 0;) {
-        twodigits acc = (twodigits)carry << PyLong_SHIFT | a[i];
-        carry = (digit)acc & mask;
+        twodigits acc = carry << PyLong_SHIFT | a[i];
+        carry = acc & mask;
         z[i] = (digit)(acc >> d);
     }
-    return carry;
+    return (digit)carry;
 }
 
 /* Divide long pin, w/ size digits, by non-zero digit n, storing quotient
@@ -3325,7 +3325,7 @@ x_divrem(PyLongObject *v1, PyLongObject *w1, PyLongObject **prem)
     int d;
     digit wm1, wm2, carry, q, r, vtop, *v0, *vk, *w0, *ak;
     twodigits vv;
-    sdigit zhi;
+    stwodigits zhi;
     stwodigits z;
 
     /* We follow Knuth [The Art of Computer Programming, Vol. 2 (3rd
@@ -3416,11 +3416,10 @@ x_divrem(PyLongObject *v1, PyLongObject *w1, PyLongObject **prem)
         for (i = 0; i < size_w; ++i) {
             /* invariants: -PyLong_BASE <= -q <= zhi <= 0;
                -PyLong_BASE * q <= z < PyLong_BASE */
-            z = (sdigit)vk[i] + zhi -
+            z = (stwodigits)(sdigit)vk[i] + zhi -
                 (stwodigits)q * (stwodigits)w0[i];
             vk[i] = (digit)z & PyLong_MASK;
-            zhi = (sdigit)Py_ARITHMETIC_RIGHT_SHIFT(stwodigits,
-                                                    z, PyLong_SHIFT);
+            zhi = Py_ARITHMETIC_RIGHT_SHIFT(stwodigits, z, PyLong_SHIFT);
         }
 
         /* add w back if q was too large (this branch taken rarely) */