ZJIT: Compile opt_case_dispatch instead of the === chain

`opt_case_dispatch` looks a `case` value up in a compile-time hash and jumps
straight to the matching `when` body. ZJIT popped the key and fell through to
the `===` chain the compiler emits after it, which costs one `Integer#===`
send per `when` clause tested.

Worse, the interpreter never executes that chain -- it jumps out of
`opt_case_dispatch` on a hit -- so none of those sends are ever profiled.
Every one of them compiled to an unprofiled, un-inlined `CCallWithFrame` to
`Integer#===`. On the rubyboy benchmark that was 840M cfunc calls, 90% of all
calls to non-inlined C methods and 63% of all calls to C from JIT code, for a
`case addr >> 12` bus decoder and a ~500-way CPU opcode dispatch.

Compile the lookup instead. When every key in the dispatch hash is a Fixnum,
emit a binary search over the sorted keys that branches directly to each
`when` body, guarded by a `BOP_EQQ`/Integer patch point (the chain calls
`Integer#===`, so the lookup only agrees with it while that is the stock
implementation). Keys that are not Fixnums take the `===` chain as before, so
behavior is unchanged for them, and hashes with non-Fixnum keys are left
alone entirely.

Also annotate `Integer#===`, which is `rb_int_equal` just like `Integer#==`,
so the `===` chains that remain inline their comparisons.

A 256-way `case` over integer literals goes from 1816ms to 150ms (YJIT: 308ms).
On rubyboy, calls to C from JIT code drop from 1.32B to 909M and the total
executed instruction count drops 27%.

Author: k0kubun

zjit.c

@@ -147,6 +147,39 @@ rb_zjit_iseq_insn_set(const rb_iseq_t *iseq, unsigned int insn_idx, enum ruby_vm
     iseq->body->iseq_encoded[insn_idx] = (VALUE)insn_table[bare_insn];
 }
 
+struct zjit_cdhash_entries {
+    // Interleaved (key, offset) pairs, both as C longs
+    long *buf;
+    long capacity;
+    long size;
+};
+
+static int
+zjit_cdhash_fixnum_entries_i(st_data_t key, st_data_t val, st_data_t data)
+{
+    struct zjit_cdhash_entries *entries = (struct zjit_cdhash_entries *)data;
+    if (!FIXNUM_P((VALUE)key) || entries->size >= entries->capacity) {
+        entries->size = -1;
+        return ST_STOP;
+    }
+    entries->buf[entries->size * 2] = FIX2LONG((VALUE)key);
+    // Values are raw jump offsets, not Ruby objects. See cdhash_set_label_replace_i().
+    entries->buf[entries->size * 2 + 1] = (long)val;
+    entries->size++;
+    return ST_CONTINUE;
+}
+
+// Write the (key, jump offset) pairs of an `opt_case_dispatch` hash into `buf` as
+// interleaved C longs. Returns the number of pairs written, or -1 if any key is not
+// a Fixnum or if there is not enough room for all of them in `buf`.
+long
+rb_zjit_cdhash_fixnum_entries(VALUE cdhash, long *buf, long capacity)
+{
+    struct zjit_cdhash_entries entries = { .buf = buf, .capacity = capacity, .size = 0 };
+    st_foreach(rb_imemo_cdhash_tbl(cdhash), zjit_cdhash_fixnum_entries_i, (st_data_t)&entries);
+    return entries.size;
+}
+
 // Get profiling information for ISEQ
 void *
 rb_iseq_get_zjit_payload(const rb_iseq_t *iseq)

zjit/src/codegen_tests.rs

@@ -6241,6 +6241,48 @@ fn test_opt_case_dispatch() {
     assert_snapshot!(assert_compiles("[test(:foo), test(1)]"), @"[true, false]");
 }
 
