bpo-39648: Expand math.gcd() and math.lcm() to handle multiple arguments. (GH-18604)

* bpo-39648: Expand math.gcd() and math.lcm() to handle multiple arguments.

* Simplify fast path.

* Difine lcm() without arguments returning 1.

* Apply suggestions from code review

Co-Authored-By: Mark Dickinson <dickinsm@gmail.com>

Co-authored-by: Mark Dickinson <dickinsm@gmail.com>

1 files changed, 109 insertions, 74 deletions

diff --git a/Modules/mathmodule.c b/Modules/mathmodule.c
index f74b7e1..77e325c 100644
--- a/Modules/mathmodule.c
+++ b/Modules/mathmodule.c
@@ -826,37 +826,125 @@ m_log10(double x)
 }
 
 
-/*[clinic input]
-math.gcd
+static PyObject *
+math_gcd(PyObject *module, PyObject * const *args, Py_ssize_t nargs)
+{
+    PyObject *res, *x;
+    Py_ssize_t i;
 
-    x as a: object
-    y as b: object
-    /
+    if (nargs == 0) {
+        return PyLong_FromLong(0);
+    }
+    res = PyNumber_Index(args[0]);
+    if (res == NULL) {
+        return NULL;
+    }
+    if (nargs == 1) {
+        Py_SETREF(res, PyNumber_Absolute(res));
+        return res;
+    }
+    for (i = 1; i < nargs; i++) {
+        x = PyNumber_Index(args[i]);
+        if (x == NULL) {
+            Py_DECREF(res);
+            return NULL;
+        }
+        if (res == _PyLong_One) {
+            /* Fast path: just check arguments.
+               It is okay to use identity comparison here. */
+            Py_DECREF(x);
+            continue;
+        }
+        Py_SETREF(res, _PyLong_GCD(res, x));
+        Py_DECREF(x);
+        if (res == NULL) {
+            return NULL;
+        }
+    }
+    return res;
+}
+
+PyDoc_STRVAR(math_gcd_doc,
+"gcd($module, *integers)\n"
+"--\n"
+"\n"
+"Greatest Common Divisor.");
 
-greatest common divisor of x and y
-[clinic start generated code]*/
 
 static PyObject *
