YJIT: implement codegen for rb_int_lshift (#8201)

* YJIT: implement codegen for rb_int_lshift

* Update yjit/src/asm/x86_64/mod.rs

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

---------

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

4 files changed, 35 insertions, 16 deletions

diff --git a/yjit/src/asm/x86_64/mod.rs b/yjit/src/asm/x86_64/mod.rs
index 3b2df6a7e6..4ff89dafe3 100644
--- a/yjit/src/asm/x86_64/mod.rs
+++ b/yjit/src/asm/x86_64/mod.rs
@@ -1277,7 +1277,7 @@ fn write_shift(cb: &mut CodeBlock, op_mem_one_pref: u8, op_mem_cl_pref: u8, op_m
         }
 
         _ => {
-            unreachable!();
+            unreachable!("unsupported operands: {:?}, {:?}", opnd0, opnd1);
         }
     }
 }
diff --git a/yjit/src/backend/ir.rs b/yjit/src/backend/ir.rs
index d29199a390..a6803f7a63 100644
--- a/yjit/src/backend/ir.rs
+++ b/yjit/src/backend/ir.rs
@@ -544,6 +544,7 @@ impl Insn {
             Insn::Jb(target) |
             Insn::Je(target) |
             Insn::Jl(target) |
+            Insn::Jg(target) |
             Insn::Jmp(target) |
             Insn::Jne(target) |
             Insn::Jnz(target) |
@@ -692,6 +693,7 @@ impl Insn {
             Insn::Jb(target) |
             Insn::Je(target) |
             Insn::Jl(target) |
+            Insn::Jg(target) |
             Insn::Jmp(target) |
             Insn::Jne(target) |
             Insn::Jnz(target) |
@@ -1843,6 +1845,10 @@ impl Assembler {
         self.push_insn(Insn::Jl(target));
     }
 
+    pub fn jg(&mut self, target: Target) {
+        self.push_insn(Insn::Jg(target));
+    }
+
     pub fn jmp(&mut self, target: Target) {
         self.push_insn(Insn::Jmp(target));
     }
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 1baa5edd64..3af6167045 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -4471,7 +4471,6 @@ fn jit_rb_int_div(
     true
 }
 
-/*
 fn jit_rb_int_lshift(
     jit: &mut JITState,
     asm: &mut Assembler,
@@ -4487,28 +4486,42 @@ fn jit_rb_int_lshift(
     }
     guard_two_fixnums(jit, asm, ocb);
 
+    let comptime_shift = jit.peek_at_stack(&asm.ctx, 0);
+
+    if !comptime_shift.fixnum_p() {
+        return false;
+    }
+
+    // Untag the fixnum shift amount
+    let shift_amt = comptime_shift.as_isize() >> 1;
+
+    if shift_amt > 63 || shift_amt < 0 {
+        return false;
+    }
+
     let rhs = asm.stack_pop(1);
     let lhs = asm.stack_pop(1);
 
-    // Ruby supports using a negative shift value
-    asm.comment("Guard shift negative");
-    let shift_val = asm.sub(rhs, 1.into());
-    asm.cmp(shift_val, 0.into());
-    asm.jl(Target::side_exit(Counter::lshift_range));
+    // Guard on the shift value we speculated on
+    asm.cmp(rhs, comptime_shift.into());
+    asm.jne(Target::side_exit(Counter::lshift_amt_changed));
+
+    let in_val = asm.sub(lhs, 1.into());
+    let shift_opnd = Opnd::UImm(shift_amt as u64);
+    let out_val = asm.lshift(in_val, shift_opnd);
+    let unshifted = asm.rshift(out_val, shift_opnd);
 
-    asm.cmp(shift_val, 63.into());
-    asm.jg(Target::side_exit(Counter::lshift_range));
+    // Guard that we did not overflow
+    asm.cmp(unshifted, in_val);
+    asm.jne(Target::side_exit(Counter::lshift_overflow));
 
-    // FIXME: we don't yet support shift with non-immediate values in the backend
-    // Do the shifting
-    let out_val = asm.lshift(lhs, shift_val);
-    asm.jo(Target::side_exit(Counter::lshift_overflow));
+    // Re-tag the output value
+    let out_val = asm.add(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,
@@ -8593,7 +8606,7 @@ impl CodegenGlobals {
 
             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_lshift);
+            self.yjit_reg_method(rb_cInteger, "<<", jit_rb_int_lshift);
             self.yjit_reg_method(rb_cInteger, "[]", jit_rb_int_aref);
 
             // rb_str_to_s() methods in string.c
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index 66b36a61c5..1c87a8e74c 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -348,7 +348,7 @@ make_counters! {
     opt_mod_zero,
     opt_div_zero,
 
-    lshift_range,
+    lshift_amt_changed,
     lshift_overflow,
 
     opt_aref_argc_not_one,