+#[test]
+fn test_opt_case_dispatch_fixnum() {
+    eval("
+        def test(x)
+          case x
+          when -3 then :a
+          when 0, 1 then :b
+          when 5 then :c
+          else :d
+          end
+        end
+        test(0)
+    ");
+    assert_contains_opcode("test", YARVINSN_opt_case_dispatch);
+    assert_snapshot!(
+        assert_compiles("[-4, -3, 0, 1, 2, 5, 1.0, :x, nil].map { |x| test(x) }"),
+        @"[:d, :a, :b, :b, :d, :c, :b, :d, :d]"
+    );
+}
+
+#[test]
+fn test_opt_case_dispatch_fixnum_redefined() {
+    eval("
+        def test(x)
+          case x
+          when 0 then :a
+          when 1 then :b
+          else :c
+          end
+        end
+        test(0)
+    ");
+    assert_contains_opcode("test", YARVINSN_opt_case_dispatch);
+    assert_snapshot!(assert_compiles_allowing_exits("
+        [test(0), test(1), test(2)].tap {
+          class Integer
+            def ===(other) = true
+          end
+        } + [test(0), test(1), test(2)] + 100.times.map { test(2) }.uniq
+    "), @"[:a, :b, :c, :a, :a, :a, :a]");
+}
+
 #[test]
 fn test_checkmatch_case() {
     eval(r#"

zjit/src/cruby.rs

@@ -127,6 +127,7 @@ unsafe extern "C" {
     pub fn rb_float_new(d: f64) -> VALUE;
 
     pub fn rb_hash_empty_p(hash: VALUE) -> VALUE;
+    pub fn rb_zjit_cdhash_fixnum_entries(cdhash: VALUE, buf: *mut c_long, capacity: c_long) -> c_long;
     pub fn rb_ary_new_from_args(n: c_long, ...) -> VALUE;
     pub fn rb_str_setbyte(str: VALUE, index: VALUE, value: VALUE) -> VALUE;
     pub fn rb_str_getbyte(str: VALUE, index: VALUE) -> VALUE;

zjit/src/cruby_methods.rs

@@ -252,6 +252,10 @@ pub fn init() -> Annotations {
     annotate!(rb_cInteger, "succ", inline_integer_succ);
     annotate!(rb_cInteger, "^", inline_integer_xor);
     annotate!(rb_cInteger, "==", inline_integer_eq);
+    // Integer#=== is rb_int_equal, the same C function as Integer#==. `case`/`when`
+    // over integer literals compiles to a chain of `===` sends when opt_case_dispatch
+    // misses, so inlining this avoids a cfunc call per `when` clause tested.
+    annotate!(rb_cInteger, "===", inline_integer_eq);
     annotate!(rb_cInteger, "+", inline_integer_plus);
     annotate!(rb_cInteger, "-", inline_integer_minus);
     annotate!(rb_cInteger, "*", inline_integer_mult);

zjit/src/hir.rs

@@ -7922,6 +7922,29 @@ struct BytecodeInfo {
     jump_targets: Vec<u32>,
 }
 
+/// The largest `opt_case_dispatch` hash we are willing to turn into an inline
+/// binary search. Bigger `case`/`when` statements keep the `===` chain.
+const MAX_CASE_DISPATCH_ENTRIES: usize = 512;
+
+/// Read the `(key, jump offset)` pairs out of an `opt_case_dispatch` hash, sorted by key.
+/// Returns None unless every key is a Fixnum, which is what lets us compile the dispatch
+/// as an integer comparison tree guarded by a single `Integer#===` redefinition check.
+fn cdhash_fixnum_entries(cdhash: VALUE) -> Option<Vec<(i64, i64)>> {
+    let mut buf = vec![0 as std::os::raw::c_long; MAX_CASE_DISPATCH_ENTRIES * 2];
+    let size = unsafe {
+        rb_zjit_cdhash_fixnum_entries(cdhash, buf.as_mut_ptr(), MAX_CASE_DISPATCH_ENTRIES as std::os::raw::c_long)
+    };
+    if size <= 0 {
+        return None;
+    }
+    let mut entries: Vec<(i64, i64)> = buf[..(size as usize) * 2]
+        .chunks_exact(2)
+        .map(|pair| (pair[0] as i64, pair[1] as i64))
+        .collect();
+    entries.sort_unstable_by_key(|&(key, _)| key);
+    Some(entries)
+}
+
 fn compute_bytecode_info(iseq: *const rb_iseq_t, opt_table: &[u32]) -> BytecodeInfo {
     let iseq_size = unsafe { get_iseq_encoded_size(iseq) };
     let mut insn_idx = 0;
@@ -7952,6 +7975,16 @@ fn compute_bytecode_info(iseq: *const rb_iseq_t, opt_table: &[u32]) -> BytecodeI
                 let offset = get_arg(pc, 1).as_i64();
                 jump_targets.insert(insn_idx_at_offset(insn_idx, offset));
             }
+            YARVINSN_opt_case_dispatch => {
+                // The `when` bodies are already jump targets of the `===` chain that
+                // follows, but the else offset is only reachable by fallthrough there.
+                if let Some(entries) = cdhash_fixnum_entries(get_arg(pc, 0)) {
+                    for (_, offset) in entries {
+                        jump_targets.insert(insn_idx_at_offset(insn_idx, offset));
+                    }
+                    jump_targets.insert(insn_idx_at_offset(insn_idx, get_arg(pc, 1).as_i64()));
+                }
+            }
             YARVINSN_leave | YARVINSN_opt_invokebuiltin_delegate_leave => {
                 if insn_idx < iseq_size {
                     jump_targets.insert(insn_idx);
@@ -8930,10 +8963,77 @@ fn add_iseq_to_hir(
                     queue.push_back((state.clone(), target, target_idx, local_inval));
                 }
                 YARVINSN_opt_case_dispatch => {
-                    // TODO: Some keys are visible at compile time, so in the future we can
-                    // compile jump targets for certain cases
-                    // Pop the key from the stack and fallback to the === branches for now
-                    state.stack_pop()?;
+                    let key = state.stack_pop()?;
+                    // The interpreter jumps straight out of `opt_case_dispatch` on a hit, so
+                    // the `===` chain that follows is dead code there and never gets profiled.
+                    // Compiling the chain therefore means an unprofiled `Integer#===` cfunc
+                    // call per `when` clause tested. Compile the hash lookup instead, as a
+                    // binary search over the (Fixnum) keys.
+                    // The chain calls `Integer#===` on each `when` literal, so the lookup only
+                    // agrees with it while that stays the stock implementation.
+                    let unredefined = unsafe { rb_BASIC_OP_UNREDEFINED_P(BOP_EQQ, INTEGER_REDEFINED_OP_FLAG) };
+                    let Some(entries) = unredefined.then(|| cdhash_fixnum_entries(get_arg(pc, 0))).flatten() else {
+                        // Fall through to the `===` chain.
+                        continue;
+                    };
+                    fun.push_insn(block, Insn::PatchPoint {
+                        invariant: Invariant::BOPRedefined { klass: INTEGER_REDEFINED_OP_FLAG, bop: BOP_EQQ },
+                        state: exit_id,
+                    });
+                    // Only Fixnum keys can match an all-Fixnum dispatch hash. Anything else
+                    // falls through to the `===` chain, which handles every type correctly.
+                    let is_fixnum = fun.push_insn(block, Insn::HasType { val: key, expected: types::Fixnum });
+                    let chain_block = fun.new_block(insn_idx);
+                    let dispatch_block = fun.new_block(insn_idx);
+                    fun.push_insn(block, Insn::CondBranch {
+                        val: is_fixnum,
+                        if_true: BranchEdge { target: dispatch_block, args: vec![] },
+                        if_false: BranchEdge { target: chain_block, args: vec![] },
+                    });
+                    let key_fixnum = fun.push_insn(dispatch_block, Insn::RefineType { val: key, new_type: types::Fixnum });
+                    let mut state = state.clone();
+                    state.replace(key, key_fixnum);
+                    let else_idx = insn_idx_at_offset(insn_idx, get_arg(pc, 1).as_i64());
+                    let else_block = insn_idx_to_block[&else_idx];
+                    // Emit a comparison tree over the sorted keys. Each leaf range is scanned
+                    // linearly; anything bigger splits on a pivot key.
+                    let mut work = vec![(0usize, entries.len(), dispatch_block)];
+                    while let Some((lo, hi, mut cur)) = work.pop() {
+                        if hi - lo > 3 {
+                            let mid = lo + (hi - lo) / 2;
+                            let pivot = fun.push_insn(cur, Insn::Const { val: Const::Value(VALUE::fixnum_from_isize(entries[mid].0 as isize)) });
+                            let less = fun.push_insn(cur, Insn::FixnumLt { left: key_fixnum, right: pivot });
+                            let less_c = fun.push_insn(cur, Insn::Test { val: less });
+                            let lo_block = fun.new_block(insn_idx);
+                            let hi_block = fun.new_block(insn_idx);
+                            fun.push_insn(cur, Insn::CondBranch {
+                                val: less_c,
+                                if_true: BranchEdge { target: lo_block, args: vec![] },
+                                if_false: BranchEdge { target: hi_block, args: vec![] },
+                            });
+                            work.push((lo, mid, lo_block));
+                            work.push((mid, hi, hi_block));
+                            continue;
+                        }
+                        for &(key_value, offset) in &entries[lo..hi] {
+                            let target_idx = insn_idx_at_offset(insn_idx, offset);
+                            let target = insn_idx_to_block[&target_idx];
+                            let expected = fun.push_insn(cur, Insn::Const { val: Const::Value(VALUE::fixnum_from_isize(key_value as isize)) });
+                            let matches = fun.push_insn(cur, Insn::IsBitEqual { left: key_fixnum, right: expected });
+                            let next = fun.new_block(insn_idx);
+                            fun.push_insn(cur, Insn::CondBranch {
+                                val: matches,
+                                if_true: BranchEdge { target, args: state.as_args(self_param) },
+                                if_false: BranchEdge { target: next, args: vec![] },
+                            });
+                            queue.push_back((state.clone(), target, target_idx, local_inval));
+                            cur = next;
+                        }
+                        fun.push_insn(cur, Insn::Jump(BranchEdge { target: else_block, args: state.as_args(self_param) }));
+                        queue.push_back((state.clone(), else_block, else_idx, local_inval));
+                    }
+                    // Keep compiling the `===` chain for non-Fixnum keys.
+                    block = chain_block;
                 }
                 YARVINSN_opt_new => {
                     let cd: *const rb_call_data = get_arg(pc, 0).as_ptr();