ZJIT: Add BoxBool and remove CCall from BasicObject#==

3 files changed, 84 insertions, 3 deletions

diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index 1179b9bf16..7f93d22601 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -405,6 +405,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::IsNil { val } => gen_isnil(asm, opnd!(val)),
         &Insn::IsMethodCfunc { val, cd, cfunc, state: _ } => gen_is_method_cfunc(jit, asm, opnd!(val), cd, cfunc),
         &Insn::IsBitEqual { left, right } => gen_is_bit_equal(asm, opnd!(left), opnd!(right)),
+        &Insn::BoxBool { val } => gen_box_bool(asm, opnd!(val)),
         Insn::Test { val } => gen_test(asm, opnd!(val)),
         Insn::GuardType { val, guard_type, state } => gen_guard_type(jit, asm, opnd!(val), *guard_type, &function.frame_state(*state)),
         Insn::GuardTypeNot { val, guard_type, state } => gen_guard_type_not(jit, asm, opnd!(val), *guard_type, &function.frame_state(*state)),
@@ -1553,6 +1554,11 @@ fn gen_is_bit_equal(asm: &mut Assembler, left: lir::Opnd, right: lir::Opnd) -> l
     asm.csel_e(Opnd::Imm(1), Opnd::Imm(0))
 }
 
