bpo-36027: Extend three-argument pow to negative second argument (GH-13266)

6 files changed, 173 insertions, 16 deletions

diff --git a/Doc/library/functions.rst b/Doc/library/functions.rst
index d5c9f18..415a65b 100644
--- a/Doc/library/functions.rst
+++ b/Doc/library/functions.rst
@@ -1277,9 +1277,24 @@ are always available.  They are listed here in alphabetical order.
    operands, the result has the same type as the operands (after coercion)
    unless the second argument is negative; in that case, all arguments are
    converted to float and a float result is delivered.  For example, ``10**2``
-   returns ``100``, but ``10**-2`` returns ``0.01``.  If the second argument is
-   negative, the third argument must be omitted.  If *z* is present, *x* and *y*
-   must be of integer types, and *y* must be non-negative.
+   returns ``100``, but ``10**-2`` returns ``0.01``.
+
+   For :class:`int` operands *x* and *y*, if *z* is present, *z* must also be
+   of integer type and *z* must be nonzero. If *z* is present and *y* is
+   negative, *x* must be relatively prime to *z*. In that case, ``pow(inv_x,
+   -y, z)`` is returned, where *inv_x* is an inverse to *x* modulo *z*.
+
+   Here's an example of computing an inverse for ``38`` modulo ``97``::
+
+      >>> pow(38, -1, 97)
+      23
+      >>> 23 * 38 % 97 == 1
+      True
+
+   .. versionchanged:: 3.8
+      For :class:`int` operands, the three-argument form of ``pow`` now allows
+      the second argument to be negative, permitting computation of modular
+      inverses.
 
 
 .. function:: print(*objects, sep=' ', end='\\n', file=sys.stdout, flush=False)
diff --git a/Doc/whatsnew/3.8.rst b/Doc/whatsnew/3.8.rst
index 591b454..74d0079 100644
--- a/Doc/whatsnew/3.8.rst
+++ b/Doc/whatsnew/3.8.rst
@@ -304,6 +304,12 @@ Other Language Changes
 * Added new ``replace()`` method to the code type (:class:`types.CodeType`).
   (Contributed by Victor Stinner in :issue:`37032`.)
 
+* For integers, the three-argument form of the :func:`pow` function now permits
+  the exponent to be negative in the case where the base is relatively prime to
+  the modulus. It then computes a modular inverse to the base when the exponent
+  is ``-1``, and a suitable power of that inverse for other negative exponents.
+  (Contributed by Mark Dickinson in :issue:`36027`.)
+
 
 New Modules
 ===========
diff --git a/Lib/test/test_builtin.py b/Lib/test/test_builtin.py
index e32fb75..b536cec 100644
--- a/Lib/test/test_builtin.py
+++ b/Lib/test/test_builtin.py
@@ -1195,7 +1195,8 @@ class BuiltinTest(unittest.TestCase):
         self.assertAlmostEqual(pow(-1, 0.5), 1j)
         self.assertAlmostEqual(pow(-1, 1/3), 0.5 + 0.8660254037844386j)
 
-        self.assertRaises(ValueError, pow, -1, -2, 3)
+        # See test_pow for additional tests for three-argument pow.
+        self.assertEqual(pow(-1, -2, 3), 1)
         self.assertRaises(ValueError, pow, 1, 2, 0)
 
         self.assertRaises(TypeError, pow)
diff --git a/Lib/test/test_pow.py b/Lib/test/test_pow.py
index cac1ae5..660ff80 100644
--- a/Lib/test/test_pow.py
+++ b/Lib/test/test_pow.py
@@ -1,3 +1,4 @@
+import math
 import unittest
 
 class PowTest(unittest.TestCase):
@@ -119,5 +120,30 @@ class PowTest(unittest.TestCase):
             eq(pow(a, -fiveto), expected)
         eq(expected, 1.0)   # else we didn't push fiveto to evenness
 
+    def test_negative_exponent(self):
+        for a in range(-50, 50):
+            for m in range(-50, 50):
+                with self.subTest(a=a, m=m):
+                    if m != 0 and math.gcd(a, m) == 1:
+                        # Exponent -1 should give an inverse, with the
+                        # same sign as m.
+                        inv = pow(a, -1, m)
+                        self.assertEqual(inv, inv % m)
+                        self.assertEqual((inv * a - 1) % m, 0)
+
+                        # Larger exponents
+                        self.assertEqual(pow(a, -2, m), pow(inv, 2, m))
+                        self.assertEqual(pow(a, -3, m), pow(inv, 3, m))
+                        self.assertEqual(pow(a, -1001, m), pow(inv, 1001, m))
+
+                    else:
+                        with self.assertRaises(ValueError):
+                            pow(a, -1, m)
+                        with self.assertRaises(ValueError):
+                            pow(a, -2, m)
+                        with self.assertRaises(ValueError):
+                            pow(a, -1001, m)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/Misc/NEWS.d/next/Core and Builtins/2019-05-12-18-46-50.bpo-36027.Q4YatQ.rst b/Misc/NEWS.d/next/Core and Builtins/2019-05-12-18-46-50.bpo-36027.Q4YatQ.rst
new file mode 100644
index 0000000..866309c
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2019-05-12-18-46-50.bpo-36027.Q4YatQ.rst
@@ -0,0 +1,3 @@
+Allow computation of modular inverses via three-argument ``pow``: the second
+argument is now permitted to be negative in the case where the first and
+third arguments are relatively prime.
diff --git a/Objects/longobject.c b/Objects/longobject.c
index 5d2b595..49f1420 100644
--- a/Objects/longobject.c
+++ b/Objects/longobject.c
@@ -4174,6 +4174,98 @@ long_divmod(PyObject *a, PyObject *b)
     return z;
 }
 
