[GR-18163] Optimize Integer#pow

PullRequest: truffleruby/4251

CHANGELOG.md

@@ -23,6 +23,7 @@ Compatibility:
 
 Performance:
 * Fix inline caching for Regexp creation from Strings (#3492, @andrykonchin, @eregon).
+* Optimize `Integer#pow` method for small modulus values (#3544, @andrykonchin).
 
 Changes:
 

src/main/ruby/truffleruby/core/integer.rb

@@ -130,13 +130,15 @@ def nobits?(mask)
   end
 
   def pow(e, m = undefined)
-    if Primitive.undefined?(m)
-      self ** e
-    else
-      raise TypeError, '2nd argument not allowed unless a 1st argument is integer' unless Primitive.is_a?(e, Integer)
-      raise TypeError, '2nd argument not allowed unless all arguments are integers' unless Primitive.is_a?(m, Integer)
-      raise RangeError, '1st argument cannot be negative when 2nd argument specified' if e.negative?
+    return self ** e if Primitive.undefined?(m)
+
+    raise TypeError, '2nd argument not allowed unless a 1st argument is integer' unless Primitive.is_a?(e, Integer)
+    raise TypeError, '2nd argument not allowed unless all arguments are integers' unless Primitive.is_a?(m, Integer)
+    raise RangeError, '1st argument cannot be negative when 2nd argument specified' if e.negative?
 
+    if Primitive.integer_fits_into_long(m)
+      Truffle::IntegerOperations.modular_exponentiation(self, e, m)
+    else
       Primitive.mod_pow(self, e, m)
     end
   end

src/main/ruby/truffleruby/core/truffle/integer_operations.rb

@@ -35,5 +35,32 @@ def self.bits_reference_range(n, range)
         0
       end
     end
+
+    # Implementation of a**b mod c operation
+    # See https://en.wikipedia.org/wiki/Modular_exponentiation
+    # The only difference with the Right-to-left binary method is a special handling of negative modulus -
+    # a**b mod -c = (a**b mod c) - c
+    # MRI: similar to int_pow_tmp1/int_pow_tmp2/int_pow_tmp3
+    def self.modular_exponentiation(base, exponent, modulus)
+      return 0 if modulus == 1
+
+      negative = modulus < 0
+      modulus = modulus.abs
+
+      result = 1
+      base %= modulus
+
+      while exponent > 0
+        if exponent.odd?
+          result = (result * base) % modulus
+        end
+
+        base = (base * base) % modulus
+        exponent >>= 1
+      end
+
+      result -= modulus if negative
+      result
+    end
   end
 end