mruby-bigint/bigint.c: int.pow(n,m) to take bigint as exponential

Author: matz

include/mruby/internal.h

@@ -199,7 +199,7 @@ mrb_value mrb_bint_mul_ii(mrb_state *mrb, mrb_int x, mrb_int y);
 mrb_value mrb_bint_mod(mrb_state *mrb, mrb_value x, mrb_value y);
 mrb_value mrb_bint_rem(mrb_state *mrb, mrb_value x, mrb_value y);
 mrb_value mrb_bint_pow(mrb_state *mrb, mrb_value x, mrb_value y);
-mrb_value mrb_bint_powm(mrb_state *mrb, mrb_value x, mrb_int y, mrb_value z);
+mrb_value mrb_bint_powm(mrb_state *mrb, mrb_value x, mrb_value y, mrb_value z);
 mrb_value mrb_bint_and(mrb_state *mrb, mrb_value x, mrb_value y);
 mrb_value mrb_bint_or(mrb_state *mrb, mrb_value x, mrb_value y);
 mrb_value mrb_bint_xor(mrb_state *mrb, mrb_value x, mrb_value y);

mrbgems/mruby-bigint/core/bigint.c

@@ -942,7 +942,41 @@ mpz_pow(mrb_state *mrb, mpz_t *zz, mpz_t *x, mrb_int e)
 }
 
 static void
-mpz_powm(mrb_state *mrb, mpz_t *zz, mpz_t *x, mrb_int ex, mpz_t *n)
+mpz_powm(mrb_state *mrb, mpz_t *zz, mpz_t *x, mpz_t *ex, mpz_t *n)
+{
+  mpz_t t, b;
+
+  if (uzero(ex)) {
+    mpz_set_int(mrb, zz, 1);
+    return;
+  }
+
+  if (ex->sn < 0) {
+    return;
+  }
+
+  mpz_init_set_int(mrb, &t, 1);
+  mpz_init_set(mrb, &b, x);
+
+  size_t len = digits(ex);
+  for (size_t i=0; i<len; i++) {
+    mp_limb e = ex->p[i];
+    for (size_t j=0; j<sizeof(mp_limb)*8; j++) {
+      if ((e & 1) == 1) {
+        mpz_mul(mrb, &t, &t, &b);
+        mpz_mod(mrb, &t, &t, n);
+      }
+      e >>= 1;
+      mpz_mul(mrb, &b, &b, &b);
+      mpz_mod(mrb, &b, &b, n);
+    }
+  }
+  mpz_move(mrb, zz, &t);
+  mpz_clear(mrb, &b);
+}
+
+static void
+mpz_powm_i(mrb_state *mrb, mpz_t *zz, mpz_t *x, mrb_int ex, mpz_t *n)
 {
   mpz_t t, b;
 
@@ -1429,31 +1463,36 @@ mrb_bint_pow(mrb_state *mrb, mrb_value x, mrb_value y)
 }
 
 mrb_value
