Implement protected

4 files changed, 37 insertions, 6 deletions

diff --git a/lib/ruby_vm/mjit/assembler.rb b/lib/ruby_vm/mjit/assembler.rb
index 9e28f6cef3..20cad72a39 100644
--- a/lib/ruby_vm/mjit/assembler.rb
+++ b/lib/ruby_vm/mjit/assembler.rb
@@ -828,6 +828,15 @@ module RubyVM::MJIT
           opcode: 0x85,
           mod_rm: ModRM[mod: Mod11, reg: right_reg, rm: left_reg],
         )
+      # TEST r/m64, r64 (Mod 11: reg)
+      in [Symbol => left_reg, Symbol => right_reg] if r64?(left_reg) && r64?(right_reg)
+        # REX.W + 85 /r
+        # MR: Operand 1: ModRM:r/m (r), Operand 2: ModRM:reg (r)
+        insn(
+          prefix: REX_W,
+          opcode: 0x85,
+          mod_rm: ModRM[mod: Mod11, reg: right_reg, rm: left_reg],
+        )
       else
         raise NotImplementedError, "test: not-implemented operands: #{left.inspect}, #{right.inspect}"
       end
diff --git a/lib/ruby_vm/mjit/insn_compiler.rb b/lib/ruby_vm/mjit/insn_compiler.rb
index dc7683078d..7c3d443dd2 100644
--- a/lib/ruby_vm/mjit/insn_compiler.rb
+++ b/lib/ruby_vm/mjit/insn_compiler.rb
@@ -2640,7 +2640,7 @@ module RubyVM::MJIT
       # and so only have to do this once at compile time this is fine to always
       # check and side exit.
       comptime_recv = jit.peek_at_stack(argc)
-      unless C.rb_obj_is_kind_of(comptime_recv, current_defined_class)
+      unless C.obj_is_kind_of(comptime_recv, current_defined_class)
         return CantCompile
       end
 
@@ -2707,8 +2707,11 @@ module RubyVM::MJIT
           return CantCompile
         end
       when C.METHOD_VISI_PROTECTED
-        asm.incr_counter(:send_protected)
-        return CantCompile # TODO: support this
+        # If the method call is an FCALL, it is always valid
+        if flags & C.VM_CALL_FCALL == 0
+          # otherwise we need an ancestry check to ensure the receiver is valid to be called as protected
+          jit_protected_callee_ancestry_guard(asm, cme, side_exit(jit, ctx))
+        end
       else
         # TODO: Change them to a constant and use case-in instead
         raise 'unreachable'
@@ -2723,6 +2726,21 @@ module RubyVM::MJIT
       jit_call_method_each_type(jit, ctx, asm, argc, flags, cme, comptime_recv, recv_opnd, block_handler, send_shift:)
     end
 
+    # Generate ancestry guard for protected callee.
+    # Calls to protected callees only go through when self.is_a?(klass_that_defines_the_callee).
+    def jit_protected_callee_ancestry_guard(asm, cme, side_exit)
+      # See vm_call_method().
+      def_class = cme.defined_class
+      # Note: PC isn't written to current control frame as rb_is_kind_of() shouldn't raise.
+      # VALUE rb_obj_is_kind_of(VALUE obj, VALUE klass);
+
+      asm.mov(C_ARGS[0], [CFP, C.rb_control_frame_t.offsetof(:self)])
+      asm.mov(C_ARGS[1], to_value(def_class))
+      asm.call(C.rb_obj_is_kind_of)
+      asm.test(C_RET, C_RET)
+      asm.jz(counted_exit(side_exit, :send_protected_check_failed))
+    end
+
     # vm_call_method_each_type
     # @param jit [RubyVM::MJIT::JITState]
     # @param ctx [RubyVM::MJIT::Context]
diff --git a/mjit_c.h b/mjit_c.h
index 4e01dff7b3..b2cf602c86 100644
--- a/mjit_c.h
+++ b/mjit_c.h
@@ -117,7 +117,7 @@ MJIT_RUNTIME_COUNTERS(
     send_kwarg,
     send_missing_cme,
     send_private,
-    send_protected,
+    send_protected_check_failed,
     send_tailcall,
     send_notimplemented,
     send_cfunc,
diff --git a/mjit_c.rb b/mjit_c.rb
index 0a75ed1650..7c10b70078 100644
--- a/mjit_c.rb
+++ b/mjit_c.rb
@@ -237,10 +237,14 @@ module RubyVM::MJIT # :nodoc: all
       Primitive.cexpr! 'ID2SYM((ID)NUM2SIZET(id))'
     end
 
-    def rb_obj_is_kind_of(obj, c)
+    def obj_is_kind_of(obj, c)
       Primitive.cexpr! 'rb_obj_is_kind_of(obj, c)'
     end
 
+    def rb_obj_is_kind_of
+      Primitive.cexpr! 'SIZET2NUM((size_t)rb_obj_is_kind_of)'
+    end
+
     def rb_vm_defined
       Primitive.cstmt! %{
         extern bool rb_vm_defined(rb_execution_context_t *ec, rb_control_frame_t *reg_cfp, rb_num_t op_type, VALUE obj, VALUE v);
@@ -1363,7 +1367,7 @@ module RubyVM::MJIT # :nodoc: all
       send_kwarg: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_kwarg)")],
       send_missing_cme: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_missing_cme)")],
       send_private: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_private)")],
-      send_protected: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_protected)")],
+      send_protected_check_failed: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_protected_check_failed)")],
       send_tailcall: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_tailcall)")],
       send_notimplemented: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_notimplemented)")],
       send_cfunc: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_cfunc)")],