bpo-46407: Optimizing power-of-2 operations

Misc/NEWS.d/next/Core and Builtins/2022-01-17-07-34-16.bpo-46407.00aN9w.rst

@@ -0,0 +1 @@
+Optimized power-of-2 operations in ``Objects/longobject.c``. Patch by Jeremiah Vivian.

Objects/longobject.c

@@ -159,6 +159,34 @@ _PyLong_New(Py_ssize_t size)
     return result;
 }
 
+/* Like _PyLong_New, but it uses `PyObject_Calloc` instead */
+PyLongObject *
+_PyLong_NewFillZero(Py_ssize_t size)
+{
+    PyLongObject *result;
+    if (size > (Py_ssize_t)MAX_LONG_DIGITS) {
+        PyErr_SetString(PyExc_OverflowError,
+                        "too many digits in integer");
+        return NULL;
+    }
+    /* Fast operations for single digit integers (including zero)
+     * assume that there is always at least one digit present. */
+    Py_ssize_t ndigits = size ? size : 1;
+    /* Number of bytes needed is: offsetof(PyLongObject, ob_digit) +
+       sizeof(digit)*size.  Previous incarnations of this code used
+       sizeof(PyVarObject) instead of the offsetof, but this risks being
+       incorrect in the presence of padding between the PyVarObject header
+       and the digits. */
+    result = PyObject_Calloc(offsetof(PyLongObject, ob_digit) +
+                             ndigits*sizeof(digit), 1);
+    if (!result) {
+        PyErr_NoMemory();
+        return NULL;
+    }
+    _PyObject_InitVar((PyVarObject*)result, &PyLong_Type, size);
+    return result;
+}
+
 PyObject *
 _PyLong_Copy(PyLongObject *src)
 {
@@ -4231,8 +4259,32 @@ long_pow(PyObject *v, PyObject *w, PyObject *x)
         c = (PyLongObject *)x;
         Py_INCREF(x);
     }
-    else if (x == Py_None)
+    else if (x == Py_None) {
         c = NULL;
+        if (IS_MEDIUM_VALUE(a) && a->ob_digit[0] == (digit)2) {
+            /* only compute positive powers of 2 */
+            if (Py_SIZE(b) > 0) {
+                /* b must fit a Py_ssize_t */
+                i = PyLong_AsSsize_t((PyObject *)b);
+                if (i == -1)
+                    goto Error;
+                z = _PyLong_NewFillZero((j = i / PyLong_SHIFT) + 1);
+                if (z == NULL)
+                    goto Error;
+                /* set the last digit of ob_digit to 1 << (i % PyLong_SHIFT) */
+                z->ob_digit[j] = 1 << (i - j*PyLong_SHIFT);
+                if (Py_SIZE(a) < 0)
+                    Py_SET_SIZE(z, (j + 1) * (-(i & 1) | 1));
+                z = maybe_small_long(z);
+                goto Done;
+            }
+            else if (Py_SIZE(b) == 0) {
+                z = (PyLongObject *)_PyLong_GetOne();
+                Py_INCREF(z);
+                goto Done;
+            }
+        }
+    }
     else {
         Py_DECREF(a);
         Py_DECREF(b);
@@ -4666,6 +4718,21 @@ long_lshift1(PyLongObject *a, Py_ssize_t wordshift, digit remshift)
         stwodigits x = m < 0 ? -(-m << remshift) : m << remshift;
         return _PyLong_FromSTwoDigits(x);
     }
+    if (IS_MEDIUM_VALUE(a) && a->ob_digit[0] == (digit)1) {
+        /* shifting is guaranteed positive, so there is
+           no need for filtering less than 0 shifting unlike
+           long_pow */
+        /* if statement for checking greater than 0 shifts */
+        if (wordshift || remshift) {
+            z = _PyLong_NewFillZero(wordshift + 1);
+            if (z == NULL)
+                return NULL;
+            z->ob_digit[wordshift] = 1 << remshift;
+            Py_SET_SIZE(z, (wordshift + 1) * Py_SIZE(a));
+            return (PyObject *)maybe_small_long(z);
+        }
+        return PyLong_FromSsize_t(Py_SIZE(a));
+    }
 
     oldsize = Py_ABS(Py_SIZE(a));
     newsize = oldsize + wordshift;