Skip integer algorithm if the power base has special components

Author: skirpichev

Lib/test/test_complex.py

@@ -1,3 +1,4 @@
+import cmath
 import unittest
 import sys
 from test import support
@@ -7,7 +8,9 @@
 
 from random import random
 from math import isnan, copysign
+from itertools import combinations_with_replacement
 import operator
+import _testcapi
 
 INF = float("inf")
 NAN = float("nan")
@@ -412,11 +415,6 @@ def test_pow(self):
                     except OverflowError:
                         pass
 
-        # Check that complex numbers with special components
-        # are correctly handled.
-        self.assertComplexesAreIdentical(complex(1, +0.0)**2, complex(1, +0.0))
-        self.assertComplexesAreIdentical(complex(1, -0.0)**2, complex(1, -0.0))
-
         # gh-113841: possible undefined division by 0 in _Py_c_pow()
         x, y = 9j, 33j**3
         with self.assertRaises(OverflowError):
@@ -450,6 +448,36 @@ def test_pow_with_small_integer_exponents(self):
                     self.assertEqual(str(float_pow), str(int_pow))
                     self.assertEqual(str(complex_pow), str(int_pow))
 
+
+        # Check that complex numbers with special components
+        # are correctly handled.
+        values = [complex(*_)
+                  for _ in combinations_with_replacement([1, -1, 0.0, 0, -0.0, 2,
+                                                          -3, INF, -INF, NAN], 2)]
+        exponents = [0, 1, 2, 3, 4, 5, 6, 19]
+        for z in values:
+            for e in exponents:
+                with self.subTest(value=z, exponent=e):
+                    if cmath.isfinite(z) and z.real and z.imag:
+                        continue
+                    try:
+                        r_pow = z**e
+                    except OverflowError:
+                        continue
+                    # Use the generic complex power algorithm.
+                    r_pro, r_pro_errno = _testcapi._py_c_pow(z, e)
+                    self.assertEqual(r_pro_errno, 0)
+                    if isnan(r_pow.real):
+                        self.assertTrue(isnan(r_pro.real))
+                    else:
+                        self.assertEqual(copysign(1, r_pow.real),
+                                         copysign(1, r_pro.real))
+                    if isnan(r_pow.imag):
+                        self.assertTrue(isnan(r_pro.imag))
+                    else:
+                        self.assertEqual(copysign(1, r_pow.imag),
+                                         copysign(1, r_pro.imag))
+
     def test_boolcontext(self):
         for i in range(100):
             self.assertTrue(complex(random() + 1e-6, random() + 1e-6))

Objects/complexobject.c

@@ -26,6 +26,8 @@ class complex "PyComplexObject *" "&PyComplex_Type"
 
 /* elementary operations on complex numbers */
 
+static Py_complex c_1 = {1., 0.};
+
 Py_complex
 _Py_c_sum(Py_complex a, Py_complex b)
 {
@@ -336,6 +338,38 @@ _Py_c_pow(Py_complex a, Py_complex b)
     return r;
 }
 
+/* Switch to exponentiation by squaring if integer exponent less that this. */
+#define INT_EXP_CUTOFF 100
+
+static Py_complex
+c_powu(Py_complex x, long n)
+{
+    Py_complex r, p;
+    long mask = 1;
+    r = c_1;
+    p = x;
+    assert(0 <= n);
+    assert(n <= C_EXP_CUTOFF);
+    while (n >= mask) {
+        assert(mask > 0);
+        if (n & mask)
+            r = _Py_c_prod(r,p);
+        mask <<= 1;
+        p = _Py_c_prod(p,p);
+    }
+    return r;
+}
+
+static Py_complex
+c_powi(Py_complex x, long n)
+{
+    if (n > 0)
+        return c_powu(x,n);
+    else
+        return _Py_c_quot(c_1, c_powu(x,-n));
+
+}
+
 double
 _Py_c_abs(Py_complex z)
 {
@@ -705,7 +739,20 @@ complex_pow(PyObject *v, PyObject *w, PyObject *z)
         return NULL;
     }
     errno = 0;
-    p = _Py_c_pow(a, b);
+    // Check whether the exponent has a small integer value, and if so use
+    // a faster and more accurate algorithm.
+    // Fallback on the generic code if the base has special
+    // components (zeros or infinities).
+    if (b.imag == 0.0 && b.real == floor(b.real) && fabs(b.real) <= INT_EXP_CUTOFF
+        && isfinite(a.real) && a.real && isfinite(a.imag) && a.imag)
+    {
+        p = c_powi(a, (long)b.real);
+        _Py_ADJUST_ERANGE2(p.real, p.imag);
+    }
+    else {
+        p = _Py_c_pow(a, b);
+    }
+
     if (errno == EDOM) {
         PyErr_SetString(PyExc_ZeroDivisionError,
                         "zero to a negative or complex power");