Merge pull request #248 from skirpichev/lcm/223

Add lcm()

Author: skirpichev

gmp.c

@@ -1713,7 +1713,7 @@ gmp_gcdext(PyObject *Py_UNUSED(module), PyObject *const *args,
     Py_XDECREF(x);
     Py_XDECREF(y);
     if (ret == ZZ_MEM) {
-        PyErr_NoMemory(); /* LCOV_EXCL_LINE */
+        return PyErr_NoMemory(); /* LCOV_EXCL_LINE */
     }
     PyObject *tup = PyTuple_Pack(3, g, s, t);
 
@@ -1730,6 +1730,37 @@ gmp_gcdext(PyObject *Py_UNUSED(module), PyObject *const *args,
     return NULL;
 }
 
+static PyObject *
+gmp_lcm(PyObject *Py_UNUSED(module), PyObject *const *args, Py_ssize_t nargs)
+{
+    MPZ_Object *res = MPZ_new(0);
+
+    if (!res || zz_from_i64(1, &res->z)) {
+        return PyErr_NoMemory(); /* LCOV_EXCL_LINE */
+    }
+    for (Py_ssize_t i = 0; i < nargs; i++) {
+        MPZ_Object *arg;
+
+        CHECK_OP_INT(arg, args[i]);
+        if (zz_cmp_i32(&res->z, 0) == ZZ_EQ) {
+            Py_DECREF(arg);
+            continue;
+        }
+        if (zz_lcm(&res->z, &arg->z, &res->z)) {
+            /* LCOV_EXCL_START */
+            Py_DECREF(res);
+            Py_DECREF(arg);
+            return PyErr_NoMemory();
+            /* LCOV_EXCL_STOP */
+        }
+        Py_DECREF(arg);
+    }
+    return (PyObject *)res;
+end:
+    Py_DECREF(res);
+    return NULL;
+}
+
 static PyObject *
 gmp_isqrt(PyObject *Py_UNUSED(module), PyObject *arg)
 {
@@ -2071,6 +2102,9 @@ static PyMethodDef gmp_functions[] = {
     {"gcdext", (PyCFunction)gmp_gcdext, METH_FASTCALL,
      ("gcdext($module, x, y, /)\n--\n\n"
       "Compute extended GCD.")},
+    {"lcm", (PyCFunction)gmp_lcm, METH_FASTCALL,
+     ("lcm($module, /, *integers)\n--\n\n"
+      "Least Common Multiple.")},
     {"isqrt", gmp_isqrt, METH_O,
      ("isqrt($module, n, /)\n--\n\n"
       "Return the integer part of the square root of n.")},
@@ -2182,7 +2216,7 @@ gmp_exec(PyObject *m)
                        "numbers.Integral.register(gmp.mpz)\n"
                        "gmp.fac = gmp.factorial\n"
                        "gmp.__all__ = ['comb', 'factorial', 'gcd', 'isqrt',\n"
-                       "               'mpz']\n"
+                       "               'lcm', 'mpz']\n"
                        "gmp.__version__ = imp.version('python-gmp')\n");
 
     PyObject *res = PyRun_String(str, Py_file_input, ns, ns);

tests/test_functions.py

@@ -14,6 +14,7 @@
     gcdext,
     isqrt,
     isqrt_rem,
+    lcm,
     mpz,
 )
 from hypothesis import example, given
@@ -77,6 +78,17 @@ def test_gcdext_binary(x, y, c):
         assert fm(x, y) == r
 
 
+@given(bigints(), bigints())
+def test_lcm_binary(x, y):
+    mx = mpz(x)
+    my = mpz(y)
+    r = math.lcm(x, y)
+    assert lcm(mx, my) == r
+    assert lcm(mx, y) == r
+    assert lcm(x, my) == r
+    assert lcm(x, y) == r
+
+
 @given(lists(bigints(), max_size=6), bigints())
 @example([], 1)
 @example([2, 3, 4], 1)
@@ -90,6 +102,15 @@ def test_gcd_nary(xs, c):
     assert gcd(*xs) == r
 
 
+@given(lists(bigints(), max_size=6))
+@example([])
+def test_lcm_nary(xs):
+    mxs = list(map(mpz, xs))
+    r = math.lcm(*xs)
+    assert lcm(*mxs) == r
+    assert lcm(*xs) == r
+
+
 @given(booleans(), bigints(min_value=0), bigints(),
        integers(min_value=1, max_value=1<<30),
        sampled_from(["n", "f", "c", "u", "d"]))
@@ -138,6 +159,10 @@ def test_interfaces():
         gcdext(2, 1j)
     with pytest.raises(TypeError):
         gcdext(2j, 2)
+    with pytest.raises(TypeError):
+        lcm(1j)
+    with pytest.raises(TypeError):
+        lcm(1, 1j)
     with pytest.raises(TypeError):
         isqrt(1j)
     with pytest.raises(TypeError):

zz.c

@@ -1951,6 +1951,26 @@ zz_inverse(const zz_t *u, const zz_t *v, zz_t *w)
     return ZZ_VAL;
 }
 
+zz_err
+zz_lcm(const zz_t *u, const zz_t *v, zz_t *w)
+{
+    zz_t g;
+
+    if (zz_init(&g) || zz_gcd(u, v, &g)) {
+        /* LCOV_EXCL_START */
+err:
+        zz_clear(&g);
+        return ZZ_MEM;
+        /* LCOV_EXCL_STOP */
+    }
+    if (zz_div(u, &g, ZZ_RNDD, w, NULL) || zz_mul(w, v, w)) {
+        goto err; /* LCOV_EXCL_LINE */
+    }
+    zz_clear(&g);
+    (void)zz_abs(w, w);
+    return ZZ_OK;
+}
+
 zz_err
 zz_powm(const zz_t *u, const zz_t *v, const zz_t *w, zz_t *res)
 {

zz.h

@@ -122,6 +122,7 @@ zz_err zz_pow(const zz_t *u, uint64_t v, zz_t *w);
 zz_err zz_gcd(const zz_t *u, const zz_t *v, zz_t *w);
 zz_err zz_gcdext(const zz_t *u, const zz_t *v, zz_t *g, zz_t *s, zz_t *t);
 zz_err zz_inverse(const zz_t *u, const zz_t *v, zz_t *w);
+zz_err zz_lcm(const zz_t *u, const zz_t *v, zz_t *w);
 
 zz_err zz_powm(const zz_t *u, const zz_t *v, const zz_t *w, zz_t *res);