YJIT: Optimize Integer#succ (#9519)

ruby · Jan 15, 2024 · 5e61cc2 · 5e61cc2
1 parent 7c6d7fb
commit 5e61cc2
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 0 deletions.
diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
@@ -4347,3 +4347,8 @@ def entry
 
   entry
 }
+
+# Integer succ and overflow
+assert_equal '[2, 4611686018427387904]', %q{
+  [1.succ, 4611686018427387903.succ]
+}
diff --git a/yjit.rb b/yjit.rb
@@ -284,6 +284,7 @@ def _print_stats_reasons(stats, out) # :nodoc:
         opt_mod
         opt_mult
         opt_plus
+        opt_succ
         setlocal
       ].each do |insn|
         print_counters(stats, out: out, prefix: "#{insn}_", prompt: "#{insn} exit reasons:", optional: true)

diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
@@ -4610,6 +4610,36 @@ fn jit_rb_int_equal(
     true
 }
 
+fn jit_rb_int_succ(
+    _jit: &mut JITState,
+    asm: &mut Assembler,
+    _ocb: &mut OutlinedCb,
+    _ci: *const rb_callinfo,
+    _cme: *const rb_callable_method_entry_t,
+    _block: Option<BlockHandler>,
+    _argc: i32,
+    _known_recv_class: *const VALUE,
+) -> bool {
+    // Guard the receiver is fixnum
+    let recv_type = asm.ctx.get_opnd_type(StackOpnd(0));
+    let recv = asm.stack_pop(1);
+    if recv_type != Type::Fixnum {
+        asm_comment!(asm, "guard object is fixnum");
+        asm.test(recv, Opnd::Imm(RUBY_FIXNUM_FLAG as i64));
+        asm.jz(Target::side_exit(Counter::opt_succ_not_fixnum));
+    }
+
+    asm_comment!(asm, "Integer#succ");
+    let out_val = asm.add(recv, Opnd::Imm(2)); // 2 is untagged Fixnum 1
+    asm.jo(Target::side_exit(Counter::opt_succ_overflow));
+
+    // Push the output onto the stack
+    let dst = asm.stack_push(Type::Fixnum);
+    asm.mov(dst, out_val);
+
+    true
+}
+
 fn jit_rb_int_div(
     jit: &mut JITState,
     asm: &mut Assembler,
@@ -8852,6 +8882,7 @@ pub fn yjit_reg_method_codegen_fns() {
         yjit_reg_method(rb_cInteger, "==", jit_rb_int_equal);
         yjit_reg_method(rb_cInteger, "===", jit_rb_int_equal);
 
+        yjit_reg_method(rb_cInteger, "succ", jit_rb_int_succ);
         yjit_reg_method(rb_cInteger, "/", jit_rb_int_div);
         yjit_reg_method(rb_cInteger, "<<", jit_rb_int_lshift);
         yjit_reg_method(rb_cInteger, "[]", jit_rb_int_aref);

diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
@@ -444,6 +444,9 @@ make_counters! {
     opt_minus_overflow,
     opt_mult_overflow,
 
+    opt_succ_not_fixnum,
+    opt_succ_overflow,
+
     opt_mod_zero,
     opt_div_zero,