Add Integer#pow. #4876

jruby · Mar 27, 2018 · 47328e0 · 47328e0
1 parent 35e67c6
commit 47328e0
Show file tree

Hide file tree

Showing 3 changed files with 164 additions and 1 deletion.
diff --git a/core/src/main/java/org/jruby/RubyFixnum.java b/core/src/main/java/org/jruby/RubyFixnum.java
@@ -59,6 +59,8 @@
 import org.jruby.util.Numeric;
 import org.jruby.util.cli.Options;
 
+import static org.jruby.util.Numeric.f_odd_p;
+
 /**
  * Implementation of the Fixnum class.
  */
@@ -873,6 +875,89 @@ private RubyNumeric powerFixnum(ThreadContext context, RubyFixnum other) {
         return Numeric.int_pow(context, a, b);
     }
 
+    // MRI: int_pow_tmp1
+    protected IRubyObject intPowTmp1(ThreadContext context, RubyInteger y, long mm, boolean negaFlg) {
+        Ruby runtime = context.runtime;
+
+        long xx = getLongValue();
+        long tmp = 1L;
+        long yy;
+
+        for (/*NOP*/; !(y instanceof RubyFixnum); y = (RubyInteger) sites(context).op_rshift.call(context, y, y, RubyFixnum.one(runtime))) {
+            if (f_odd_p(context, y)) {
+                tmp = (tmp * xx) % mm;
+            }
+            xx = (xx * xx) % mm;
+        }
+        for (yy = y.convertToInteger().getLongValue(); yy != 0L; yy >>= 1L) {
+            if ((yy & 1L) != 0L) {
+                tmp = (tmp * xx) % mm;
+            }
+            xx = (xx * xx) % mm;
+        }
+
+        if (negaFlg && (tmp != 0)) {
+            tmp -= mm;
+        }
+        return runtime.newFixnum(tmp);
+    }
+
+    // MRI: int_pow_tmp2
+    protected IRubyObject intPowTmp2(ThreadContext context, IRubyObject y, long mm, boolean negaFlg) {
+        Ruby runtime = context.runtime;
+
+        long tmp = 1L;
+        long yy;
+
+        final IRubyObject m = runtime.newFixnum(mm);
+        RubyFixnum tmp2 = runtime.newFixnum(tmp);
+        RubyFixnum xx = (RubyFixnum) this;
+
+        for (/*NOP*/; !(y instanceof RubyFixnum); y = sites(context).op_rshift.call(context, y, y, RubyFixnum.one(runtime))) {
+            if (f_odd_p(context, y)) {
+                tmp2 = mulModulo(context, tmp2, xx, m);
+            }
+            xx = mulModulo(context, xx, xx, m);
+        }
+        for (yy = ((RubyFixnum) y).getLongValue(); yy != 0; yy >>= 1L) {
+            if ((yy & 1L) != 0) {
+                tmp2 = mulModulo(context, tmp2, xx, m);
+            }
+            xx = mulModulo(context, xx, xx, m);
+        }
+
+        tmp = tmp2.getLongValue();
+        if (negaFlg && (tmp != 0)) {
+            tmp -= mm;
+        }
+        return runtime.newFixnum(tmp);
+    }
+
+    protected IRubyObject intPowTmp3(ThreadContext context, RubyInteger y, RubyBignum m, boolean negaFlg) {
+        Ruby runtime = context.runtime;
+
+        BigInteger xn, yn, mn, zn;
+
+        xn = BigInteger.valueOf(this.getLongValue());
+        if (y instanceof RubyFixnum) {
+            yn = BigInteger.valueOf(y.getLongValue());
+        } else {
+            yn = y.getBigIntegerValue();
+        }
+        mn = m.getBigIntegerValue();
+
+        zn = xn.modPow(yn, mn);
+        if (negaFlg & zn.signum() == 1) {
+            zn = zn.negate();
+        }
+        return RubyBignum.bignorm(runtime, zn);
+    }
+
+    // MRI: MUL_MODULO macro defined within int_pow_tmp2 in numeric.c
+    private static RubyFixnum mulModulo(ThreadContext context, RubyFixnum a, RubyFixnum b, IRubyObject c) {
+        return (RubyFixnum) ((RubyInteger) a.op_mul(context, b.getLongValue())).modulo(context, c);
+    }
+
     /** fix_abs
      *
      */

