From 3681dbbdb03607a007ca7d5f84bb39ae5da48dc7 Mon Sep 17 00:00:00 2001
From: Andrew Konchin <andrii.konchyn@oracle.com>
Date: Tue, 23 Apr 2024 20:13:26 +0300
Subject: [PATCH] Optimize Integer#pow

---
 CHANGELOG.md                                  |  1 +
 src/main/ruby/truffleruby/core/integer.rb     | 14 +++++-----
 .../core/truffle/integer_operations.rb        | 27 +++++++++++++++++++
 3 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 26b096cd0ac..704202f5572 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -22,6 +22,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:
 
diff --git a/src/main/ruby/truffleruby/core/integer.rb b/src/main/ruby/truffleruby/core/integer.rb
index 695234dc292..456cb4f464b 100644
--- a/src/main/ruby/truffleruby/core/integer.rb
+++ b/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
diff --git a/src/main/ruby/truffleruby/core/truffle/integer_operations.rb b/src/main/ruby/truffleruby/core/truffle/integer_operations.rb
index 1890dafb4b9..7a250b863d6 100644
--- a/src/main/ruby/truffleruby/core/truffle/integer_operations.rb
+++ b/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