From b6f6fc6e870d00e5151647f3f14eaa16b8fe145b Mon Sep 17 00:00:00 2001
From: Dave Schwantes <dorkrawk@github.com>
Date: Thu, 30 Jun 2022 09:26:46 -0500
Subject: [PATCH] YJIT: Refactor gen_opt_mod (#6078)

Refactor gen_opt_mod in YJIT
---
 vm_insnhelper.c     |  6 -----
 yjit.c              |  7 ++++++
 yjit/src/codegen.rs | 56 ++++++++++++++++++++++++++++++---------------
 yjit/src/cruby.rs   |  4 +++-
 4 files changed, 48 insertions(+), 25 deletions(-)

diff --git a/vm_insnhelper.c b/vm_insnhelper.c
index da04e3521f..894c44693a 100644
--- a/vm_insnhelper.c
+++ b/vm_insnhelper.c
@@ -5253,12 +5253,6 @@ vm_opt_mod(VALUE recv, VALUE obj)
     }
 }
 
-VALUE
-rb_vm_opt_mod(VALUE recv, VALUE obj)
-{
-    return vm_opt_mod(recv, obj);
-}
-
 static VALUE
 vm_opt_neq(const rb_iseq_t *iseq, CALL_DATA cd, CALL_DATA cd_eq, VALUE recv, VALUE obj)
 {
diff --git a/yjit.c b/yjit.c
index a30e3bad94..108a376282 100644
--- a/yjit.c
+++ b/yjit.c
@@ -13,6 +13,7 @@
 #include "internal/variable.h"
 #include "internal/compile.h"
 #include "internal/class.h"
+#include "internal/fixnum.h"
 #include "gc.h"
 #include "vm_core.h"
 #include "vm_callinfo.h"
@@ -714,6 +715,12 @@ rb_yarv_ary_entry_internal(VALUE ary, long offset)
     return rb_ary_entry_internal(ary, offset);
 }
 
+VALUE
+rb_yarv_fix_mod_fix(VALUE recv, VALUE obj)
+{
+    return rb_fix_mod_fix(recv, obj);
+}
+
 // Print the Ruby source location of some ISEQ for debugging purposes
 void
 rb_yjit_dump_iseq_loc(const rb_iseq_t *iseq, uint32_t insn_idx)
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 1b8b99d530..cdd01221d4 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -2951,30 +2951,50 @@ fn gen_opt_mod(
     cb: &mut CodeBlock,
     ocb: &mut OutlinedCb,
 ) -> CodegenStatus {
-    // Save the PC and SP because the callee may allocate bignums
-    // Note that this modifies REG_SP, which is why we do it first
-    jit_prepare_routine_call(jit, ctx, cb, REG0);
+    // Defer compilation so we can specialize on a runtime `self`
+    if !jit_at_current_insn(jit) {
+        defer_compilation(jit, ctx, cb, ocb);
+        return EndBlock;
+    }
 
-    let side_exit = get_side_exit(jit, ocb, ctx);
+    let comptime_a = jit_peek_at_stack(jit, ctx, 1);
+    let comptime_b = jit_peek_at_stack(jit, ctx, 0);
 
-    // Get the operands from the stack
-    let arg1 = ctx.stack_pop(1);
-    let arg0 = ctx.stack_pop(1);
+    if comptime_a.fixnum_p() && comptime_b.fixnum_p() {
+        // Create a side-exit to fall back to the interpreter
+        // Note: we generate the side-exit before popping operands from the stack
+        let side_exit = get_side_exit(jit, ocb, ctx);
 
-    // Call rb_vm_opt_mod(VALUE recv, VALUE obj)
-    mov(cb, C_ARG_REGS[0], arg0);
-    mov(cb, C_ARG_REGS[1], arg1);
-    call_ptr(cb, REG0, rb_vm_opt_mod as *const u8);
+        if !assume_bop_not_redefined(jit, ocb, INTEGER_REDEFINED_OP_FLAG, BOP_MOD) {
+            return CantCompile;
+        }
 
-    // If val == Qundef, bail to do a method call
-    cmp(cb, RAX, imm_opnd(Qundef.as_i64()));
-    je_ptr(cb, side_exit);
+        // Check that both operands are fixnums
+        guard_two_fixnums(ctx, cb, side_exit);
 
-    // Push the return value onto the stack
-    let stack_ret = ctx.stack_push(Type::Unknown);
-    mov(cb, stack_ret, RAX);
+        // Get the operands and destination from the stack
+        let arg1 = ctx.stack_pop(1);
+        let arg0 = ctx.stack_pop(1);
 
-    KeepCompiling
+        mov(cb, C_ARG_REGS[0], arg0);
+        mov(cb, C_ARG_REGS[1], arg1);
+
+        // Check for arg0 % 0
+        cmp(cb, C_ARG_REGS[1], imm_opnd(VALUE::fixnum_from_usize(0).as_i64()));
+        je_ptr(cb, side_exit);
+
+        // Call rb_fix_mod_fix(VALUE recv, VALUE obj)
+        call_ptr(cb, REG0, rb_fix_mod_fix as *const u8);
+
+        // Push the return value onto the stack
+        let stack_ret = ctx.stack_push(Type::Unknown);
+        mov(cb, stack_ret, RAX);
+
+        KeepCompiling
+    } else {
+        // Delegate to send, call the method on the recv
+        gen_opt_send_without_block(jit, ctx, cb, ocb)
+    }
 }
 
 fn gen_opt_ltlt(
diff --git a/yjit/src/cruby.rs b/yjit/src/cruby.rs
index 51ba9c1531..8cebfd58a5 100644
--- a/yjit/src/cruby.rs
+++ b/yjit/src/cruby.rs
@@ -235,6 +235,9 @@ extern "C" {
     #[link_name = "rb_yarv_ary_entry_internal"]
     pub fn rb_ary_entry_internal(ary: VALUE, offset: c_long) -> VALUE;
 
+    #[link_name = "rb_yarv_fix_mod_fix"]
+    pub fn rb_fix_mod_fix(recv: VALUE, obj: VALUE) -> VALUE;
+
     #[link_name = "rb_FL_TEST"]
     pub fn FL_TEST(obj: VALUE, flags: VALUE) -> VALUE;
 
@@ -255,7 +258,6 @@ extern "C" {
 
     // Ruby only defines these in vm_insnhelper.c, not in any header.
     // Parsing it would result in a lot of duplicate definitions.
-    pub fn rb_vm_opt_mod(recv: VALUE, obj: VALUE) -> VALUE;
     pub fn rb_vm_splat_array(flag: VALUE, ary: VALUE) -> VALUE;
     pub fn rb_vm_defined(
         ec: EcPtr,