From 4e40fdbcee51767de888c81ce4ad2baae639e35f Mon Sep 17 00:00:00 2001
From: Maxime Chevalier-Boisvert <maxime.chevalierboisvert@shopify.com>
Date: Thu, 22 Sep 2022 17:47:54 -0400
Subject: [PATCH] YJIT: add chain guards in `guard_two_fixnums` (#6422)

* Add chain guards in guard_two_fixnums, opt_eq with symbols

* Remove symbol comparison in gen_equality_specialized
---
 yjit/src/codegen.rs | 50 +++++++++++++++++++++++++++++++++------------
 1 file changed, 37 insertions(+), 13 deletions(-)

diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 15ff2a467f..c246c7b48f 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -1150,7 +1150,7 @@ fn gen_opt_plus(
         }
 
         // Check that both operands are fixnums
-        guard_two_fixnums(ctx, asm, side_exit);
+        guard_two_fixnums(jit, ctx, asm, ocb, side_exit);
 
         // Get the operands from the stack
         let arg1 = ctx.stack_pop(1);
@@ -2320,7 +2320,13 @@ fn gen_concatstrings(
     KeepCompiling
 }
 
-fn guard_two_fixnums(ctx: &mut Context, asm: &mut Assembler, side_exit: CodePtr) {
+fn guard_two_fixnums(
+    jit: &mut JITState,
+    ctx: &mut Context,
+    asm: &mut Assembler,
+    ocb: &mut OutlinedCb,
+    side_exit: CodePtr
+) {
     // Get the stack operand types
     let arg1_type = ctx.get_opnd_type(StackOpnd(0));
     let arg0_type = ctx.get_opnd_type(StackOpnd(1));
@@ -2352,16 +2358,34 @@ fn guard_two_fixnums(ctx: &mut Context, asm: &mut Assembler, side_exit: CodePtr)
     let arg1 = ctx.stack_opnd(0);
     let arg0 = ctx.stack_opnd(1);
 
-    // If not fixnums, fall back
+    // If not fixnums at run-time, fall back
     if arg0_type != Type::Fixnum {
         asm.comment("guard arg0 fixnum");
         asm.test(arg0, Opnd::UImm(RUBY_FIXNUM_FLAG as u64));
-        asm.jz(side_exit.into());
+
+        jit_chain_guard(
+            JCC_JZ,
+            jit,
+            &ctx,
+            asm,
+            ocb,
+            SEND_MAX_DEPTH,
+            side_exit,
+        );
     }
     if arg1_type != Type::Fixnum {
         asm.comment("guard arg1 fixnum");
         asm.test(arg1, Opnd::UImm(RUBY_FIXNUM_FLAG as u64));
-        asm.jz(side_exit.into());
+
+        jit_chain_guard(
+            JCC_JZ,
+            jit,
+            &ctx,
+            asm,
+            ocb,
+            SEND_MAX_DEPTH,
+            side_exit,
+        );
     }
 
     // Set stack types in context
@@ -2398,7 +2422,7 @@ fn gen_fixnum_cmp(
         }
 
         // Check that both operands are fixnums
-        guard_two_fixnums(ctx, asm, side_exit);
+        guard_two_fixnums(jit, ctx, asm, ocb, side_exit);
 
         // Get the operands from the stack
         let arg1 = ctx.stack_pop(1);
@@ -2475,7 +2499,7 @@ fn gen_equality_specialized(
             return false;
         }
 
-        guard_two_fixnums(ctx, asm, side_exit);
+        guard_two_fixnums(jit, ctx, asm, ocb, side_exit);
 
         asm.cmp(a_opnd, b_opnd);
 
@@ -2487,7 +2511,8 @@ fn gen_equality_specialized(
         asm.mov(dst, val);
 
         true
-    } else if unsafe { comptime_a.class_of() == rb_cString && comptime_b.class_of() == rb_cString }
+    }
+    else if unsafe { comptime_a.class_of() == rb_cString && comptime_b.class_of() == rb_cString }
     {
         if !assume_bop_not_redefined(jit, ocb, STRING_REDEFINED_OP_FLAG, BOP_EQ) {
             // if overridden, emit the generic version
@@ -2851,7 +2876,7 @@ fn gen_opt_and(
         }
 
         // Check that both operands are fixnums
-        guard_two_fixnums(ctx, asm, side_exit);
+        guard_two_fixnums(jit, ctx, asm, ocb, side_exit);
 
         // Get the operands and destination from the stack
         let arg1 = ctx.stack_pop(1);
@@ -2896,7 +2921,7 @@ fn gen_opt_or(
         }
 
         // Check that both operands are fixnums
-        guard_two_fixnums(ctx, asm, side_exit);
+        guard_two_fixnums(jit, ctx, asm, ocb, side_exit);
 
         // Get the operands and destination from the stack
         let arg1 = ctx.stack_pop(1);
@@ -2941,7 +2966,7 @@ fn gen_opt_minus(
         }
 
         // Check that both operands are fixnums
-        guard_two_fixnums(ctx, asm, side_exit);
+        guard_two_fixnums(jit, ctx, asm, ocb, side_exit);
 
         // Get the operands and destination from the stack
         let arg1 = ctx.stack_pop(1);
@@ -3008,7 +3033,7 @@ fn gen_opt_mod(
         }
 
         // Check that both operands are fixnums
-        guard_two_fixnums(ctx, asm, side_exit);
+        guard_two_fixnums(jit, ctx, asm, ocb, side_exit);
 
         // Get the operands and destination from the stack
         let arg1 = ctx.stack_pop(1);
@@ -5068,7 +5093,6 @@ fn gen_send_general(
         return CantCompile;
     }
 
-
     if flags & VM_CALL_ARGS_BLOCKARG != 0 {
         gen_counter_incr!(asm, send_block_arg);
         return CantCompile;