From fc0b2a8df2327b40a2f6667c1567b3f8264b58a5 Mon Sep 17 00:00:00 2001
From: Maxime Chevalier-Boisvert <maxime.chevalierboisvert@shopify.com>
Date: Fri, 4 Aug 2023 10:09:43 -0400
Subject: [PATCH] YJIT: guard for array_len >= num in expandarray (#8169)

Avoid generating long dispatch chains for all array lengths seen.
---
 yjit/src/backend/ir.rs |  6 ++++++
 yjit/src/codegen.rs    | 18 ++++--------------
 yjit/src/core.rs       |  6 ++++++
 3 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/yjit/src/backend/ir.rs b/yjit/src/backend/ir.rs
index 888922579d..51fc24677f 100644
--- a/yjit/src/backend/ir.rs
+++ b/yjit/src/backend/ir.rs
@@ -538,6 +538,7 @@ impl Insn {
     pub(super) fn target_mut(&mut self) -> Option<&mut Target> {
         match self {
             Insn::Jbe(target) |
+            Insn::Jb(target) |
             Insn::Je(target) |
             Insn::Jl(target) |
             Insn::Jmp(target) |
@@ -682,6 +683,7 @@ impl Insn {
     pub fn target(&self) -> Option<&Target> {
         match self {
             Insn::Jbe(target) |
+            Insn::Jb(target) |
             Insn::Je(target) |
             Insn::Jl(target) |
             Insn::Jmp(target) |
@@ -1821,6 +1823,10 @@ impl Assembler {
         self.push_insn(Insn::Jbe(target));
     }
 
+    pub fn jb(&mut self, target: Target) {
+        self.push_insn(Insn::Jb(target));
+    }
+
     pub fn je(&mut self, target: Target) {
         self.push_insn(Insn::Je(target));
     }
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index ee299574a7..721b2a9b07 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -241,8 +241,10 @@ pub enum JCCKinds {
     JCC_JNZ,
     JCC_JZ,
     JCC_JE,
+    JCC_JB,
     JCC_JBE,
     JCC_JNA,
+    JCC_JNAE,
 }
 
 #[inline(always)]
@@ -1535,8 +1537,6 @@ fn gen_expandarray(
     let array_reg = asm.load(array_opnd);
     let array_len_opnd = get_array_len(asm, array_reg);
 
-    /*
-    // FIXME: JCC_JB not implemented
     // Guard on the comptime/expected array length
     if comptime_len >= num {
         asm.comment(&format!("guard array length >= {}", num));
@@ -1549,6 +1549,7 @@ fn gen_expandarray(
             OPT_AREF_MAX_CHAIN_DEPTH,
             Counter::expandarray_chain_max_depth,
         );
+
     } else {
         asm.comment(&format!("guard array length == {}", comptime_len));
         asm.cmp(array_len_opnd, comptime_len.into());
@@ -1561,18 +1562,6 @@ fn gen_expandarray(
             Counter::expandarray_chain_max_depth,
         );
     }
-    */
-
-    asm.comment(&format!("guard array length == {}", comptime_len));
-    asm.cmp(array_len_opnd, comptime_len.into());
-    jit_chain_guard(
-        JCC_JNE,
-        jit,
-        asm,
-        ocb,
-        OPT_AREF_MAX_CHAIN_DEPTH,
-        Counter::expandarray_chain_max_depth,
-    );
 
     let array_opnd = asm.stack_pop(1); // pop after using the type info
 
@@ -1908,6 +1897,7 @@ fn jit_chain_guard(
         JCC_JNE | JCC_JNZ => BranchGenFn::JNZToTarget0,
         JCC_JZ | JCC_JE => BranchGenFn::JZToTarget0,
         JCC_JBE | JCC_JNA => BranchGenFn::JBEToTarget0,
+        JCC_JB | JCC_JNAE => BranchGenFn::JBToTarget0,
     };
 
     if (asm.ctx.get_chain_depth() as i32) < depth_limit {
diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index 97a84c6306..32324deb97 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -474,6 +474,7 @@ pub enum BranchGenFn {
     JNZToTarget0,
     JZToTarget0,
     JBEToTarget0,
+    JBToTarget0,
     JITReturn,
 }
 
@@ -527,6 +528,9 @@ impl BranchGenFn {
             BranchGenFn::JBEToTarget0 => {
                 asm.jbe(target0)
             }
+            BranchGenFn::JBToTarget0 => {
+                asm.jb(target0)
+            }
             BranchGenFn::JITReturn => {
                 asm.comment("update cfp->jit_return");
                 asm.mov(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_JIT_RETURN), Opnd::const_ptr(target0.unwrap_code_ptr().raw_ptr()));
@@ -543,6 +547,7 @@ impl BranchGenFn {
             BranchGenFn::JNZToTarget0 |
             BranchGenFn::JZToTarget0 |
             BranchGenFn::JBEToTarget0 |
+            BranchGenFn::JBToTarget0 |
             BranchGenFn::JITReturn => BranchShape::Default,
         }
     }
@@ -563,6 +568,7 @@ impl BranchGenFn {
             BranchGenFn::JNZToTarget0 |
             BranchGenFn::JZToTarget0 |
             BranchGenFn::JBEToTarget0 |
+            BranchGenFn::JBToTarget0 |
             BranchGenFn::JITReturn => {
                 assert_eq!(new_shape, BranchShape::Default);
             }