-mrb_bint_powm(mrb_state *mrb, mrb_value x, mrb_int exp, mrb_value mod)
+mrb_bint_powm(mrb_state *mrb, mrb_value x, mrb_value exp, mrb_value mod)
 {
   struct RBigint *b = RBIGINT(x);
-  switch (mrb_type(mod)) {
-  case MRB_TT_INTEGER:
-    {
-      mrb_int m = mrb_integer(mod);
-      if (m == 0) mrb_int_zerodiv(mrb);
-      struct RBigint *b2 = bint_new_int(mrb, m);
-      struct RBigint *b3 = bint_new(mrb);
-      mpz_powm(mrb, &b3->mp, &b->mp, exp, &b2->mp);
-      return bint_norm(mrb, b3);
+  struct RBigint *b2, *b3;
+
+  if (mrb_bigint_p(mod)) {
+    b2 = RBIGINT(mod);
+    if (uzero(&b2->mp)) mrb_int_zerodiv(mrb);
+  }
+  else {
+    mrb_int m = mrb_integer(mod);
+    if (m == 0) mrb_int_zerodiv(mrb);
+    b2 = bint_new_int(mrb, m);
+  }
+  b3 = bint_new(mrb);
+  if (mrb_bigint_p(exp)) {
+    struct RBigint *be = RBIGINT(exp);
+    if (be->mp.sn < 0) {
+      mrb_raise(mrb, E_ARGUMENT_ERROR, "int.pow(n,m): n must be positive");
     }
-  case MRB_TT_BIGINT:
-    {
-      struct RBigint *b2 = RBIGINT(mod);
-      struct RBigint *b3 = bint_new(mrb);
-      if (uzero(&b2->mp)) mrb_int_zerodiv(mrb);
-      mpz_powm(mrb, &b3->mp, &b->mp, exp, &b2->mp);
-      return bint_norm(mrb, b3);
+    mpz_powm(mrb, &b3->mp, &b->mp, &be->mp, &b2->mp);
+  }
+  else {
+    mrb_int e = mrb_integer(exp);
+    if (e < 0) {
+      mrb_raise(mrb, E_ARGUMENT_ERROR, "int.pow(n,m): n must be positive");
     }
-  default:
-    mrb_raisef(mrb, E_TYPE_ERROR, "%v cannot be convert to integer", mod);
+    mpz_powm_i(mrb, &b3->mp, &b->mp, e, &b2->mp);
   }
-  return mrb_nil_value();
+  return bint_norm(mrb, b3);
 }
 
 mrb_value

mrbgems/mruby-numeric-ext/src/numeric_ext.c

@@ -73,16 +73,16 @@ int_powm(mrb_state *mrb, mrb_value x)
       ) {
     mrb_raise(mrb, E_TYPE_ERROR, "int.pow(n,m): 2nd argument not allowed unless 1st argument is an integer");
   }
-  exp = mrb_as_int(mrb, e);
-  if (exp < 0) mrb_raise(mrb, E_ARGUMENT_ERROR, "int.pow(n,m): n must be positive");
 #ifdef MRB_USE_BIGINT
   if (mrb_bigint_p(x)) {
-    return mrb_bint_powm(mrb, x, exp, m);
+    return mrb_bint_powm(mrb, x, e, m);
   }
-  if (mrb_bigint_p(m)) {
-    return mrb_bint_powm(mrb, mrb_bint_new_int(mrb, mrb_integer(x)), exp, m);
+  if (mrb_bigint_p(e) || mrb_bigint_p(m)) {
+    return mrb_bint_powm(mrb, mrb_bint_new_int(mrb, mrb_integer(x)), e, m);
   }
 #endif
+  exp = mrb_integer(e);
+  if (exp < 0) mrb_raise(mrb, E_ARGUMENT_ERROR, "int.pow(n,m): n must be positive");
   if (!mrb_integer_p(m)) mrb_raise(mrb, E_TYPE_ERROR, "int.pow(n,m): m must be integer");
   mod = mrb_integer(m);
   if (mod < 0) mrb_raise(mrb, E_ARGUMENT_ERROR, "int.pow(n,m): m must be positive when 2nd argument specified");
@@ -96,7 +96,7 @@ int_powm(mrb_state *mrb, mrb_value x)
         result %= mod; base %= mod;
         if (mrb_int_mul_overflow(result, base, &tmp)) {
 #ifdef MRB_USE_BIGINT
-          return mrb_bint_powm(mrb, mrb_bint_new_int(mrb, mrb_integer(x)), mrb_integer(e), m);
+          return mrb_bint_powm(mrb, mrb_bint_new_int(mrb, mrb_integer(x)), e, m);
 #else
           mrb_int_overflow(mrb, "pow");
 #endif
@@ -110,7 +110,7 @@ int_powm(mrb_state *mrb, mrb_value x)
       base %= mod;
       if (mrb_int_mul_overflow(base, base, &tmp)) {
 #ifdef MRB_USE_BIGINT
-        return mrb_bint_powm(mrb, mrb_bint_new_int(mrb, mrb_integer(x)), mrb_integer(e), m);
+        return mrb_bint_powm(mrb, mrb_bint_new_int(mrb, mrb_integer(x)), e, m);
 #else
         mrb_int_overflow(mrb, "pow");
 #endif