-math_gcd_impl(PyObject *module, PyObject *a, PyObject *b)
-/*[clinic end generated code: output=7b2e0c151bd7a5d8 input=c2691e57fb2a98fa]*/
+long_lcm(PyObject *a, PyObject *b)
 {
-    PyObject *g;
+    PyObject *g, *m, *f, *ab;
 
-    a = PyNumber_Index(a);
-    if (a == NULL)
+    if (Py_SIZE(a) == 0 || Py_SIZE(b) == 0) {
+        return PyLong_FromLong(0);
+    }
+    g = _PyLong_GCD(a, b);
+    if (g == NULL) {
         return NULL;
-    b = PyNumber_Index(b);
-    if (b == NULL) {
-        Py_DECREF(a);
+    }
+    f = PyNumber_FloorDivide(a, g);
+    Py_DECREF(g);
+    if (f == NULL) {
         return NULL;
     }
-    g = _PyLong_GCD(a, b);
-    Py_DECREF(a);
-    Py_DECREF(b);
-    return g;
+    m = PyNumber_Multiply(f, b);
+    Py_DECREF(f);
+    if (m == NULL) {
+        return NULL;
+    }
+    ab = PyNumber_Absolute(m);
+    Py_DECREF(m);
+    return ab;
 }
 
 
+static PyObject *
+math_lcm(PyObject *module, PyObject * const *args, Py_ssize_t nargs)
+{
+    PyObject *res, *x;
+    Py_ssize_t i;
+
+    if (nargs == 0) {
+        return PyLong_FromLong(1);
+    }
+    res = PyNumber_Index(args[0]);
+    if (res == NULL) {
+        return NULL;
+    }
+    if (nargs == 1) {
+        Py_SETREF(res, PyNumber_Absolute(res));
+        return res;
+    }
+    for (i = 1; i < nargs; i++) {
+        x = PyNumber_Index(args[i]);
+        if (x == NULL) {
+            Py_DECREF(res);
+            return NULL;
+        }
+        if (res == _PyLong_Zero) {
+            /* Fast path: just check arguments.
+               It is okay to use identity comparison here. */
+            Py_DECREF(x);
+            continue;
+        }
+        Py_SETREF(res, long_lcm(res, x));
+        Py_DECREF(x);
+        if (res == NULL) {
+            return NULL;
+        }
+    }
+    return res;
+}
+
+
+PyDoc_STRVAR(math_lcm_doc,
+"lcm($module, *integers)\n"
+"--\n"
+"\n"
+"Least Common Multiple.");
+
+
 /* Call is_error when errno != 0, and where x is the result libm
  * returned.  is_error will usually set up an exception and return
  * true (1), but may return false (0) without setting up an exception.
@@ -2018,59 +2106,6 @@ math_factorial(PyObject *module, PyObject *arg)
 
 
 /*[clinic input]
-math.lcm
-    x as a: object
-    y as b: object
-    /
-least common multiple of x and y
-[clinic start generated code]*/
-
-static PyObject *
-math_lcm_impl(PyObject *module, PyObject *a, PyObject *b)
-/*[clinic end generated code: output=6f83fb6d671074ba input=efb3d7b7334b7118]*/
-{
-    PyObject *g, *m, *f, *ab;
-
-    a = PyNumber_Index(a);
-    if (a == NULL) {
-        return NULL;
-    }
-    b = PyNumber_Index(b);
-    if (b == NULL) {
-        Py_DECREF(a);
-        return NULL;
-    }
-    if (_PyLong_Sign(a) == 0 || _PyLong_Sign(b) == 0) {
-        Py_DECREF(a);
-        Py_DECREF(b);
-        return PyLong_FromLong(0);
-    }
-    g = _PyLong_GCD(a, b);
-    if (g == NULL) {
-        Py_DECREF(a);
-        Py_DECREF(b);
-        return NULL;
-    }
-    f = PyNumber_FloorDivide(a, g);
-    Py_DECREF(g);
-    Py_DECREF(a);
-    if (f == NULL) {
-        Py_DECREF(b);
-        return NULL;
-    }
-    m = PyNumber_Multiply(f, b);
-    Py_DECREF(f);
-    Py_DECREF(b);
-    if (m == NULL) {
-        return NULL;
-    }
-    ab = PyNumber_Absolute(m);
-    Py_DECREF(m);
-    return ab;
-}
-
-
-/*[clinic input]
 math.trunc
 
     x: object
@@ -3408,14 +3443,14 @@ static PyMethodDef math_methods[] = {
     MATH_FREXP_METHODDEF
     MATH_FSUM_METHODDEF
     {"gamma",           math_gamma,     METH_O,         math_gamma_doc},
-    MATH_GCD_METHODDEF
+    {"gcd",             (PyCFunction)(void(*)(void))math_gcd,       METH_FASTCALL,  math_gcd_doc},
     {"hypot",           (PyCFunction)(void(*)(void))math_hypot,     METH_FASTCALL,  math_hypot_doc},
     MATH_ISCLOSE_METHODDEF
     MATH_ISFINITE_METHODDEF
     MATH_ISINF_METHODDEF
     MATH_ISNAN_METHODDEF
     MATH_ISQRT_METHODDEF
-    MATH_LCM_METHODDEF
+    {"lcm",             (PyCFunction)(void(*)(void))math_lcm,       METH_FASTCALL,  math_lcm_doc},
     MATH_LDEXP_METHODDEF
     {"lgamma",          math_lgamma,    METH_O,         math_lgamma_doc},
     MATH_LOG_METHODDEF