From 0d3fc08ff46cdfa389170db6141f1747d5d9d283 Mon Sep 17 00:00:00 2001
From: Takashi Kokubun <takashikkbn@gmail.com>
Date: Wed, 30 Nov 2022 13:16:11 -0800
Subject: [PATCH] YJIT: Optimize rb_int_equal (#6838)

---
 test/ruby/test_yjit.rb | 16 ++++++++++++++++
 yjit/src/codegen.rs    | 31 +++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/test/ruby/test_yjit.rb b/test/ruby/test_yjit.rb
index a4e584a07c6fc8..aa29cf471fa550 100644
--- a/test/ruby/test_yjit.rb
+++ b/test/ruby/test_yjit.rb
@@ -827,6 +827,22 @@ def foo
     RUBY
   end
 
+  def test_int_equal
+    assert_compiles(<<~'RUBY', exits: :any, result: [true, false, true, false, true, false, true, false])
+      def eq(a, b)
+        a == b
+      end
+
+      def eqq(a, b)
+        a === b
+      end
+
+      big1 = 2 ** 65
+      big2 = big1 + 1
+      [eq(1, 1), eq(1, 2), eq(big1, big1), eq(big1, big2), eqq(1, 1), eqq(1, 2), eqq(big1, big1), eqq(big1, big2)]
+    RUBY
+  end
+
   def test_code_gc
     assert_compiles(code_gc_helpers + <<~'RUBY', exits: :any, result: :ok)
       return :not_paged unless add_pages(100) # prepare freeable pages
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index d1f40e0b6dbca3..bb4091c5c29df6 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -3651,6 +3651,35 @@ fn jit_rb_obj_equal(
     true
 }
 
+// Codegen for rb_int_equal()
+fn jit_rb_int_equal(
+    jit: &mut JITState,
+    ctx: &mut Context,
+    asm: &mut Assembler,
+    ocb: &mut OutlinedCb,
+    _ci: *const rb_callinfo,
+    _cme: *const rb_callable_method_entry_t,
+    _block: Option<IseqPtr>,
+    _argc: i32,
+    _known_recv_class: *const VALUE,
+) -> bool {
+    let side_exit = get_side_exit(jit, ocb, ctx);
+
+    // Check that both operands are fixnums
+    guard_two_fixnums(jit, ctx, asm, ocb, side_exit);
+
+    // Compare the arguments
+    asm.comment("rb_int_equal");
+    let arg1 = ctx.stack_pop(1);
+    let arg0 = ctx.stack_pop(1);
+    asm.cmp(arg0, arg1);
+    let ret_opnd = asm.csel_e(Qtrue.into(), Qfalse.into());
+
+    let stack_ret = ctx.stack_push(Type::UnknownImm);
+    asm.mov(stack_ret, ret_opnd);
+    true
+}
+
 /// If string is frozen, duplicate it to get a non-frozen string. Otherwise, return it.
 fn jit_rb_str_uplus(
     _jit: &mut JITState,
@@ -6880,6 +6909,8 @@ impl CodegenGlobals {
             self.yjit_reg_method(rb_cModule, "==", jit_rb_obj_equal);
             self.yjit_reg_method(rb_cSymbol, "==", jit_rb_obj_equal);
             self.yjit_reg_method(rb_cSymbol, "===", jit_rb_obj_equal);
+            self.yjit_reg_method(rb_cInteger, "==", jit_rb_int_equal);
+            self.yjit_reg_method(rb_cInteger, "===", jit_rb_int_equal);
 
             // rb_str_to_s() methods in string.c
             self.yjit_reg_method(rb_cString, "to_s", jit_rb_str_to_s);