+fn gen_box_bool(asm: &mut Assembler, val: lir::Opnd) -> lir::Opnd {
+    asm.test(val, val);
+    asm.csel_nz(Opnd::Value(Qtrue), Opnd::Value(Qfalse))
+}
+
 fn gen_anytostring(asm: &mut Assembler, val: lir::Opnd, str: lir::Opnd, state: &FrameState) -> lir::Opnd {
     gen_prepare_leaf_call_with_gc(asm, state);
 
diff --git a/zjit/src/cruby_methods.rs b/zjit/src/cruby_methods.rs
index b434cfe25a..e7be5ab445 100644
--- a/zjit/src/cruby_methods.rs
+++ b/zjit/src/cruby_methods.rs
@@ -217,7 +217,7 @@ pub fn init() -> Annotations {
     annotate!(rb_cNilClass, "nil?", inline_nilclass_nil_p);
     annotate!(rb_mKernel, "nil?", inline_kernel_nil_p);
     annotate!(rb_mKernel, "respond_to?", inline_kernel_respond_to_p);
-    annotate!(rb_cBasicObject, "==", types::BoolExact, no_gc, leaf, elidable);
+    annotate!(rb_cBasicObject, "==", inline_basic_object_eq, types::BoolExact, no_gc, leaf, elidable);
     annotate!(rb_cBasicObject, "!", types::BoolExact, no_gc, leaf, elidable);
     annotate!(rb_cBasicObject, "initialize", inline_basic_object_initialize);
     annotate!(rb_cInteger, "succ", inline_integer_succ);
@@ -361,6 +361,13 @@ fn inline_integer_xor(fun: &mut hir::Function, block: hir::BlockId, recv: hir::I
     None
 }
 
+fn inline_basic_object_eq(fun: &mut hir::Function, block: hir::BlockId, recv: hir::InsnId, args: &[hir::InsnId], _state: hir::InsnId) -> Option<hir::InsnId> {
+    let &[other] = args else { return None; };
+    let c_result = fun.push_insn(block, hir::Insn::IsBitEqual { left: recv, right: other });
+    let result = fun.push_insn(block, hir::Insn::BoxBool { val: c_result });
+    Some(result)
+}
+
 fn inline_basic_object_initialize(fun: &mut hir::Function, block: hir::BlockId, _recv: hir::InsnId, args: &[hir::InsnId], _state: hir::InsnId) -> Option<hir::InsnId> {
     if !args.is_empty() { return None; }
     let result = fun.push_insn(block, hir::Insn::Const { val: hir::Const::Value(Qnil) });
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index be8f060e67..6d0b120a37 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -610,6 +610,7 @@ pub enum Insn {
     IsMethodCfunc { val: InsnId, cd: *const rb_call_data, cfunc: *const u8, state: InsnId },
     /// Return C `true` if left == right
     IsBitEqual { left: InsnId, right: InsnId },
+    BoxBool { val: InsnId },
     // TODO(max): In iseq body types that are not ISEQ_TYPE_METHOD, rewrite to Constant false.
     Defined { op_type: usize, obj: VALUE, pushval: VALUE, v: InsnId, state: InsnId },
     GetConstantPath { ic: *const iseq_inline_constant_cache, state: InsnId },
@@ -997,6 +998,7 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
             Insn::IsNil { val } => { write!(f, "IsNil {val}") }
             Insn::IsMethodCfunc { val, cd, .. } => { write!(f, "IsMethodCFunc {val}, :{}", ruby_call_method_name(*cd)) }
             Insn::IsBitEqual { left, right } => write!(f, "IsBitEqual {left}, {right}"),
+            Insn::BoxBool { val } => write!(f, "BoxBool {val}"),
             Insn::Jump(target) => { write!(f, "Jump {target}") }
             Insn::IfTrue { val, target } => { write!(f, "IfTrue {val}, {target}") }
             Insn::IfFalse { val, target } => { write!(f, "IfFalse {val}, {target}") }
@@ -1623,6 +1625,7 @@ impl Function {
             &IsNil { val } => IsNil { val: find!(val) },
             &IsMethodCfunc { val, cd, cfunc, state } => IsMethodCfunc { val: find!(val), cd, cfunc, state },
             &IsBitEqual { left, right } => IsBitEqual { left: find!(left), right: find!(right) },
+            &BoxBool { val } => BoxBool { val: find!(val) },
             Jump(target) => Jump(find_branch_edge!(target)),
             &IfTrue { val, ref target } => IfTrue { val: find!(val), target: find_branch_edge!(target) },
             &IfFalse { val, ref target } => IfFalse { val: find!(val), target: find_branch_edge!(target) },
@@ -1810,6 +1813,7 @@ impl Function {
             Insn::IsNil { .. } => types::CBool,
             Insn::IsMethodCfunc { .. } => types::CBool,
             Insn::IsBitEqual { .. } => types::CBool,
+            Insn::BoxBool { .. } => types::BoolExact,
             Insn::StringCopy { .. } => types::StringExact,
             Insn::StringIntern { .. } => types::Symbol,
             Insn::StringConcat { .. } => types::StringExact,
@@ -3094,6 +3098,7 @@ impl Function {
             | &Insn::Return { val }
             | &Insn::Test { val }
             | &Insn::SetLocal { val, .. }
+            | &Insn::BoxBool { val }
             | &Insn::IsNil { val } =>
                 worklist.push_back(val),
             &Insn::SetGlobal { val, state, .. }
@@ -13084,7 +13089,7 @@ mod opt_tests {
     }
 
     #[test]
-    fn test_specialize_basic_object_eq_to_ccall() {
+    fn test_specialize_basic_object_eq() {
         eval("
             class C; end
             def test(a, b) = a == b
@@ -13106,14 +13111,77 @@ mod opt_tests {
           PatchPoint MethodRedefined(C@0x1000, ==@0x1008, cme:0x1010)
           PatchPoint NoSingletonClass(C@0x1000)
           v28:HeapObject[class_exact:C] = GuardType v11, HeapObject[class_exact:C]
+          v29:CBool = IsBitEqual v28, v12
+          v30:BoolExact = BoxBool v29
           IncrCounter inline_cfunc_optimized_send_count
-          v30:BoolExact = CCall ==@0x1038, v28, v12
           CheckInterrupts
           Return v30
         ");
     }
 
     #[test]
+    fn test_specialize_basic_object_eqq() {
+        eval("
+            class C; end
+            def test(a, b) = a === b
+
+            test(C.new, C.new)
+        ");
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:3:
+        bb0():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          v2:BasicObject = GetLocal l0, SP@5
+          v3:BasicObject = GetLocal l0, SP@4
+          Jump bb2(v1, v2, v3)
+        bb1(v6:BasicObject, v7:BasicObject, v8:BasicObject):
+          EntryPoint JIT(0)
+          Jump bb2(v6, v7, v8)
+        bb2(v10:BasicObject, v11:BasicObject, v12:BasicObject):
+          PatchPoint MethodRedefined(C@0x1000, ===@0x1008, cme:0x1010)
+          PatchPoint NoSingletonClass(C@0x1000)
+          v26:HeapObject[class_exact:C] = GuardType v11, HeapObject[class_exact:C]
+          PatchPoint MethodRedefined(C@0x1000, ==@0x1038, cme:0x1040)
+          PatchPoint NoSingletonClass(C@0x1000)
+          v30:CBool = IsBitEqual v26, v12
+          v31:BoolExact = BoxBool v30
+          IncrCounter inline_cfunc_optimized_send_count
+          CheckInterrupts
+          Return v31
+        ");
+    }
+
+    #[test]
+    fn test_specialize_nil_eq() {
+        eval("
+            def test(a, b) = a == b
+
+            test(nil, 5)
+        ");
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:2:
+        bb0():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          v2:BasicObject = GetLocal l0, SP@5
+          v3:BasicObject = GetLocal l0, SP@4
+          Jump bb2(v1, v2, v3)
+        bb1(v6:BasicObject, v7:BasicObject, v8:BasicObject):
+          EntryPoint JIT(0)
+          Jump bb2(v6, v7, v8)
+        bb2(v10:BasicObject, v11:BasicObject, v12:BasicObject):
+          PatchPoint MethodRedefined(NilClass@0x1000, ==@0x1008, cme:0x1010)
+          v27:NilClass = GuardType v11, NilClass
+          v28:CBool = IsBitEqual v27, v12
+          v29:BoolExact = BoxBool v28
+          IncrCounter inline_cfunc_optimized_send_count
+          CheckInterrupts
+          Return v29
+        ");
+    }
+
+    #[test]
     fn test_guard_fixnum_and_fixnum() {
         eval("
             def test(x, y) = x & y