From df0117c8ae213a0652c3b19a969edc7fd994eeab Mon Sep 17 00:00:00 2001 From: Nicko van Someren Date: Wed, 1 Feb 2017 16:41:22 -0700 Subject: [PATCH] py: Added optimised support for 3-argument calls to builtin.pow() Updated modbuiltin.c to add conditional support for 3-arg calls to pow() using MICROPY_PY_BUILTINS_POW3 config parameter. Added support in objint_mpz.c for for optimised implementation. --- py/modbuiltins.c | 9 ++++++++- py/mpconfig.h | 5 +++++ py/mpz.c | 6 +++--- py/mpz.h | 1 + py/objint.h | 1 + py/objint_mpz.c | 33 +++++++++++++++++++++++++++++++++ tests/basics/builtin_pow.py | 28 ++++++++++++++++++++++++++++ unix/mpconfigport.h | 1 + 8 files changed, 80 insertions(+), 4 deletions(-) diff --git a/py/modbuiltins.c b/py/modbuiltins.c index f62afd8074..a0c68930d8 100644 --- a/py/modbuiltins.c +++ b/py/modbuiltins.c @@ -378,7 +378,14 @@ MP_DEFINE_CONST_FUN_OBJ_1(mp_builtin_ord_obj, mp_builtin_ord); STATIC mp_obj_t mp_builtin_pow(size_t n_args, const mp_obj_t *args) { switch (n_args) { case 2: return mp_binary_op(MP_BINARY_OP_POWER, args[0], args[1]); - default: return mp_binary_op(MP_BINARY_OP_MODULO, mp_binary_op(MP_BINARY_OP_POWER, args[0], args[1]), args[2]); // TODO optimise... + default: +#if !MICROPY_PY_BUILTINS_POW3 + mp_raise_msg(&mp_type_NotImplementedError, "3-arg pow() not supported"); +#elif MICROPY_LONGINT_IMPL != MICROPY_LONGINT_IMPL_MPZ + return mp_binary_op(MP_BINARY_OP_MODULO, mp_binary_op(MP_BINARY_OP_POWER, args[0], args[1]), args[2]); +#else + return mp_obj_int_pow3(args[0], args[1], args[2]); +#endif } } MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mp_builtin_pow_obj, 2, 3, mp_builtin_pow); diff --git a/py/mpconfig.h b/py/mpconfig.h index 993ad1db8a..13af4c62bd 100644 --- a/py/mpconfig.h +++ b/py/mpconfig.h @@ -490,6 +490,11 @@ #define MICROPY_LONGINT_IMPL (MICROPY_LONGINT_IMPL_NONE) #endif +// Support for calls to pow() with 3 integer arguments +#ifndef MICROPY_PY_BUILTINS_POW3 +#define MICROPY_PY_BUILTINS_POW3 (0) +#endif + #if MICROPY_LONGINT_IMPL == MICROPY_LONGINT_IMPL_LONGLONG typedef long long mp_longint_impl_t; #endif diff --git a/py/mpz.c b/py/mpz.c index 6477c3f8d8..230eb921c6 100644 --- a/py/mpz.c +++ b/py/mpz.c @@ -1395,9 +1395,6 @@ void mpz_pow_inpl(mpz_t *dest, const mpz_t *lhs, const mpz_t *rhs) { mpz_free(n); } -#if 0 -these functions are unused - /* computes dest = (lhs ** rhs) % mod can have dest, lhs, rhs the same; mod can't be the same as dest */ @@ -1436,6 +1433,9 @@ void mpz_pow3_inpl(mpz_t *dest, const mpz_t *lhs, const mpz_t *rhs, const mpz_t mpz_free(n); } +#if 0 +these functions are unused + /* computes gcd(z1, z2) based on Knuth's modified gcd algorithm (I think?) gcd(z1, z2) >= 0 diff --git a/py/mpz.h b/py/mpz.h index a26cbea5cb..8facb1a0f3 100644 --- a/py/mpz.h +++ b/py/mpz.h @@ -123,6 +123,7 @@ void mpz_add_inpl(mpz_t *dest, const mpz_t *lhs, const mpz_t *rhs); void mpz_sub_inpl(mpz_t *dest, const mpz_t *lhs, const mpz_t *rhs); void mpz_mul_inpl(mpz_t *dest, const mpz_t *lhs, const mpz_t *rhs); void mpz_pow_inpl(mpz_t *dest, const mpz_t *lhs, const mpz_t *rhs); +void mpz_pow3_inpl(mpz_t *dest, const mpz_t *lhs, const mpz_t *rhs, const mpz_t *mod); void mpz_and_inpl(mpz_t *dest, const mpz_t *lhs, const mpz_t *rhs); void mpz_or_inpl(mpz_t *dest, const mpz_t *lhs, const mpz_t *rhs); void mpz_xor_inpl(mpz_t *dest, const mpz_t *lhs, const mpz_t *rhs); diff --git a/py/objint.h b/py/objint.h index a84a33fa50..7205761ad2 100644 --- a/py/objint.h +++ b/py/objint.h @@ -66,5 +66,6 @@ mp_obj_t mp_obj_int_abs(mp_obj_t self_in); mp_obj_t mp_obj_int_unary_op(mp_uint_t op, mp_obj_t o_in); mp_obj_t mp_obj_int_binary_op(mp_uint_t op, mp_obj_t lhs_in, mp_obj_t rhs_in); mp_obj_t mp_obj_int_binary_op_extra_cases(mp_uint_t op, mp_obj_t lhs_in, mp_obj_t rhs_in); +mp_obj_t mp_obj_int_pow3(mp_obj_t base, mp_obj_t exponent, mp_obj_t modulus); #endif // __MICROPY_INCLUDED_PY_OBJINT_H__ diff --git a/py/objint_mpz.c b/py/objint_mpz.c index d465ef9650..2b27df4f63 100644 --- a/py/objint_mpz.c +++ b/py/objint_mpz.c @@ -326,6 +326,39 @@ mp_obj_t mp_obj_int_binary_op(mp_uint_t op, mp_obj_t lhs_in, mp_obj_t rhs_in) { } } +#if MICROPY_PY_BUILTINS_POW3 +STATIC mpz_t *mp_mpz_for_int(mp_obj_t arg, mpz_t *temp) { + if (MP_OBJ_IS_SMALL_INT(arg)) { + mpz_init_from_int(temp, MP_OBJ_SMALL_INT_VALUE(arg)); + return temp; + } else { + mp_obj_int_t *arp_p = MP_OBJ_TO_PTR(arg); + return &(arp_p->mpz); + } +} + +mp_obj_t mp_obj_int_pow3(mp_obj_t base, mp_obj_t exponent, mp_obj_t modulus) { + if (!MP_OBJ_IS_INT(base) || !MP_OBJ_IS_INT(exponent) || !MP_OBJ_IS_INT(modulus)) { + mp_raise_TypeError("pow() with 3 arguments requires integers"); + } else { + mp_obj_t result = mp_obj_new_int_from_ull(0); // Use the _from_ull version as this forces an mpz int + mp_obj_int_t *res_p = (mp_obj_int_t *) MP_OBJ_TO_PTR(result); + + mpz_t l_temp, r_temp, m_temp; + mpz_t *lhs = mp_mpz_for_int(base, &l_temp); + mpz_t *rhs = mp_mpz_for_int(exponent, &r_temp); + mpz_t *mod = mp_mpz_for_int(modulus, &m_temp); + + mpz_pow3_inpl(&(res_p->mpz), lhs, rhs, mod); + + if (lhs == &l_temp) { mpz_deinit(lhs); } + if (rhs == &r_temp) { mpz_deinit(rhs); } + if (mod == &m_temp) { mpz_deinit(mod); } + return result; + } +} +#endif + mp_obj_t mp_obj_new_int(mp_int_t value) { if (MP_SMALL_INT_FITS(value)) { return MP_OBJ_NEW_SMALL_INT(value); diff --git a/tests/basics/builtin_pow.py b/tests/basics/builtin_pow.py index a19ab8c843..faf75f0df1 100644 --- a/tests/basics/builtin_pow.py +++ b/tests/basics/builtin_pow.py @@ -8,4 +8,32 @@ print(pow(3, 8)) # 3 arg version print(pow(3, 4, 7)) +print(pow(555557, 1000002, 1000003)) +# 3 arg pow is defined to only work on integers +try: + print(pow("x", 5, 6)) +except TypeError: + print("TypeError expected") + +try: + print(pow(4, "y", 6)) +except TypeError: + print("TypeError expected") + +try: + print(pow(4, 5, "z")) +except TypeError: + print("TypeError expected") + +# Tests for 3 arg pow with large values + +# This value happens to be prime +x = 0xd48a1e2a099b1395895527112937a391d02d4a208bce5d74b281cf35a57362502726f79a632f063a83c0eba66196712d963aa7279ab8a504110a668c0fc38a7983c51e6ee7a85cae87097686ccdc359ee4bbf2c583bce524e3f7836bded1c771a4efcb25c09460a862fc98e18f7303df46aaeb34da46b0c4d61d5cd78350f3edb60e6bc4befa712a849 +y = 0x3accf60bb1a5365e4250d1588eb0fe6cd81ad495e9063f90880229f2a625e98c59387238670936afb2cafc5b79448e4414d6cd5e9901aa845aa122db58ddd7b9f2b17414600a18c47494ed1f3d49d005a5 + +print(hex(pow(2, 200, x))) # Should not overflow, just 1 << 200 +print(hex(pow(2, x-1, x))) # Should be 1, since x is prime +print(hex(pow(y, x-1, x))) # Should be 1, since x is prime +print(hex(pow(y, y-1, x))) # Should be a 'big value' +print(hex(pow(y, y-1, y))) # Should be a 'big value' diff --git a/unix/mpconfigport.h b/unix/mpconfigport.h index ba2b5ce98d..66de0fa966 100644 --- a/unix/mpconfigport.h +++ b/unix/mpconfigport.h @@ -80,6 +80,7 @@ #define MICROPY_PY_BUILTINS_FROZENSET (1) #define MICROPY_PY_BUILTINS_COMPILE (1) #define MICROPY_PY_BUILTINS_NOTIMPLEMENTED (1) +#define MICROPY_PY_BUILTINS_POW3 (1) #define MICROPY_PY_MICROPYTHON_MEM_INFO (1) #define MICROPY_PY_ALL_SPECIAL_METHODS (1) #define MICROPY_PY_ARRAY_SLICE_ASSIGN (1)