bpo-44376 - reduce pow() overhead for small exponents (GH-26662)

Greatly reduce pow() overhead for small exponents.

2 files changed, 46 insertions, 5 deletions

diff --git a/Misc/NEWS.d/next/Core and Builtins/2021-06-11-17-37-15.bpo-44376.zhM1UW.rst b/Misc/NEWS.d/next/Core and Builtins/2021-06-11-17-37-15.bpo-44376.zhM1UW.rst
new file mode 100644
index 0000000..f854d56
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2021-06-11-17-37-15.bpo-44376.zhM1UW.rst
@@ -0,0 +1 @@
+Exact integer exponentiation (like ``i**2`` or ``pow(i, 2)``) with a small exponent is much faster, due to reducing overhead in such cases.
\ No newline at end of file
diff --git a/Objects/longobject.c b/Objects/longobject.c
index e1c1191..5e29e9a 100644
--- a/Objects/longobject.c
+++ b/Objects/longobject.c
@@ -4239,17 +4239,57 @@ long_pow(PyObject *v, PyObject *w, PyObject *x)
         REDUCE(result);                         \
     } while(0)
 
-    if (Py_SIZE(b) <= FIVEARY_CUTOFF) {
+    i = Py_SIZE(b);
+    digit bi = i ? b->ob_digit[i-1] : 0;
+    digit bit;
+    if (i <= 1 && bi <= 3) {
+        /* aim for minimal overhead */
+        if (bi >= 2) {
+            MULT(a, a, z);
+            if (bi == 3) {
+                MULT(z, a, z);
+            }
+        }
+        else if (bi == 1) {
+            /* Multiplying by 1 serves two purposes: if `a` is of an int
+             * subclass, makes the result an int (e.g., pow(False, 1) returns
+             * 0 instead of False), and potentially reduces `a` by the modulus.
+             */
+            MULT(a, z, z);
+        }
+        /* else bi is 0, and z==1 is correct */
+    }
+    else if (i <= FIVEARY_CUTOFF) {
         /* Left-to-right binary exponentiation (HAC Algorithm 14.79) */
         /* http://www.cacr.math.uwaterloo.ca/hac/about/chap14.pdf    */
-        for (i = Py_SIZE(b) - 1; i >= 0; --i) {
-            digit bi = b->ob_digit[i];
 
-            for (j = (digit)1 << (PyLong_SHIFT-1); j != 0; j >>= 1) {
+        /* Find the first significant exponent bit. Search right to left
+         * because we're primarily trying to cut overhead for small powers.
+         */
+        assert(bi);  /* else there is no significant bit */
+        Py_INCREF(a);
+        Py_DECREF(z);
+        z = a;
+        for (bit = 2; ; bit <<= 1) {
+            if (bit > bi) { /* found the first bit */
+                assert((bi & bit) == 0);
+                bit >>= 1;
+                assert(bi & bit);
+                break;
+            }
+        }
+        for (--i, bit >>= 1;;) {
+            for (; bit != 0; bit >>= 1) {
                 MULT(z, z, z);
-                if (bi & j)
+                if (bi & bit) {
                     MULT(z, a, z);
+                }
+            }
+            if (--i < 0) {
+                break;
             }
+            bi = b->ob_digit[i];
+            bit = (digit)1 << (PyLong_SHIFT-1);
         }
     }
     else {