add small_power

Author: cocolato

Lib/test/test_opcache.py

@@ -1536,6 +1536,44 @@ def binary_op_bitwise_extend():
         self.assert_specialized(binary_op_bitwise_extend, "BINARY_OP_EXTEND")
         self.assert_no_opcode(binary_op_bitwise_extend, "BINARY_OP")
 
+    @cpython_only
+    @requires_specialization
+    def test_binary_op_small_power(self):
+        def small_power():
+            for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
+                a = 7
+                self.assertEqual(a ** 2, 49)
+                b = -9
+                self.assertEqual(b ** 3, -729)
+                c = 2 ** 20 - 1  # largest base accepted by the exponent-3 bound
+                self.assertEqual(c ** 3, (2 ** 20 - 1) ** 3)
+                z = 0  # result is the immortal cached 0
+                self.assertEqual(z ** 2, 0)
+
+        small_power()
+        self.assert_specialized(small_power, "BINARY_OP_EXTEND")
+        self.assert_no_opcode(small_power, "BINARY_OP")
+
+        def small_power_inplace():
+            for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
+                a = 7
+                a **= 2
+                self.assertEqual(a, 49)
+
+        small_power_inplace()
+        self.assert_specialized(small_power_inplace, "BINARY_OP_EXTEND")
+
+        def power_not_extended(x):
+            for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
+                self.assertEqual(x ** 4, 2401)
+                self.assertEqual((x + 2 ** 20) ** 3, (7 + 2 ** 20) ** 3)
+                self.assertEqual((x + (1 << 40)) ** 2, (7 + (1 << 40)) ** 2)
+                self.assertEqual(bool(x) ** 2, True)
+                self.assertEqual(float(x) ** 3, 343.0)
+
+        power_not_extended(7)
+        self.assert_no_opcode(power_not_extended, "BINARY_OP_EXTEND")
+
     @cpython_only
     @requires_specialization
     def test_load_super_attr(self):

Python/specialize.c

@@ -2179,6 +2179,40 @@ BITWISE_LONGS_ACTION(compactlongs_and, &)
 BITWISE_LONGS_ACTION(compactlongs_xor, ^)
 #undef BITWISE_LONGS_ACTION
 
+static int
+small_power_guard(PyObject *lhs, PyObject *rhs)
+{
+    if (!is_compactlong(lhs) || !is_compactlong(rhs)) {
+        return false;
+    }
+    Py_ssize_t exponent = _PyLong_CompactValue((PyLongObject *)rhs);
+    if (exponent == 2) {
+        /* |x| < 2^30 (compact), so x*x < 2^60: fits a 64-bit
+           intermediate even where Py_ssize_t is 32 bits. */
+        return true;
+    }
+    if (exponent == 3) {
+        /* Need |x*x*x| < 2^60, so |x| < 2^20. */
+        Py_ssize_t x = _PyLong_CompactValue((PyLongObject *)lhs);
+        return x < (1 << 20) && x > -(1 << 20);
+    }
+    return false;
+}
+
+static PyObject *
+small_power(PyObject *lhs, PyObject *rhs)
+{
+    /* Compute in int64_t (not Py_ssize_t): the guard bounds keep the
+       product below 2^60, which overflows a 32-bit Py_ssize_t. */
+    int64_t x = (int64_t)_PyLong_CompactValue((PyLongObject *)lhs);
+    int64_t exponent = (int64_t)_PyLong_CompactValue((PyLongObject *)rhs);
+    assert(exponent == 2 || exponent == 3);
+    if (exponent == 2) {
+        return PyLong_FromLongLong(x * x);
+    }
+    return PyLong_FromLongLong(x * x * x);
+}
+
 /* float-long */
 
 static inline int
@@ -2259,6 +2293,11 @@ static _PyBinaryOpSpecializationDescr binaryop_extend_descrs[] = {
     {NB_INPLACE_AND, compactlongs_guard, compactlongs_and, &PyLong_Type, 1, NULL, NULL},
     {NB_INPLACE_XOR, compactlongs_guard, compactlongs_xor, &PyLong_Type, 1, NULL, NULL},
 
+    /* x ** 2 / x ** 3 for compact ints; PyLong_FromLongLong may return
+       the cached immortal small ints, which result_unique accepts. */
+    {NB_POWER, small_power_guard, small_power, &PyLong_Type, 1, NULL, NULL},
+    {NB_INPLACE_POWER, small_power_guard, small_power, &PyLong_Type, 1, NULL, NULL},
+
     /* float-long arithmetic: guards also check NaN and compactness. */
     {NB_ADD, float_compactlong_guard, float_compactlong_add, &PyFloat_Type, 1, NULL, NULL},
     {NB_SUBTRACT, float_compactlong_guard, float_compactlong_subtract, &PyFloat_Type, 1, NULL, NULL},