uJIT: add guards for protected opt_send_without_block calls

These account for about 12% of the time when we were bailing from calls
in railsbench. `ratio_in_ujit` went up 0.1% with this change.

2 files changed, 29 insertions, 8 deletions

diff --git a/ujit_codegen.c b/ujit_codegen.c
index 8249c1b1d5..4b28fa1dbd 100644
--- a/ujit_codegen.c
+++ b/ujit_codegen.c
@@ -1120,7 +1120,23 @@ gen_jump(jitstate_t* jit, ctx_t* ctx)
     return UJIT_END_BLOCK;
 }
 
-static codegen_status_t
+static void
+jit_protected_guard(jitstate_t *jit, codeblock_t *cb, const rb_callable_method_entry_t *cme, uint8_t *side_exit)
+{
+    // Callee is protected. Generate ancestry guard.
+    // See vm_call_method().
+    ujit_save_regs(cb);
+    mov(cb, C_ARG_REGS[0], member_opnd(REG_CFP, rb_control_frame_t, self));
+    jit_mov_gc_ptr(jit, cb, C_ARG_REGS[1], 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);
+    call_ptr(cb, REG0, (void *)&rb_obj_is_kind_of);
+    ujit_load_regs(cb);
+    cmp(cb, RAX, imm_opnd(0));
+    jz_ptr(cb, COUNTED_EXIT(side_exit, oswb_se_protected_check_failed));
+}
+
+static bool
 gen_oswb_cfunc(jitstate_t* jit, ctx_t* ctx, struct rb_call_data * cd, const rb_callable_method_entry_t *cme, int32_t argc)
 {
     const rb_method_cfunc_t *cfunc = UNALIGNED_MEMBER_PTR(cme->def, body.cfunc);
@@ -1190,6 +1206,11 @@ gen_oswb_cfunc(jitstate_t* jit, ctx_t* ctx, struct rb_call_data * cd, const rb_c
     mov(cb, REG0, const_ptr_opnd(jit->pc + insn_len(BIN(opt_send_without_block))));
     mov(cb, mem_opnd(64, REG_CFP, offsetof(rb_control_frame_t, pc)), REG0);
 
+    if (METHOD_ENTRY_VISI(cme) == METHOD_VISI_PROTECTED) {
+        // Generate ancestry guard for protected callee.
+        jit_protected_guard(jit, cb, cme, side_exit);
+    }
+
     // If this function needs a Ruby stack frame
     if (cfunc_needs_frame(cfunc))
     {
@@ -1416,6 +1437,12 @@ gen_oswb_iseq(jitstate_t* jit, ctx_t* ctx, struct rb_call_data * cd, const rb_ca
     cmp(cb, klass_opnd, REG1);
     jne_ptr(cb, COUNTED_EXIT(side_exit, oswb_se_cc_klass_differ));
 
+
+    if (METHOD_ENTRY_VISI(cme) == METHOD_VISI_PROTECTED) {
+        // Generate ancestry guard for protected callee.
+        jit_protected_guard(jit, cb, cme, side_exit);
+    }
+
     // Store the updated SP on the current frame (pop arguments and receiver)
     lea(cb, REG0, ctx_sp_opnd(ctx, sizeof(VALUE) * -(argc + 1)));
     mov(cb, member_opnd(REG_CFP, rb_control_frame_t, sp), REG0);
@@ -1563,12 +1590,6 @@ gen_opt_send_without_block(jitstate_t* jit, ctx_t* ctx)
         return false;
     }
 
-    // We don't generate code to check protected method calls
-    if (METHOD_ENTRY_VISI(cme) == METHOD_VISI_PROTECTED) {
-        GEN_COUNTER_INC(cb, oswb_protected);
-        return false;
-    }
-
     switch (cme->def->type) {
     case VM_METHOD_TYPE_ISEQ:
         return gen_oswb_iseq(jit, ctx, cd, cme, argc);
diff --git a/ujit_iface.h b/ujit_iface.h
index 2542e21590..e0addb3574 100644
--- a/ujit_iface.h
+++ b/ujit_iface.h
@@ -33,7 +33,6 @@ UJIT_DECLARE_COUNTERS(
     oswb_kw_splat,
     oswb_ic_empty,
     oswb_invalid_cme,
-    oswb_protected,
     oswb_ivar_set_method,
     oswb_ivar_get_method,
     oswb_zsuper_method,
@@ -54,6 +53,7 @@ UJIT_DECLARE_COUNTERS(
     oswb_se_receiver_not_heap,
     oswb_se_cf_overflow,
     oswb_se_cc_klass_differ,
+    oswb_se_protected_check_failed,
 
     // Member with known name for iterating over counters
     last_member