+
+/* Compute an inverse to a modulo n, or raise ValueError if a is not
+   invertible modulo n. Assumes n is positive. The inverse returned
+   is whatever falls out of the extended Euclidean algorithm: it may
+   be either positive or negative, but will be smaller than n in
+   absolute value.
+
+   Pure Python equivalent for long_invmod:
+
+        def invmod(a, n):
+            b, c = 1, 0
+            while n:
+                q, r = divmod(a, n)
+                a, b, c, n = n, c, b - q*c, r
+
+            # at this point a is the gcd of the original inputs
+            if a == 1:
+                return b
+            raise ValueError("Not invertible")
+*/
+
+static PyLongObject *
+long_invmod(PyLongObject *a, PyLongObject *n)
+{
+    PyLongObject *b, *c;
+
+    /* Should only ever be called for positive n */
+    assert(Py_SIZE(n) > 0);
+
+    b = (PyLongObject *)PyLong_FromLong(1L);
+    if (b == NULL) {
+        return NULL;
+    }
+    c = (PyLongObject *)PyLong_FromLong(0L);
+    if (c == NULL) {
+        Py_DECREF(b);
+        return NULL;
+    }
+    Py_INCREF(a);
+    Py_INCREF(n);
+
+    /* references now owned: a, b, c, n */
+    while (Py_SIZE(n) != 0) {
+        PyLongObject *q, *r, *s, *t;
+
+        if (l_divmod(a, n, &q, &r) == -1) {
+            goto Error;
+        }
+        Py_DECREF(a);
+        a = n;
+        n = r;
+        t = (PyLongObject *)long_mul(q, c);
+        Py_DECREF(q);
+        if (t == NULL) {
+            goto Error;
+        }
+        s = (PyLongObject *)long_sub(b, t);
+        Py_DECREF(t);
+        if (s == NULL) {
+            goto Error;
+        }
+        Py_DECREF(b);
+        b = c;
+        c = s;
+    }
+    /* references now owned: a, b, c, n */
+
+    Py_DECREF(c);
+    Py_DECREF(n);
+    if (long_compare(a, _PyLong_One)) {
+        /* a != 1; we don't have an inverse. */
+        Py_DECREF(a);
+        Py_DECREF(b);
+        PyErr_SetString(PyExc_ValueError,
+                        "base is not invertible for the given modulus");
+        return NULL;
+    }
+    else {
+        /* a == 1; b gives an inverse modulo n */
+        Py_DECREF(a);
+        return b;
+    }
+
+  Error:
+    Py_DECREF(a);
+    Py_DECREF(b);
+    Py_DECREF(c);
+    Py_DECREF(n);
+    return NULL;
+}
+
+
 /* pow(v, w, x) */
 static PyObject *
 long_pow(PyObject *v, PyObject *w, PyObject *x)
@@ -4207,20 +4299,14 @@ long_pow(PyObject *v, PyObject *w, PyObject *x)
         Py_RETURN_NOTIMPLEMENTED;
     }
 
-    if (Py_SIZE(b) < 0) {  /* if exponent is negative */
-        if (c) {
-            PyErr_SetString(PyExc_ValueError, "pow() 2nd argument "
-                            "cannot be negative when 3rd argument specified");
-            goto Error;
-        }
-        else {
-            /* else return a float.  This works because we know
+    if (Py_SIZE(b) < 0 && c == NULL) {
+        /* if exponent is negative and there's no modulus:
+               return a float.  This works because we know
                that this calls float_pow() which converts its
                arguments to double. */
-            Py_DECREF(a);
-            Py_DECREF(b);
-            return PyFloat_Type.tp_as_number->nb_power(v, w, x);
-        }
+        Py_DECREF(a);
+        Py_DECREF(b);
+        return PyFloat_Type.tp_as_number->nb_power(v, w, x);
     }
 
     if (c) {
@@ -4255,6 +4341,26 @@ long_pow(PyObject *v, PyObject *w, PyObject *x)
             goto Done;
         }
 
+        /* if exponent is negative, negate the exponent and
+           replace the base with a modular inverse */
+        if (Py_SIZE(b) < 0) {
+            temp = (PyLongObject *)_PyLong_Copy(b);
+            if (temp == NULL)
+                goto Error;
+            Py_DECREF(b);
+            b = temp;
+            temp = NULL;
+            _PyLong_Negate(&b);
+            if (b == NULL)
+                goto Error;
+
+            temp = long_invmod(a, c);
+            if (temp == NULL)
+                goto Error;
+            Py_DECREF(a);
+            a = temp;
+        }
+
         /* Reduce base by modulus in some cases:
            1. If base < 0.  Forcing the base non-negative makes things easier.
            2. If base is obviously larger than the modulus.  The "small