YJIT: Add codegen for Integer methods

bootstraptest/test_yjit.rb

@@ -3817,3 +3817,8 @@ def calling_my_func
   end
   calling_my_func
 }
+
+# Test Integer#[] with 2 args
+assert_equal '0', %q{
+  3[0, 0]
+}

common.mk

@@ -18201,6 +18201,8 @@ yjit.$(OBJEXT): $(CCAN_DIR)/str/str.h
 yjit.$(OBJEXT): $(hdrdir)/ruby/ruby.h
 yjit.$(OBJEXT): $(top_srcdir)/internal/array.h
 yjit.$(OBJEXT): $(top_srcdir)/internal/basic_operators.h
+yjit.$(OBJEXT): $(top_srcdir)/internal/bignum.h
+yjit.$(OBJEXT): $(top_srcdir)/internal/bits.h
 yjit.$(OBJEXT): $(top_srcdir)/internal/class.h
 yjit.$(OBJEXT): $(top_srcdir)/internal/compile.h
 yjit.$(OBJEXT): $(top_srcdir)/internal/compilers.h
@@ -18209,6 +18211,7 @@ yjit.$(OBJEXT): $(top_srcdir)/internal/fixnum.h
 yjit.$(OBJEXT): $(top_srcdir)/internal/gc.h
 yjit.$(OBJEXT): $(top_srcdir)/internal/hash.h
 yjit.$(OBJEXT): $(top_srcdir)/internal/imemo.h
+yjit.$(OBJEXT): $(top_srcdir)/internal/numeric.h
 yjit.$(OBJEXT): $(top_srcdir)/internal/sanitizers.h
 yjit.$(OBJEXT): $(top_srcdir)/internal/serial.h
 yjit.$(OBJEXT): $(top_srcdir)/internal/static_assert.h
@@ -18378,6 +18381,7 @@ yjit.$(OBJEXT): {$(VPATH)}internal/memory.h
 yjit.$(OBJEXT): {$(VPATH)}internal/method.h
 yjit.$(OBJEXT): {$(VPATH)}internal/module.h
 yjit.$(OBJEXT): {$(VPATH)}internal/newobj.h
+yjit.$(OBJEXT): {$(VPATH)}internal/numeric.h
 yjit.$(OBJEXT): {$(VPATH)}internal/scan_args.h
 yjit.$(OBJEXT): {$(VPATH)}internal/special_consts.h
 yjit.$(OBJEXT): {$(VPATH)}internal/static_assert.h

yjit.c

@@ -14,6 +14,7 @@
 #include "internal/compile.h"
 #include "internal/class.h"
 #include "internal/fixnum.h"
+#include "internal/numeric.h"
 #include "internal/gc.h"
 #include "vm_core.h"
 #include "vm_callinfo.h"
@@ -855,11 +856,23 @@ rb_yjit_rb_ary_subseq_length(VALUE ary, long beg)
 }
 
 VALUE
-rb_yarv_fix_mod_fix(VALUE recv, VALUE obj)
+rb_yjit_fix_div_fix(VALUE recv, VALUE obj)
+{
+    return rb_fix_div_fix(recv, obj);
+}
+
+VALUE
+rb_yjit_fix_mod_fix(VALUE recv, VALUE obj)
 {
     return rb_fix_mod_fix(recv, obj);
 }
 
+VALUE
+rb_yjit_fix_mul_fix(VALUE recv, VALUE obj)
+{
+    return rb_fix_mul_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)

yjit/bindgen/src/main.rs

