From afba09d30f4c2ce6d0d77e06d5250967411cd74a Mon Sep 17 00:00:00 2001
From: Maxime Chevalier-Boisvert <maxime.chevalierboisvert@shopify.com>
Date: Wed, 17 Jan 2024 10:35:48 -0500
Subject: YJIT: specialized codegen for integer right shift (#9564)

* YJIT: specialized codegen for integer right shift

Used in optcarrot. May also be used to write pure-Ruby gems.
No overflow check or fixnum untagging required.

* Update yjit/src/codegen.rs

Co-authored-by: Takashi Kokubun <takashikkbn@gmail.com>

---------

Co-authored-by: Takashi Kokubun <takashikkbn@gmail.com>
---
 bootstraptest/test_yjit.rb |  5 +++++
 yjit.rb                    |  1 +
 yjit/src/codegen.rs        | 52 ++++++++++++++++++++++++++++++++++++++++++++++
 yjit/src/stats.rs          |  2 ++
 4 files changed, 60 insertions(+)

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index 70a5501be2..5b53e8089b 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -4352,3 +4352,8 @@ assert_equal '0', %q{
 assert_equal '[2, 4611686018427387904]', %q{
   [1.succ, 4611686018427387903.succ]
 }
+
+# Integer right shift
+assert_equal '[0, 1, -4]', %q{
+  [0 >> 1, 2 >> 1, -7 >> 1]
+}
diff --git a/yjit.rb b/yjit.rb
index 84d4d99b1f..485b4062fa 100644
--- a/yjit.rb
+++ b/yjit.rb
@@ -290,6 +290,7 @@ module RubyVM::YJIT
         print_counters(stats, out: out, prefix: "#{insn}_", prompt: "#{insn} exit reasons:", optional: true)
       end
       print_counters(stats, out: out, prefix: 'lshift_', prompt: 'left shift (opt_ltlt) exit reasons: ')
+      print_counters(stats, out: out, prefix: 'rshift_', prompt: 'right shift (>>) exit reasons: ')
       print_counters(stats, out: out, prefix: 'invalidate_', prompt: 'invalidation reasons: ')
     end
 
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 1ccb55d3b4..7d42030345 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -4737,6 +4737,57 @@ fn jit_rb_int_lshift(
     true
 }
 
+fn jit_rb_int_rshift(
+    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 {
+    if asm.ctx.two_fixnums_on_stack(jit) != Some(true) {
+        return false;
+    }
+    guard_two_fixnums(jit, asm, ocb);
+
+    let comptime_shift = jit.peek_at_stack(&asm.ctx, 0);
+
+    // Untag the fixnum shift amount
+    let shift_amt = comptime_shift.as_isize() >> 1;
+    if shift_amt > 63 || shift_amt < 0 {
+        return false;
+    }
+
+    // Fallback to a C call if the shift amount varies
+    if asm.ctx.get_chain_depth() > 2 {
+        return false;
+    }
+
+    let rhs = asm.stack_pop(1);
+    let lhs = asm.stack_pop(1);
+
+    // Guard on the shift amount we speculated on
+    asm.cmp(rhs, comptime_shift.into());
+    jit_chain_guard(
+        JCC_JNE,
+        jit,
+        asm,
+        ocb,
+        2, // defer_compilation increments chain_depth
+        Counter::rshift_amount_changed,
+    );
+
+    let shift_opnd = Opnd::UImm(shift_amt as u64);
+    let out_val = asm.rshift(lhs, shift_opnd);
+    let out_val = asm.or(out_val, 1.into());
+
+    let ret_opnd = asm.stack_push(Type::Fixnum);
+    asm.mov(ret_opnd, out_val);
+    true
+}
+
 fn jit_rb_int_aref(
     jit: &mut JITState,
     asm: &mut Assembler,
@@ -8909,6 +8960,7 @@ pub fn yjit_reg_method_codegen_fns() {
         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_rshift);
         yjit_reg_method(rb_cInteger, "[]", jit_rb_int_aref);
 
         yjit_reg_method(rb_cString, "empty?", jit_rb_str_empty_p);
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index 924628e13a..4b77dceaf1 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -458,6 +458,8 @@ make_counters! {
     lshift_amount_changed,
     lshift_overflow,
 
+    rshift_amount_changed,
+
     opt_aref_argc_not_one,
     opt_aref_arg_not_fixnum,
     opt_aref_not_array,
-- 
cgit v1.2.3