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>
diff --git a/Modules/clinic/mathmodule.c.h b/Modules/clinic/mathmodule.c.h
index df45a1a..65f3dd4 100644
--- a/Modules/clinic/mathmodule.c.h
+++ b/Modules/clinic/mathmodule.c.h
@@ -2,36 +2,6 @@
preserve
[clinic start generated code]*/
-PyDoc_STRVAR(math_gcd__doc__,
-"gcd($module, x, y, /)\n"
-"--\n"
-"\n"
-"greatest common divisor of x and y");
-
-#define MATH_GCD_METHODDEF \
- {"gcd", (PyCFunction)(void(*)(void))math_gcd, METH_FASTCALL, math_gcd__doc__},
-
-static PyObject *
-math_gcd_impl(PyObject *module, PyObject *a, PyObject *b);
-
-static PyObject *
-math_gcd(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
-{
- PyObject *return_value = NULL;
- PyObject *a;
- PyObject *b;
-
- if (!_PyArg_CheckPositional("gcd", nargs, 2, 2)) {
- goto exit;
- }
- a = args[0];
- b = args[1];
- return_value = math_gcd_impl(module, a, b);
-
-exit:
- return return_value;
-}
-
PyDoc_STRVAR(math_ceil__doc__,
"ceil($module, x, /)\n"
"--\n"
@@ -85,36 +55,6 @@
#define MATH_FACTORIAL_METHODDEF \
{"factorial", (PyCFunction)math_factorial, METH_O, math_factorial__doc__},
-PyDoc_STRVAR(math_lcm__doc__,
-"lcm($module, x, y, /)\n"
-"--\n"
-"\n"
-"least common multiple of x and y");
-
-#define MATH_LCM_METHODDEF \
- {"lcm", (PyCFunction)(void(*)(void))math_lcm, METH_FASTCALL, math_lcm__doc__},
-
-static PyObject *
-math_lcm_impl(PyObject *module, PyObject *a, PyObject *b);
-
-static PyObject *
-math_lcm(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
-{
- PyObject *return_value = NULL;
- PyObject *a;
- PyObject *b;
-
- if (!_PyArg_CheckPositional("lcm", nargs, 2, 2)) {
- goto exit;
- }
- a = args[0];
- b = args[1];
- return_value = math_lcm_impl(module, a, b);
-
-exit:
- return return_value;
-}
-
PyDoc_STRVAR(math_trunc__doc__,
"trunc($module, x, /)\n"
"--\n"
@@ -925,4 +865,4 @@
exit:
return return_value;
}
-/*[clinic end generated code: output=f8daa185c043a7b7 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=1eae2b3ef19568fa input=a9049054013a1b77]*/
diff --git a/Modules/mathmodule.c b/Modules/mathmodule.c
index f74b7e1..77e325c 100644
--- a/Modules/mathmodule.c
+++ b/Modules/mathmodule.c
@@ -826,36 +826,124 @@
}
-/*[clinic input]
-math.gcd
-
- x as a: object
- y as b: object
- /
-
-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]*/
+math_gcd(PyObject *module, PyObject * const *args, Py_ssize_t nargs)
{
- PyObject *g;
+ PyObject *res, *x;
+ Py_ssize_t i;
- a = PyNumber_Index(a);
- if (a == NULL)
- return NULL;
- b = PyNumber_Index(b);
- if (b == NULL) {
- Py_DECREF(a);
+ if (nargs == 0) {
+ return PyLong_FromLong(0);
+ }
+ res = PyNumber_Index(args[0]);
+ if (res == NULL) {
return NULL;
}
- g = _PyLong_GCD(a, b);
- Py_DECREF(a);
- Py_DECREF(b);
- return g;
+ 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.");
+
+
+static PyObject *
+long_lcm(PyObject *a, PyObject *b)
+{
+ PyObject *g, *m, *f, *ab;
+
+ if (Py_SIZE(a) == 0 || Py_SIZE(b) == 0) {
+ return PyLong_FromLong(0);
+ }
+ g = _PyLong_GCD(a, b);
+ if (g == NULL) {
+ return NULL;
+ }
+ f = PyNumber_FloorDivide(a, g);
+ Py_DECREF(g);
+ if (f == NULL) {
+ return NULL;
+ }
+ 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
@@ -2018,59 +2106,6 @@
/*[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 @@
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