YJIT: Reorder branches for Fixnum opt_case_dispatch (#6841)

* YJIT: Reorder branches for Fixnum opt_case_dispatch

Co-authored-by: Maxime Chevalier-Boisvert <maxime.chevalierboisvert@shopify.com>
Co-authored-by: Alan Wu <alansi.xingwu@shopify.com>

* YJIT: Don't support too large values

Co-authored-by: Maxime Chevalier-Boisvert <maxime.chevalierboisvert@shopify.com>
Co-authored-by: Alan Wu <alansi.xingwu@shopify.com>

3 files changed, 60 insertions, 5 deletions

diff --git a/yjit/bindgen/src/main.rs b/yjit/bindgen/src/main.rs
index 9afba864d2..ffcc148685 100644
--- a/yjit/bindgen/src/main.rs
+++ b/yjit/bindgen/src/main.rs
@@ -105,6 +105,7 @@ fn main() {
         .allowlist_function("rb_hash_aset")
         .allowlist_function("rb_hash_aref")
         .allowlist_function("rb_hash_bulk_insert")
+        .allowlist_function("rb_hash_stlike_lookup")
 
         // From include/ruby/internal/intern/array.h
         .allowlist_function("rb_ary_new_capa")
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index bb4091c5c2..5b7eb8f67e 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -1895,6 +1895,9 @@ pub const SEND_MAX_DEPTH: i32 = 5;
 // up to 20 different methods for send
 pub const SEND_MAX_CHAIN_DEPTH: i32 = 20;
 
+// up to 20 different offsets for case-when
+pub const CASE_WHEN_MAX_DEPTH: i32 = 20;
+
 // Codegen for setting an instance variable.
 // Preconditions:
 //   - receiver is in REG0
@@ -3126,10 +3129,10 @@ fn gen_opt_regexpmatch2(
 }
 
 fn gen_opt_case_dispatch(
-    _jit: &mut JITState,
+    jit: &mut JITState,
     ctx: &mut Context,
-    _asm: &mut Assembler,
-    _ocb: &mut OutlinedCb,
+    asm: &mut Assembler,
+    ocb: &mut OutlinedCb,
 ) -> CodegenStatus {
     // Normally this instruction would lookup the key in a hash and jump to an
     // offset based on that.
@@ -3138,10 +3141,54 @@ fn gen_opt_case_dispatch(
     // We'd hope that our jitted code will be sufficiently fast without the
     // hash lookup, at least for small hashes, but it's worth revisiting this
     // assumption in the future.
+    if !jit_at_current_insn(jit) {
+        defer_compilation(jit, ctx, asm, ocb);
+        return EndBlock;
+    }
+    let starting_context = *ctx;
 
-    ctx.stack_pop(1);
+    let case_hash = jit_get_arg(jit, 0);
+    let else_offset = jit_get_arg(jit, 1).as_u32();
+
+    // Try to reorder case/else branches so that ones that are actually used come first.
+    // Supporting only Fixnum for now so that the implementation can be an equality check.
+    let key_opnd = ctx.stack_pop(1);
+    let comptime_key = jit_peek_at_stack(jit, ctx, 0);
+    if comptime_key.fixnum_p() && comptime_key.0 <= u32::MAX.as_usize() {
+        if !assume_bop_not_redefined(jit, ocb, INTEGER_REDEFINED_OP_FLAG, BOP_EQQ) {
+            return CantCompile;
+        }
 
-    KeepCompiling // continue with the next instruction
+        // Check if the key is the same value
+        asm.cmp(key_opnd, comptime_key.into());
+        let side_exit = get_side_exit(jit, ocb, &starting_context);
+        jit_chain_guard(
+            JCC_JNE,
+            jit,
+            &starting_context,
+            asm,
+            ocb,
+            CASE_WHEN_MAX_DEPTH as i32,
+            side_exit,
+        );
+
+        // Get the offset for the compile-time key
+        let mut offset = 0;
+        unsafe { rb_hash_stlike_lookup(case_hash, comptime_key.0 as _, &mut offset) };
+        let jump_offset = if offset == 0 {
+            // NOTE: If we hit the else branch with various values, it could negatively impact the performance.
+            else_offset
+        } else {
+            (offset as u32) >> 1 // FIX2LONG
+        };
+
+        // Jump to the offset of case or else
+        let jump_block = BlockId { iseq: jit.iseq, idx: jit_next_insn_idx(jit) + jump_offset };
+        gen_direct_jump(jit, &ctx, jump_block, asm);
+        EndBlock
+    } else {
+        KeepCompiling // continue with === branches
+    }
 }
 
 fn gen_branchif_branch(
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index ea57bd69d8..671c0ad353 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@@ -1030,6 +1030,13 @@ extern "C" {
     pub fn rb_hash_resurrect(hash: VALUE) -> VALUE;
 }
 extern "C" {
+    pub fn rb_hash_stlike_lookup(
+        hash: VALUE,
+        key: st_data_t,
+        pval: *mut st_data_t,
+    ) -> ::std::os::raw::c_int;
+}
+extern "C" {
     pub fn rb_gvar_get(arg1: ID) -> VALUE;
 }
 extern "C" {