ZJIT: Load immediate into register before masking (#16677)

This otherwise causes a crash on e.g. https://github.com/ruby/ruby/pull/15220:

    thread '<unnamed>' (18071) panicked at zjit/src/codegen.rs:2511:21:
    internal error: entered unreachable code: with_num_bits should not be used for: Value(VALUE(19239180))
    stack backtrace:
       0: __rustc::rust_begin_unwind
                 at /rustc/e408947bfd200af42db322daf0fadfe7e26d3bd1/library/std/src/panicking.rs:689:5
       1: core::panicking::panic_fmt
                 at /rustc/e408947bfd200af42db322daf0fadfe7e26d3bd1/library/core/src/panicking.rs:80:14
       2: zjit::backend::lir::Opnd::with_num_bits
                 at /home/runner/work/ruby/ruby/src/zjit/src/backend/lir.rs:457:18
       3: zjit::codegen::gen_guard_type
                 at /home/runner/work/ruby/ruby/src/zjit/src/codegen.rs:2511:21
       4: zjit::codegen::gen_insn
                 at /home/runner/work/ruby/ruby/src/zjit/src/codegen.rs:690:55

2 files changed, 39 insertions, 4 deletions

diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index 5df020f4f5..da15d30d03 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -2458,8 +2458,9 @@ fn gen_has_type(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, ty: Typ
         asm.csel_e(Opnd::Imm(1), Opnd::Imm(0))
     } else if ty.is_subtype(types::StaticSymbol) {
         // Static symbols have (val & 0xff) == RUBY_SYMBOL_FLAG
-        // Use 8-bit comparison like YJIT does. GuardType should not be used
-        // for a known VALUE, which with_num_bits() does not support.
+        // Use 8-bit comparison like YJIT does.
+        // If `val` is a constant (rare but possible), put it in a register to allow masking.
+        let val = asm.load_imm(val);
         asm.cmp(val.with_num_bits(8), Opnd::UImm(RUBY_SYMBOL_FLAG as u64));
         asm.csel_e(Opnd::Imm(1), Opnd::Imm(0))
     } else if ty.is_subtype(types::NilClass) {
@@ -2528,8 +2529,9 @@ fn gen_guard_type(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, guard
         asm.jne(jit, side_exit(jit, state, GuardType(guard_type)));
     } else if guard_type.is_subtype(types::StaticSymbol) {
         // Static symbols have (val & 0xff) == RUBY_SYMBOL_FLAG
-        // Use 8-bit comparison like YJIT does. GuardType should not be used
-        // for a known VALUE, which with_num_bits() does not support.
+        // Use 8-bit comparison like YJIT does.
+        // If `val` is a constant (rare but possible), put it in a register to allow masking.
+        let val = asm.load_imm(val);
         asm.cmp(val.with_num_bits(8), Opnd::UImm(RUBY_SYMBOL_FLAG as u64));
         asm.jne(jit, side_exit(jit, state, GuardType(guard_type)));
     } else if guard_type.is_subtype(types::NilClass) {
@@ -3543,6 +3545,15 @@ impl Assembler {
         }
     }
 
+    /// Emits a load for constant based operands and returns a vreg,
+    /// otherwise returns recv.
+    fn load_imm(&mut self, recv: Opnd) -> Opnd {
+        match recv {
+            Opnd::Value { .. } | Opnd::UImm(_) | Opnd::Imm(_) => self.load(recv),
+            _ => recv,
+        }
+    }
+
     /// Make a C call while marking the start and end positions for IseqCall
     fn ccall_with_iseq_call(&mut self, fptr: *const u8, opnds: Vec<Opnd>, iseq_call: &IseqCallRef) -> Opnd {
         // We need to create our own branch rc objects so that we can move the closure below
diff --git a/zjit/src/codegen_tests.rs b/zjit/src/codegen_tests.rs
index d57efdc698..e19d365057 100644
--- a/zjit/src/codegen_tests.rs
+++ b/zjit/src/codegen_tests.rs
@@ -5586,3 +5586,27 @@ fn test_send_block_unused_warning_emitted_from_jit() {
         test
     "#), @"true");
 }
+
+#[test]
+fn test_load_immediates_into_registers_before_masking() {
+    // See https://github.com/ruby/ruby/pull/16669 -- this is a reduced reproduction from a Ruby
+    // spec.
+    set_call_threshold(2);
+    assert_snapshot!(inspect(r#"
+        def test
+          klass = Class.new do
+            def ===(o)
+              true
+            end
+          end
+
+          case 1
+          when klass.new
+            :called
+          end == :called
+        end
+
+        test
+        test
+    "#), @"true");
+}