From a0eecfb5bae66470ccf40e27b9193fbf5c76618f Mon Sep 17 00:00:00 2001
From: Takashi Kokubun <takashikkbn@gmail.com>
Date: Mon, 8 Jan 2024 09:34:57 -0800
Subject: [PATCH] YJIT: Fallback Integer#<< if a shift amount varies (#9426)

* YJIT: Fallback Integer#<< if a shift amount varies

* YJIT: Do not fallback lshift in the first chain
---
 yjit.rb             |  2 +-
 yjit/src/codegen.rs | 17 ++++++++++++++---
 yjit/src/stats.rs   |  2 +-
 3 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/yjit.rb b/yjit.rb
index a2286117a8..a96db9b641 100644
--- a/yjit.rb
+++ b/yjit.rb
@@ -288,7 +288,7 @@ module RubyVM::YJIT
       ].each do |insn|
         print_counters(stats, out: out, prefix: "#{insn}_", prompt: "#{insn} exit reasons:", optional: true)
       end
-      print_counters(stats, out: out, prefix: 'lshift_', prompt: 'left shift (ltlt) exit reasons: ')
+      print_counters(stats, out: out, prefix: 'lshift_', prompt: 'left shift (opt_ltlt) 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 6d85ca4266..7e44c3add0 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -4637,17 +4637,28 @@ fn jit_rb_int_lshift(
 
     // 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() > 1 {
+        return false;
+    }
+
     let rhs = asm.stack_pop(1);
     let lhs = asm.stack_pop(1);
 
-    // Guard on the shift value we speculated on
+    // Guard on the shift amount we speculated on
     asm.cmp(rhs, comptime_shift.into());
-    asm.jne(Target::side_exit(Counter::lshift_amt_changed));
+    jit_chain_guard(
+        JCC_JNE,
+        jit,
+        asm,
+        ocb,
+        2, // defer_compilation increments chain_depth
+        Counter::lshift_amount_changed,
+    );
 
     let in_val = asm.sub(lhs, 1.into());
     let shift_opnd = Opnd::UImm(shift_amt as u64);
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index ac247d2fa8..ba1479e152 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -447,7 +447,7 @@ make_counters! {
     opt_mod_zero,
     opt_div_zero,
 
-    lshift_amt_changed,
+    lshift_amount_changed,
     lshift_overflow,
 
     opt_aref_argc_not_one,