ZJIT: Codegen for defined?(yield)

Lots of stdlib methods such as Integer#times and Kernel#then use this,
so at least this will make writing tests slightly easier.

jit.c

@@ -173,6 +173,12 @@ rb_get_iseq_body_local_iseq(const rb_iseq_t *iseq)
     return iseq->body->local_iseq;
 }
 
+const rb_iseq_t *
+rb_get_iseq_body_parent_iseq(const rb_iseq_t *iseq)
+{
+    return iseq->body->parent_iseq;
+}
+
 unsigned int
 rb_get_iseq_body_local_table_size(const rb_iseq_t *iseq)
 {

test/ruby/test_zjit.rb

@@ -781,6 +781,29 @@ class TestZJIT < Test::Unit::TestCase
     }, call_threshold: 2
   end
 
+  def test_defined_yield
+    assert_compiles "nil", "defined?(yield)"
+    assert_compiles '[nil, nil, "yield"]', %q{
+      def test = defined?(yield)
+      [test, test, test{}]
+    }, call_threshold: 2, insns: [:defined]
+  end
+
+  def test_defined_yield_from_block
+    # This will do some EP hopping to find the local EP,
+    # so it's slightly different than doing it outside of a block.
+
+    omit 'Test fails at the moment due to missing Send codegen'
+
+    assert_compiles '[nil, nil, "yield"]', %q{
+      def test
+        yield_self { yield_self { defined?(yield) } }
+      end
+
+      [test, test, test{}]
+    }, call_threshold: 2, insns: [:defined]
+  end
+
   def test_putspecialobject_vm_core_and_cbase
     assert_compiles '10', %q{
       def test

yjit.c

@@ -454,12 +454,6 @@ rb_get_def_bmethod_proc(rb_method_definition_t *def)
     return def->body.bmethod.proc;
 }
 
