Add comb()

Closes #224

Author: skirpichev

gmp.c

@@ -1834,6 +1834,52 @@ MAKE_MPZ_UI_FUN(fac)
 MAKE_MPZ_UI_FUN(fac2)
 MAKE_MPZ_UI_FUN(fib)
 
+static PyObject *
+gmp_comb(PyObject *self, PyObject *const *args, Py_ssize_t nargs)
+{
+    if (nargs != 2) {
+        PyErr_SetString(PyExc_TypeError, "two arguments required");
+        return NULL;
+    }
+
+    MPZ_Object *x, *y, *res = MPZ_new(0);
+
+    if (!res) {
+        return NULL; /* LCOV_EXCL_LINE */
+    }
+    CHECK_OP_INT(x, args[0]);
+    CHECK_OP_INT(y, args[1]);
+    if (zz_isneg(&x->z) || zz_isneg(&y->z)) {
+        PyErr_SetString(PyExc_ValueError,
+                        "comb() not defined for negative values");
+        goto err;
+    }
+
+    int64_t n, k;
+
+    if ((zz_to_i64(&x->z, &n) || n > ULONG_MAX)
+        || (zz_to_i64(&y->z, &k) || k > ULONG_MAX))
+    {
+        PyErr_Format(PyExc_OverflowError,
+                     "comb() arguments should not exceed %ld",
+                     ULONG_MAX);
+        goto err;
+    }
+    Py_XDECREF(x);
+    Py_XDECREF(y);
+    if (zz_bin((uint64_t)n, (uint64_t)k, &res->z)) {
+        /* LCOV_EXCL_START */
+        PyErr_NoMemory();
+        goto err;
+        /* LCOV_EXCL_STOP */
+    }
+    return (PyObject *)res;
+err:
+end:
+    Py_DECREF(res);
+    return NULL;
+}
+
 static zz_rnd
 get_round_mode(PyObject *rndstr)
 {
@@ -2040,6 +2086,10 @@ static PyMethodDef gmp_functions[] = {
     {"fib", gmp_fib, METH_O,
      ("fib($module, n, /)\n--\n\n"
       "Return the n-th Fibonacci number.")},
+    {"comb", (PyCFunction)gmp_comb, METH_FASTCALL,
+     ("comb($module, n, k, /)\n--\n\nNumber of ways to choose k"
+      " items from n items without repetition and order.\n\n"
+      "Also called the binomial coefficient.")},
     {"_mpmath_normalize", (PyCFunction)gmp__mpmath_normalize, METH_FASTCALL,
      NULL},
     {"_mpmath_create", (PyCFunction)gmp__mpmath_create, METH_FASTCALL, NULL},
@@ -2131,7 +2181,8 @@ gmp_exec(PyObject *m)
     const char *str = ("import numbers, importlib.metadata as imp\n"
                        "numbers.Integral.register(gmp.mpz)\n"
                        "gmp.fac = gmp.factorial\n"
-                       "gmp.__all__ = ['factorial', 'gcd', 'isqrt', 'mpz']\n"
+                       "gmp.__all__ = ['comb', 'factorial', 'gcd', 'isqrt',\n"
+                       "               'mpz']\n"
                        "gmp.__version__ = imp.version('python-gmp')\n");
 
     PyObject *res = PyRun_String(str, Py_file_input, ns, ns);

tests/test_functions.py

@@ -5,6 +5,7 @@
 from gmp import (
     _mpmath_create,
     _mpmath_normalize,
+    comb,
     double_fac,
     fac,
     factorial,
@@ -45,6 +46,18 @@ def test_factorials(x):
         assert fm(x) == r
 
 
+@given(integers(min_value=0, max_value=12345),
+       integers(min_value=0, max_value=12345))
+def test_comb(x, y):
+    mx = mpz(x)
+    my = mpz(y)
+    r = math.comb(x, y)
+    assert comb(mx, my) == r
+    assert comb(mx, y) == r
+    assert comb(x, my) == r
+    assert comb(x, y) == r
+
+
 @given(bigints(), bigints(), bigints())
 @example(1<<(67*2), 1<<65, 1)
 @example(123, 1<<70, 1)
@@ -139,6 +152,16 @@ def test_interfaces():
         fac(-1)
     with pytest.raises(OverflowError):
         fac(2**1000)
+    with pytest.raises(TypeError):
+        comb(123)
+    with pytest.raises(ValueError, match="not defined for negative values"):
+        comb(-1, 2)
+    with pytest.raises(ValueError, match="not defined for negative values"):
+        comb(2, -1)
+    with pytest.raises(OverflowError):
+        comb(2**1000, 1)
+    with pytest.raises(OverflowError):
+        comb(1, 2**1000)
     with pytest.raises(TypeError):
         _mpmath_create(1j)
     with pytest.raises(TypeError):

zz.c

@@ -2131,6 +2131,31 @@ MK_ZZ_FUNC_UL(fac, fac_ui)
 MK_ZZ_FUNC_UL(fac2, 2fac_ui)
 MK_ZZ_FUNC_UL(fib, fib_ui)
 
+zz_err
+zz_bin(uint64_t n, uint64_t k, zz_t *v)
+{
+    if (n > ULONG_MAX || k > ULONG_MAX) {
+        return ZZ_BUF;
+    }
+    if (TMP_OVERFLOW) {
+        return ZZ_MEM; /* LCOV_EXCL_LINE */
+    }
+
+    mpz_t z;
+
+    mpz_init(z);
+    mpz_bin_uiui(z, (unsigned long)n, (unsigned long)k);
+    if (zz_resize(z->_mp_size, v) == ZZ_MEM) {
+        /* LCOV_EXCL_START */
+        mpz_clear(z);
+        return ZZ_MEM;
+        /* LCOV_EXCL_STOP */
+    }
+    mpn_copyi(v->digits, z->_mp_d, z->_mp_size);
+    mpz_clear(z);
+    return ZZ_OK;
+}
+
 zz_err
 _zz_mpmath_normalize(zz_bitcnt_t prec, zz_rnd rnd, bool *negative,
                      zz_t *man, zz_t *exp, zz_bitcnt_t *bc)

zz.h

@@ -131,6 +131,8 @@ zz_err zz_fac(uint64_t u, zz_t *v);
 zz_err zz_fac2(uint64_t u, zz_t *v);
 zz_err zz_fib(uint64_t u, zz_t *v);
 
+zz_err zz_bin(uint64_t n, uint64_t k, zz_t *v);
+
 zz_err _zz_mpmath_normalize(zz_bitcnt_t prec, zz_rnd rnd, bool *negative,
                             zz_t *man, zz_t *exp, zz_bitcnt_t *bc);