diff --git a/extra.c b/extra.c index 7cab482f..767c525a 100644 --- a/extra.c +++ b/extra.c @@ -371,3 +371,31 @@ zz_mpmath_normalize(zz_bitcnt_t prec, zz_rnd rnd, bool *negative, } return ZZ_OK; } + +#define MAKE_MIX_LOGIC(func, op) \ + zz_err \ + zz_##func##_##i64(const zz_t *u, int64_t v, zz_t *w) \ + { \ + int64_t iu; \ + \ + if (zz_get_##i64(u, &iu)) { \ + zz_t tmp; \ + \ + if (zz_init(&tmp) || zz_set_##i64(v, &tmp)) { \ + /* LCOV_EXCL_START */ \ + zz_clear(&tmp); \ + return ZZ_MEM; \ + /* LCOV_EXCL_STOP */ \ + } \ + \ + zz_err ret = zz_##func(u, &tmp, w); \ + \ + zz_clear(&tmp); \ + return ret; \ + } \ + return zz_set_##i64(iu op v, w); \ + } + +MAKE_MIX_LOGIC(and, &) +MAKE_MIX_LOGIC(or, |) +MAKE_MIX_LOGIC(xor, ^) diff --git a/extra.h b/extra.h index eba1b89e..ba4afaa1 100644 --- a/extra.h +++ b/extra.h @@ -30,4 +30,68 @@ typedef enum { 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_err zz_and_i64(const zz_t *u, int64_t v, zz_t *w); +zz_err zz_or_i64(const zz_t *u, int64_t v, zz_t *w); +zz_err zz_xor_i64(const zz_t *u, int64_t v, zz_t *w); + +static inline zz_err +zz_i64_and(int64_t u, const zz_t *v, zz_t *w) +{ + return zz_and_i64(v, u, w); +} + +#define zz_and(U, V, W) \ + _Generic((U), \ + int: _Generic((V), \ + default: zz_i64_and), \ + long: _Generic((V), \ + default: zz_i64_and), \ + long long: _Generic((V), \ + default: zz_i64_and), \ + default: _Generic((V), \ + int: zz_and_i64, \ + long: zz_and_i64, \ + long long: zz_and_i64, \ + default: zz_and))(U, V, W) + +static inline zz_err +zz_i64_or(int64_t u, const zz_t *v, zz_t *w) +{ + return zz_or_i64(v, u, w); +} + +#define zz_or(U, V, W) \ + _Generic((U), \ + int: _Generic((V), \ + default: zz_i64_or), \ + long: _Generic((V), \ + default: zz_i64_or), \ + long long: _Generic((V), \ + default: zz_i64_or), \ + default: _Generic((V), \ + int: zz_or_i64, \ + long: zz_or_i64, \ + long long: zz_or_i64, \ + default: zz_or))(U, V, W) + +static inline zz_err +zz_i64_xor(int64_t u, const zz_t *v, zz_t *w) +{ + return zz_xor_i64(v, u, w); +} + +#define zz_xor(U, V, W) \ + _Generic((U), \ + int: _Generic((V), \ + default: zz_i64_xor), \ + long: _Generic((V), \ + default: zz_i64_xor), \ + long long: _Generic((V), \ + default: zz_i64_xor), \ + default: _Generic((V), \ + int: zz_xor_i64, \ + long: zz_xor_i64, \ + long long: zz_xor_i64, \ + default: zz_xor))(U, V, W) + #endif /* EXTRA_H */ diff --git a/gmp.c b/gmp.c index 57b24bbb..7a829e97 100644 --- a/gmp.c +++ b/gmp.c @@ -817,14 +817,16 @@ to_bool(PyObject *self) \ zz_err ret = ZZ_OK; \ \ - if (!u) { \ - ret = zz_##suff(ival, &v->z, &res->z); \ - } \ - else if (!v) { \ - ret = zz_##suff(&u->z, ival, &res->z); \ + if (u) { \ + if (v) { \ + ret = zz_##suff(&u->z, &v->z, &res->z); \ + } \ + else { \ + ret = zz_##suff(&u->z, ival, &res->z); \ + } \ } \ else { \ - ret = zz_##suff(&u->z, &v->z, &res->z); \ + ret = zz_##suff(ival, &v->z, &res->z); \ } \ if (ret) { \ Py_CLEAR(res); \ @@ -898,6 +900,23 @@ BINOP(rem, PyNumber_Remainder) } \ } \ +#define CHECK_OP_INTv2(u, a) \ + if (MPZ_Check(a)) { \ + u = (MPZ_Object *)Py_NewRef(a); \ + } \ + else { \ + int error; \ + \ + ival = PyLong_AsInt64AndOverflow(a, &error); \ + if (error) { \ + u = MPZ_from_int(a); \ + if (!u) { \ + goto end; \ + } \ + } \ + } \ + + Py_LOCAL(PyObject *) nb_divmod(PyObject *self, PyObject *other) { @@ -1016,6 +1035,55 @@ nb_truediv(PyObject *self, PyObject *other) return res; } +#define BINOP_INT_MIXED(suff) \ + Py_LOCAL(PyObject *) \ + nb_##suff(PyObject *self, PyObject *other) \ + { \ + MPZ_Object *u = NULL, *v = NULL, *res = NULL; \ + int64_t ival = 0; \ + \ + CHECK_OP_INTv2(u, self); \ + CHECK_OP_INTv2(v, other); \ + res = MPZ_new(); \ + if (!res) { \ + goto end; \ + } \ + \ + zz_err ret = ZZ_OK; \ + \ + if (u) { \ + if (v) { \ + ret = zz_##suff(&u->z, &v->z, &res->z); \ + } \ + else { \ + ret = zz_##suff(&u->z, ival, &res->z); \ + } \ + } \ + else { \ + ret = zz_##suff(ival, &v->z, &res->z); \ + } \ + if (ret) { \ + /* LCOV_EXCL_START */ \ + Py_CLEAR(res); \ + if (ret == ZZ_VAL) { \ + PyErr_SetString(PyExc_ValueError, \ + "negative shift count"); \ + } \ + else if (ret == ZZ_BUF) { \ + PyErr_SetString(PyExc_OverflowError, \ + "too many digits in integer"); \ + } \ + else { \ + PyErr_NoMemory(); \ + } \ + /* LCOV_EXCL_STOP */ \ + } \ + end: \ + Py_XDECREF((PyObject *)u); \ + Py_XDECREF((PyObject *)v); \ + return (PyObject *)res; \ + } + #define BINOP_INT(suff) \ Py_LOCAL(PyObject *) \ nb_##suff(PyObject *self, PyObject *other) \ @@ -1050,9 +1118,9 @@ nb_truediv(PyObject *self, PyObject *other) return (PyObject *)res; \ } -BINOP_INT(and) -BINOP_INT(or) -BINOP_INT(xor) +BINOP_INT_MIXED(and) +BINOP_INT_MIXED(or) +BINOP_INT_MIXED(xor) BINOP_INT(lshift) BINOP_INT(rshift)