@@ -213,6 +213,9 @@ fn main() {
         .allowlist_function("rb_sym2id")
         .allowlist_function("rb_str_intern")
 
+        // From internal/numeric.h
+        .allowlist_function("rb_fix_aref")
+
         // From internal/string.h
         .allowlist_function("rb_ec_str_resurrect")
         .allowlist_function("rb_str_concat_literals")
@@ -413,7 +416,9 @@ fn main() {
         .allowlist_function("rb_yarv_str_eql_internal")
         .allowlist_function("rb_str_neq_internal")
         .allowlist_function("rb_yarv_ary_entry_internal")
-        .allowlist_function("rb_yarv_fix_mod_fix")
+        .allowlist_function("rb_yjit_fix_div_fix")
+        .allowlist_function("rb_yjit_fix_mod_fix")
+        .allowlist_function("rb_yjit_fix_mul_fix")
         .allowlist_function("rb_FL_TEST")
         .allowlist_function("rb_FL_TEST_RAW")
         .allowlist_function("rb_RB_TYPE_P")

yjit/src/codegen.rs

@@ -4354,6 +4354,98 @@ fn jit_rb_int_equal(
     true
 }
 
+fn jit_rb_int_mul(
+    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 {
+    if ctx.two_fixnums_on_stack(jit) != Some(true) {
+        return false;
+    }
+    guard_two_fixnums(jit, ctx, asm, ocb);
+
+    // rb_fix_mul_fix may allocate memory for Bignum
+    jit_prepare_routine_call(jit, ctx, asm);
+
+    asm.comment("Integer#*");
+    let obj = ctx.stack_pop(1);
+    let recv = ctx.stack_pop(1);
+    let ret = asm.ccall(rb_fix_mul_fix as *const u8, vec![recv, obj]);
+
+    let ret_opnd = ctx.stack_push(asm, Type::Unknown);
+    asm.mov(ret_opnd, ret);
+    true
+}
+
+fn jit_rb_int_div(
+    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 {
+    if ctx.two_fixnums_on_stack(jit) != Some(true) {
+        return false;
+    }
+    guard_two_fixnums(jit, ctx, asm, ocb);
+
+    asm.comment("Integer#/");
+    asm.spill_temps(ctx); // for ccall (must be done before stack_pop)
+    let obj = ctx.stack_pop(1);
+    let recv = ctx.stack_pop(1);
+
+    // Check for arg0 % 0
+    asm.cmp(obj, VALUE::fixnum_from_usize(0).as_i64().into());
+    asm.je(side_exit(jit, ctx, ocb));
+
+    let ret = asm.ccall(rb_fix_div_fix as *const u8, vec![recv, obj]);
+
+    let ret_opnd = ctx.stack_push(asm, Type::Fixnum);
+    asm.mov(ret_opnd, ret);
+    true
+}
+
+fn jit_rb_int_aref(
+    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 {
+    if argc != 1 {
+        return false;
+    }
+    if ctx.two_fixnums_on_stack(jit) != Some(true) {
+        return false;
+    }
+    guard_two_fixnums(jit, ctx, asm, ocb);
+
+    asm.comment("Integer#[]");
+    asm.spill_temps(ctx); // for ccall (must be done before stack_pop)
+    let obj = ctx.stack_pop(1);
+    let recv = ctx.stack_pop(1);
+
+    let ret = asm.ccall(rb_fix_aref as *const u8, vec![recv, obj]);
+
+    let ret_opnd = ctx.stack_push(asm, Type::Fixnum);
+    asm.mov(ret_opnd, ret);
+    true
+}
+
 /// If string is frozen, duplicate it to get a non-frozen string. Otherwise, return it.
 fn jit_rb_str_uplus(
     jit: &mut JITState,
@@ -8136,6 +8228,10 @@ impl CodegenGlobals {
             self.yjit_reg_method(rb_cInteger, "==", jit_rb_int_equal);
             self.yjit_reg_method(rb_cInteger, "===", jit_rb_int_equal);
 
+            self.yjit_reg_method(rb_cInteger, "*", jit_rb_int_mul);
+            self.yjit_reg_method(rb_cInteger, "/", jit_rb_int_div);
+            self.yjit_reg_method(rb_cInteger, "[]", jit_rb_int_aref);
+
             // rb_str_to_s() methods in string.c
             self.yjit_reg_method(rb_cString, "empty?", jit_rb_str_empty_p);
             self.yjit_reg_method(rb_cString, "to_s", jit_rb_str_to_s);

yjit/src/cruby.rs

@@ -184,7 +184,9 @@ pub use rb_get_cikw_keywords_idx as get_cikw_keywords_idx;
 pub use rb_get_call_data_ci as get_call_data_ci;
 pub use rb_yarv_str_eql_internal as rb_str_eql_internal;
 pub use rb_yarv_ary_entry_internal as rb_ary_entry_internal;
-pub use rb_yarv_fix_mod_fix as rb_fix_mod_fix;
+pub use rb_yjit_fix_div_fix as rb_fix_div_fix;
+pub use rb_yjit_fix_mod_fix as rb_fix_mod_fix;
+pub use rb_yjit_fix_mul_fix as rb_fix_mul_fix;
 pub use rb_FL_TEST as FL_TEST;
 pub use rb_FL_TEST_RAW as FL_TEST_RAW;
 pub use rb_RB_TYPE_P as RB_TYPE_P;

yjit/src/cruby_bindings.inc.rs

@@ -1183,6 +1183,7 @@ extern "C" {
     ) -> ::std::os::raw::c_int;
     pub fn rb_insn_len(insn: VALUE) -> ::std::os::raw::c_int;
     pub fn rb_vm_insn_decode(encoded: VALUE) -> ::std::os::raw::c_int;
+    pub fn rb_fix_aref(fix: VALUE, idx: VALUE) -> VALUE;
     pub fn rb_vm_insn_addr2opcode(addr: *const ::std::os::raw::c_void) -> ::std::os::raw::c_int;
     pub fn rb_iseq_line_no(iseq: *const rb_iseq_t, pos: usize) -> ::std::os::raw::c_uint;
     pub fn rb_iseqw_to_iseq(iseqw: VALUE) -> *const rb_iseq_t;
@@ -1299,7 +1300,9 @@ extern "C" {
     pub fn rb_yarv_ary_entry_internal(ary: VALUE, offset: ::std::os::raw::c_long) -> VALUE;
     pub fn rb_ary_unshift_m(argc: ::std::os::raw::c_int, argv: *mut VALUE, ary: VALUE) -> VALUE;
     pub fn rb_yjit_rb_ary_subseq_length(ary: VALUE, beg: ::std::os::raw::c_long) -> VALUE;
-    pub fn rb_yarv_fix_mod_fix(recv: VALUE, obj: VALUE) -> VALUE;
+    pub fn rb_yjit_fix_div_fix(recv: VALUE, obj: VALUE) -> VALUE;
+    pub fn rb_yjit_fix_mod_fix(recv: VALUE, obj: VALUE) -> VALUE;
+    pub fn rb_yjit_fix_mul_fix(recv: VALUE, obj: VALUE) -> VALUE;
     pub fn rb_yjit_dump_iseq_loc(iseq: *const rb_iseq_t, insn_idx: u32);
     pub fn rb_FL_TEST(obj: VALUE, flags: VALUE) -> VALUE;
     pub fn rb_FL_TEST_RAW(obj: VALUE, flags: VALUE) -> VALUE;