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.

1 files changed, 33 insertions, 0 deletions

diff --git a/py/objint_mpz.c b/py/objint_mpz.c
index d465ef965..2b27df4f6 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);