diff --git a/core/src/main/java/org/jruby/RubyInteger.java b/core/src/main/java/org/jruby/RubyInteger.java
@@ -752,6 +752,80 @@ public IRubyObject op_mod(ThreadContext context, long other) {
     @JRubyMethod(name = "**")
     public abstract IRubyObject op_pow(ThreadContext context, IRubyObject other);
 
+    @JRubyMethod(name = "pow")
+    public IRubyObject pow(ThreadContext context, IRubyObject other) {
+        return sites(context).op_pow.call(context, this, this, other);
+    }
+
+    private final long HALF_LONG_MSB = 0x80000000L;
+
+    // MRI: rb_int_powm
+    @JRubyMethod(name = "pow")
+    public IRubyObject pow(ThreadContext context, IRubyObject b, IRubyObject m) {
+        Ruby runtime = context.runtime;
+
+        RubyInteger a = this;
+
+        boolean negaFlg = false;
+        if (!(b instanceof RubyInteger)) {
+            throw runtime.newTypeError("Integer#pow() 2nd argument not allowed unless a 1st argument is integer");
+        }
+        RubyInteger intB = (RubyInteger) b;
+        if (negativeInt(context, intB)) {
+            throw runtime.newRangeError("Integer#pow() 1st argument cannot be negative when 2nd argument specified");
+        }
+        if (!(m instanceof RubyInteger)) {
+            throw runtime.newTypeError("Integer#pow() 2nd argument not allowed unless all arguments are integers");
+        }
+
+        if (negativeInt(context, m)) {
+            m = sites(context).op_uminus.call(context, m, m);
+            negaFlg = true;
+        }
+
+        if (!positiveInt(context, m)) {
+            throw runtime.newZeroDivisionError();
+        }
+        if (m instanceof RubyFixnum) {
+            long halfVal = HALF_LONG_MSB;
+            long mm = m.convertToInteger().getLongValue();
+            RubyFixnum modulo = (RubyFixnum) a.modulo(context, m);
+            if (mm <= halfVal) {
+                return modulo.intPowTmp1(context, intB, mm, negaFlg);
+            } else {
+                return modulo.intPowTmp2(context, intB, mm, negaFlg);
+            }
+        } else if (m instanceof RubyBignum) {
+            return ((RubyInteger) a.modulo(context, m)).intPowTmp3(context, intB, (RubyBignum) m, negaFlg);
+        }
+        // not reached
+        throw new RuntimeException("BUG: unexpected type " + m.getType());
+    }
+
+    protected IRubyObject intPowTmp3(ThreadContext context, RubyInteger y, RubyBignum m, boolean negaFlg) {
+        Ruby runtime = context.runtime;
+
+        BigInteger xn, yn, mn, zn;
+
+        if (this instanceof RubyFixnum) {
+            xn = BigInteger.valueOf(this.getLongValue());
+        } else {
+            xn = this.getBigIntegerValue();
+        }
+        if (y instanceof RubyFixnum) {
+            yn = BigInteger.valueOf(y.getLongValue());
+        } else {
+            yn = y.getBigIntegerValue();
+        }
+        mn = m.getBigIntegerValue();
+
+        zn = xn.modPow(yn, mn);
+        if (negaFlg & zn.signum() == 1) {
+            zn = zn.negate();
+        }
+        return RubyBignum.bignorm(runtime, zn);
+    }
+
     @JRubyMethod(name = "abs")
     public abstract IRubyObject abs(ThreadContext context);
 

diff --git a/core/src/main/java/org/jruby/runtime/JavaSites.java b/core/src/main/java/org/jruby/runtime/JavaSites.java
@@ -211,7 +211,9 @@ public static class IntegerSites {
         public final CallSite op_minus = new FunctionalCachingCallSite("-");
         public final CallSite op_quo = new FunctionalCachingCallSite("/");
         public final CallSite op_mod = new FunctionalCachingCallSite("%");
-        public final CallSite size = new FunctionalCachingCallSite("size");
+        public final CallSite size = new FunctionalCachingCallSite("**");
+        public final CallSite op_pow = new FunctionalCachingCallSite("size");
+        public final CallSite op_uminus = new FunctionalCachingCallSite("-@");
         public final CheckedSites to_i_checked = new CheckedSites("to_i");
     }
 
@@ -236,6 +238,8 @@ public static class FixnumSites {
         public final CallSite op_lt_bignum = new FunctionalCachingCallSite("<");
         public final CallSite op_exp_rational = new FunctionalCachingCallSite("**");
         public final CallSite fdiv = new FunctionalCachingCallSite("fdiv");
+        public final CallSite op_uminus = new FunctionalCachingCallSite("-@");
+        public final CallSite op_rshift = new FunctionalCachingCallSite(">>");
         public final CheckedSites checked_op_and = new CheckedSites("&");
         public final CheckedSites checked_op_or = new CheckedSites("|");
         public final CheckedSites checked_op_xor = new CheckedSites("^");