Support more types in jit_guard_known_class

3 files changed, 20 insertions, 23 deletions

diff --git a/lib/ruby_vm/mjit/insn_compiler.rb b/lib/ruby_vm/mjit/insn_compiler.rb
index da516a4476..e3f7d90759 100644
--- a/lib/ruby_vm/mjit/insn_compiler.rb
+++ b/lib/ruby_vm/mjit/insn_compiler.rb
@@ -765,7 +765,7 @@ module RubyVM::MJIT
     # @param asm [RubyVM::MJIT::Assembler]
     def jit_chain_guard(opcode, jit, ctx, asm, side_exit, limit: 10)
       case opcode
-      when :je, :jne, :jnz
+      when :je, :jne, :jnz, :jz
         # ok
       else
         raise ArgumentError, "jit_chain_guard: unexpected opcode #{opcode.inspect}"
@@ -803,18 +803,25 @@ module RubyVM::MJIT
     # @param ctx [RubyVM::MJIT::Context]
     # @param asm [RubyVM::MJIT::Assembler]
     def jit_guard_known_class(jit, ctx, asm, known_klass, obj_opnd, comptime_obj, side_exit, limit: 5)
+      # Only memory operand is supported for now
+      assert_equal(true, obj_opnd.is_a?(Array))
+
       if known_klass == NilClass
-        asm.incr_counter(:send_guard_nil)
-        return CantCompile
+        asm.comment('guard object is nil')
+        asm.cmp(obj_opnd, Qnil)
+        jit_chain_guard(:jne, jit, ctx, asm, side_exit, limit:)
       elsif known_klass == TrueClass
-        asm.incr_counter(:send_guard_true)
-        return CantCompile
+        asm.comment('guard object is true')
+        asm.cmp(obj_opnd, Qtrue)
+        jit_chain_guard(:jne, jit, ctx, asm, side_exit, limit:)
       elsif known_klass == FalseClass
-        asm.incr_counter(:send_guard_false)
-        return CantCompile
-      elsif known_klass == Integer
-        asm.incr_counter(:send_guard_integer)
-        return CantCompile
+        asm.comment('guard object is false')
+        asm.cmp(obj_opnd, Qfalse)
+        jit_chain_guard(:jne, jit, ctx, asm, side_exit, limit:)
+      elsif known_klass == Integer && fixnum?(comptime_obj)
+        asm.comment('guard object is fixnum')
+        asm.test(obj_opnd, C.RUBY_FIXNUM_FLAG)
+        jit_chain_guard(:jz, jit, ctx, asm, side_exit, limit:)
       elsif known_klass == Symbol
         asm.incr_counter(:send_guard_symbol)
         return CantCompile
@@ -827,11 +834,9 @@ module RubyVM::MJIT
         asm.cmp(obj_opnd, :rax)
         jit_chain_guard(:jne, jit, ctx, asm, side_exit, limit:)
       else
-        # If obj_opnd isn't already a register, load it.
-        if obj_opnd.is_a?(Array)
-          asm.mov(:rax, obj_opnd)
-          obj_opnd = :rax
-        end
+        # Load memory to a register
+        asm.mov(:rax, obj_opnd)
+        obj_opnd = :rax
 
         # Check that the receiver is a heap object
         # Note: if we get here, the class doesn't have immediate instances.
diff --git a/mjit_c.h b/mjit_c.h
index 181d9662f5..9e1386821b 100644
--- a/mjit_c.h
+++ b/mjit_c.h
@@ -133,10 +133,6 @@ MJIT_RUNTIME_COUNTERS(
     send_ivar_opt_send,
     send_ivar_blockarg,
 
-    send_guard_nil,
-    send_guard_true,
-    send_guard_false,
-    send_guard_integer,
     send_guard_symbol,
     send_guard_float,
 
diff --git a/mjit_c.rb b/mjit_c.rb
index 150000fd1e..509d5327bb 100644
--- a/mjit_c.rb
+++ b/mjit_c.rb
@@ -952,10 +952,6 @@ module RubyVM::MJIT # :nodoc: all
       send_ivar_arity: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_ivar_arity)")],
       send_ivar_opt_send: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_ivar_opt_send)")],
       send_ivar_blockarg: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_ivar_blockarg)")],
-      send_guard_nil: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_guard_nil)")],
-      send_guard_true: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_guard_true)")],
-      send_guard_false: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_guard_false)")],
-      send_guard_integer: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_guard_integer)")],
       send_guard_symbol: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_guard_symbol)")],
       send_guard_float: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_guard_float)")],
       getivar_megamorphic: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), getivar_megamorphic)")],