Rework to remove duplication and move the shift==0 check earlier.

python · markshannon · May 2, 2022 · Dec 27, 2021 · Dec 28, 2021 · Dec 28, 2021
commit ad69e48167b2719a5435a5915a57357a24f9add7
diff --git a/Objects/longobject.c b/Objects/longobject.c
@@ -4486,19 +4486,20 @@ divmod_shift(PyObject *shiftby, Py_ssize_t *wordshift, digit *remshift)
     return 0;
 }
 
+/* Inner function for both long_rshift and _PyLong_Rshift, shifting an
+   integer right by a strictly positive shift. */
+
 static PyObject *
 long_rshift1(PyLongObject *a, Py_ssize_t wordshift, digit remshift)
 {
     PyLongObject *z = NULL;
-    Py_ssize_t newsize, hishift, i, j;
+    Py_ssize_t newsize, hishift, i, j, size_a;
     twodigits accum;
-    digit sticky;
+    int a_negative;
 
-    /* Special-case a shift of zero. */
-    if (wordshift == 0 && remshift == 0) {
-        return long_long((PyObject *)a);
-    }
+    assert (wordshift > 0 || remshift > 0);
 
+    /* Fast path for small a. */
     if (IS_MEDIUM_VALUE(a)) {
         stwodigits m, x;
         digit shift;
@@ -4508,79 +4509,65 @@ long_rshift1(PyLongObject *a, Py_ssize_t wordshift, digit remshift)
         return _PyLong_FromSTwoDigits(x);
     }
 
-    if (Py_SIZE(a) < 0) {
-        /*
-            Right shifting negative numbers is harder. For a positive
-            integer a and nonnegative shift, we have:
-
-                (-a) >> shift == -((a + 2**shift - 1) >> shift).
-
-            With shift == wordshift*SHIFT + remshift, 0 <= remshift <= SHIFT,
-            (a + 2**shift - 1) >> shift is equal to
-
-                (a + 2**shift - 1) >> wordshift*SHIFT >> remshift.
-
-            If the bottom `wordshift` digits of a are all zero, this is the
-            equal to
-
-                ((a >> wordshift*SHIFT) + 2**remshift - 1) >> remshift.
-
-            Otherwise, it's equal to:
-
-                ((a >> wordshift*SHIFT) + 2**remshift) >> remshift
-        */
+    a_negative = Py_SIZE(a) < 0;
+    size_a = Py_ABS(Py_SIZE(a));
 
-        /* It's convenient for remshift to be positive below. Note that
-           we dealt with the case remshift == wordshift == 0 earlier. */
+    if (a_negative) {
+        /* For shifting negative integers, it's convenient to adjust so that
+           0 < remshift <= PyLong_SHIFT. */
         if (remshift == 0) {
             remshift = PyLong_SHIFT;
             wordshift -= 1;
         }
         assert(wordshift >= 0);
+    }
 
-        newsize = -Py_SIZE(a) - wordshift;
-        if (newsize <= 0) {
-            return PyLong_FromLong(-1);
-        }
-        hishift = PyLong_SHIFT - remshift;
-        z = _PyLong_New(newsize);
-        if (z == NULL) {
-            return NULL;
-        }
+    newsize = size_a - wordshift;
+    if (newsize <= 0) {
+        return PyLong_FromLong(-a_negative);
+    }
+    z = _PyLong_New(newsize);
+    if (z == NULL) {
+        return NULL;
+    }
+    hishift = PyLong_SHIFT - remshift;
+
+    if (a_negative) {
+        /*
+            For a positive integer a and nonnegative shift, we have:
+
+                (-a) >> shift == -((a + 2**shift - 1) >> shift).
+
+            In the addition `a + (2**shift - 1)`, the low `wordshift` digits of
+            `2**shift - 1` all have value `PyLong_MASK`, so we get a carry out
+            from the bottom `wordshift` digits when at least one of the least
+            significant `wordshift` digits of `a` is nonzero. Digit `wordshift`
+            of `2**shift - 1` has value `PyLong_MASK >> hishift`.
+        */
         Py_SET_SIZE(z, -newsize);
 
-        /* sticky is used to determine whether all dropped words are zero. */
-        sticky = 0;
+        digit sticky = 0;
         for (j = 0; j < wordshift; j++) {
             sticky |= a->ob_digit[j];
         }
-        accum = a->ob_digit[j++] + ((digit)1U << remshift) - (sticky == 0);
-        accum >>= remshift;
-        for (i = 0; j < -Py_SIZE(a); i++, j++) {
-            accum += (twodigits)a->ob_digit[j] << hishift;
-            z->ob_digit[i] = (digit)(accum & PyLong_MASK);
-            accum >>= PyLong_SHIFT;
-        }
-        assert(i == newsize - 1);
-        z->ob_digit[i] = (digit)accum;
+        accum = a->ob_digit[j++] + (PyLong_MASK >> hishift)
+              + (digit)(sticky != 0);
     }
     else {
-        newsize = Py_SIZE(a) - wordshift;
-        if (newsize <= 0)
-            return PyLong_FromLong(0);
-        hishift = PyLong_SHIFT - remshift;
-        z = _PyLong_New(newsize);
-        if (z == NULL)
-            return NULL;
         j = wordshift;
-        accum = a->ob_digit[j++] >> remshift;
-        for (i = 0; j < Py_SIZE(a); i++, j++) {
-            accum |= (twodigits)a->ob_digit[j] << hishift;
-            z->ob_digit[i] = (digit)(accum & PyLong_MASK);
-            accum >>= PyLong_SHIFT;
-        }
-        z->ob_digit[i] = (digit)accum;
+        accum = a->ob_digit[j++];
     }
+
+    accum >>= remshift;
+    assert(j == wordshift+1);
+    for (i = 0; j < size_a; i++, j++) {
+        accum += (twodigits)a->ob_digit[j] << hishift;
+        z->ob_digit[i] = (digit)(accum & PyLong_MASK);
+        accum >>= PyLong_SHIFT;
+    }
+    assert(i == newsize - 1);
+    assert(accum <= PyLong_MASK);
+    z->ob_digit[i] = (digit)accum;
     z = maybe_small_long(long_normalize(z));
     return (PyObject *)z;
 }
@@ -4597,6 +4584,9 @@ long_rshift(PyObject *a, PyObject *b)
         PyErr_SetString(PyExc_ValueError, "negative shift count");
         return NULL;
     }
+    if (Py_SIZE(b) == 0) {
+        return long_long(a);
+    }
     if (Py_SIZE(a) == 0) {
         return PyLong_FromLong(0);
     }
@@ -4613,9 +4603,14 @@ _PyLong_Rshift(PyObject *a, size_t shiftby)
     digit remshift;
 
     assert(PyLong_Check(a));
+
+    if (shiftby == 0) {
+        return long_long(a);
+    }
     if (Py_SIZE(a) == 0) {
         return PyLong_FromLong(0);
     }
+
     wordshift = shiftby / PyLong_SHIFT;
     remshift = shiftby % PyLong_SHIFT;
     return long_rshift1((PyLongObject *)a, wordshift, remshift);