YJIT: Check correct BOP on gen_fixnum_cmp (#7303)

2 files changed, 45 insertions, 5 deletions

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index 66b663b062..ebb19c8d3a 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -374,6 +374,45 @@ assert_equal 'false', %q{
   less_than 2
 }
 
+# BOP redefinition works on Integer#<=
+assert_equal 'false', %q{
+  def le(x, y) = x <= y
+
+  le(2, 2)
+
+  class Integer
+    def <=(_) = false
+  end
+
+  le(2, 2)
+}
+
+# BOP redefinition works on Integer#>
+assert_equal 'false', %q{
+  def gt(x, y) = x > y
+
+  gt(3, 2)
+
+  class Integer
+    def >(_) = false
+  end
+
+  gt(3, 2)
+}
+
+# BOP redefinition works on Integer#>=
+assert_equal 'false', %q{
+  def ge(x, y) = x >= y
+
+  ge(2, 2)
+
+  class Integer
+    def >=(_) = false
+  end
+
+  ge(2, 2)
+}
+
 # Putobject, less-than operator, fixnums
 assert_equal '2', %q{
     def check_index(index)
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 9eb192b66b..4183085d2f 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -2602,6 +2602,7 @@ fn gen_fixnum_cmp(
     asm: &mut Assembler,
     ocb: &mut OutlinedCb,
     cmov_op: CmovFn,
+    bop: ruby_basic_operators,
 ) -> CodegenStatus {
     let two_fixnums = match ctx.two_fixnums_on_stack(jit) {
         Some(two_fixnums) => two_fixnums,
@@ -2617,7 +2618,7 @@ fn gen_fixnum_cmp(
         // Note: we generate the side-exit before popping operands from the stack
         let side_exit = get_side_exit(jit, ocb, ctx);
 
-        if !assume_bop_not_redefined(jit, ocb, INTEGER_REDEFINED_OP_FLAG, BOP_LT) {
+        if !assume_bop_not_redefined(jit, ocb, INTEGER_REDEFINED_OP_FLAG, bop) {
             return CantCompile;
         }
 
@@ -2648,7 +2649,7 @@ fn gen_opt_lt(
     asm: &mut Assembler,
     ocb: &mut OutlinedCb,
 ) -> CodegenStatus {
-    gen_fixnum_cmp(jit, ctx, asm, ocb, Assembler::csel_l)
+    gen_fixnum_cmp(jit, ctx, asm, ocb, Assembler::csel_l, BOP_LT)
 }
 
 fn gen_opt_le(
@@ -2657,7 +2658,7 @@ fn gen_opt_le(
     asm: &mut Assembler,
     ocb: &mut OutlinedCb,
 ) -> CodegenStatus {
-    gen_fixnum_cmp(jit, ctx, asm, ocb, Assembler::csel_le)
+    gen_fixnum_cmp(jit, ctx, asm, ocb, Assembler::csel_le, BOP_LE)
 }
 
 fn gen_opt_ge(
@@ -2666,7 +2667,7 @@ fn gen_opt_ge(
     asm: &mut Assembler,
     ocb: &mut OutlinedCb,
 ) -> CodegenStatus {
-    gen_fixnum_cmp(jit, ctx, asm, ocb, Assembler::csel_ge)
+    gen_fixnum_cmp(jit, ctx, asm, ocb, Assembler::csel_ge, BOP_GE)
 }
 
 fn gen_opt_gt(
@@ -2675,7 +2676,7 @@ fn gen_opt_gt(
     asm: &mut Assembler,
     ocb: &mut OutlinedCb,
 ) -> CodegenStatus {
-    gen_fixnum_cmp(jit, ctx, asm, ocb, Assembler::csel_g)
+    gen_fixnum_cmp(jit, ctx, asm, ocb, Assembler::csel_g, BOP_GT)
 }
 
 // Implements specialized equality for either two fixnum or two strings