-const rb_iseq_t *
-rb_get_iseq_body_parent_iseq(const rb_iseq_t *iseq)
-{
-    return iseq->body->parent_iseq;
-}
-
 VALUE
 rb_optimized_call(VALUE *recv, rb_execution_context_t *ec, int argc, VALUE *argv, int kw_splat, VALUE block_handler)
 {

yjit/src/cruby_bindings.inc.rs

@@ -1189,7 +1189,6 @@ extern "C" {
     pub fn rb_yjit_get_proc_ptr(procv: VALUE) -> *mut rb_proc_t;
     pub fn rb_get_symbol_id(namep: VALUE) -> ID;
     pub fn rb_get_def_bmethod_proc(def: *mut rb_method_definition_t) -> VALUE;
-    pub fn rb_get_iseq_body_parent_iseq(iseq: *const rb_iseq_t) -> *const rb_iseq_t;
     pub fn rb_optimized_call(
         recv: *mut VALUE,
         ec: *mut rb_execution_context_t,
@@ -1286,6 +1285,7 @@ extern "C" {
     pub fn rb_get_mct_func(mct: *const rb_method_cfunc_t) -> *mut ::std::os::raw::c_void;
     pub fn rb_get_def_iseq_ptr(def: *mut rb_method_definition_t) -> *const rb_iseq_t;
     pub fn rb_get_iseq_body_local_iseq(iseq: *const rb_iseq_t) -> *const rb_iseq_t;
+    pub fn rb_get_iseq_body_parent_iseq(iseq: *const rb_iseq_t) -> *const rb_iseq_t;
     pub fn rb_get_iseq_body_local_table_size(iseq: *const rb_iseq_t) -> ::std::os::raw::c_uint;
     pub fn rb_get_iseq_body_iseq_encoded(iseq: *const rb_iseq_t) -> *mut VALUE;
     pub fn rb_get_iseq_body_stack_max(iseq: *const rb_iseq_t) -> ::std::os::raw::c_uint;

zjit/src/codegen.rs

@@ -287,6 +287,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::SideExit { state } => return gen_side_exit(jit, asm, &function.frame_state(*state)),
         Insn::PutSpecialObject { value_type } => gen_putspecialobject(asm, *value_type),
         Insn::AnyToString { val, str, state } => gen_anytostring(asm, opnd!(val), opnd!(str), &function.frame_state(*state))?,
+        Insn::Defined { op_type, obj, pushval, v } => gen_defined(jit, asm, *op_type, *obj, *pushval, opnd!(v))?,
         _ => {
             debug!("ZJIT: gen_function: unexpected insn {insn}");
             return None;
@@ -301,6 +302,25 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
     Some(())
 }
 
+/// Gets the EP of the ISeq of the containing method, or "local level".
+/// Equivalent of GET_LEP() macro.
+fn gen_get_lep(jit: &JITState, asm: &mut Assembler) -> Opnd {
+    // Equivalent of get_lvar_level() in compile.c
+    fn get_lvar_level(mut iseq: IseqPtr) -> u32 {
+        let local_iseq = unsafe { rb_get_iseq_body_local_iseq(iseq) };
+        let mut level = 0;
+        while iseq != local_iseq {
+            iseq = unsafe { rb_get_iseq_body_parent_iseq(iseq) };
+            level += 1;
+        }
+
+        level
+    }
+
+    let level = get_lvar_level(jit.iseq);
+    gen_get_ep(asm, level)
+}
+
 // Get EP at `level` from CFP
 fn gen_get_ep(asm: &mut Assembler, level: u32) -> Opnd {
     // Load environment pointer EP from CFP into a register
@@ -320,6 +340,27 @@ fn gen_get_ep(asm: &mut Assembler, level: u32) -> Opnd {
     ep_opnd
 }
 
+fn gen_defined(jit: &JITState, asm: &mut Assembler, op_type: usize, _obj: VALUE, pushval: VALUE, _tested_value: Opnd) -> Option<Opnd> {
+    match op_type as defined_type {
+        DEFINED_YIELD => {
+            // `yield` goes to the block handler stowed in the "local" iseq which is
+            // the current iseq or a parent. Only the "method" iseq type can be passed a
+            // block handler. (e.g. `yield` in the top level script is a syntax error.)
+            let local_iseq = unsafe { rb_get_iseq_body_local_iseq(jit.iseq) };
+            if unsafe { rb_get_iseq_body_type(local_iseq) } == ISEQ_TYPE_METHOD {
+                let lep = gen_get_lep(jit, asm);
+                let block_handler = asm.load(Opnd::mem(64, lep, SIZEOF_VALUE_I32 * VM_ENV_DATA_INDEX_SPECVAL));
+                let pushval = asm.load(pushval.into());
+                asm.cmp(block_handler, VM_BLOCK_HANDLER_NONE.into());
+                Some(asm.csel_e(Qnil.into(), pushval.into()))
+            } else {
+                Some(Qnil.into())
+            }
+        }
+        _ => None
+    }
+}
+
 /// Get a local variable from a higher scope. `local_ep_offset` is in number of VALUEs.
 fn gen_nested_getlocal(asm: &mut Assembler, local_ep_offset: u32, level: NonZeroU32) -> Option<lir::Opnd> {
     let ep = gen_get_ep(asm, level.get());

zjit/src/cruby_bindings.inc.rs

@@ -960,6 +960,7 @@ unsafe extern "C" {
     pub fn rb_get_mct_func(mct: *const rb_method_cfunc_t) -> *mut ::std::os::raw::c_void;
     pub fn rb_get_def_iseq_ptr(def: *mut rb_method_definition_t) -> *const rb_iseq_t;
     pub fn rb_get_iseq_body_local_iseq(iseq: *const rb_iseq_t) -> *const rb_iseq_t;
+    pub fn rb_get_iseq_body_parent_iseq(iseq: *const rb_iseq_t) -> *const rb_iseq_t;
     pub fn rb_get_iseq_body_local_table_size(iseq: *const rb_iseq_t) -> ::std::os::raw::c_uint;
     pub fn rb_get_iseq_body_iseq_encoded(iseq: *const rb_iseq_t) -> *mut VALUE;
     pub fn rb_get_iseq_body_stack_max(iseq: *const rb_iseq_t) -> ::std::os::raw::c_uint;