YJIT: Implement getblockparam

This implements the getblockparam instruction.

There are two cases we need to handle depending on whether or not
VM_FRAME_FLAG_MODIFIED_BLOCK_PARAM is set in the environment flag.

When the modified flag is unset, we need to call rb_vm_bh_to_procval to
get a proc from our passed block, save the proc in the environment, and
set the modified flag.

In the case that the modified flag is set we are able to just use the
existing proc in the environment.

One quirk of this is that we need to call jit_prepare_routine_call early
and ensure we update PC and SP regardless of the branch taken, so that
we have a consistent SP offset at the start of the next instruction.

We considered using a chain guard to generate these two paths
separately, but decided against it because it's very common to see both
and the modified case is basically a subset of the instructions in the
unmodified case.

This includes tests for both getblockparam and getblockparamproxy which
was previously missing a test.

3 files changed, 97 insertions, 0 deletions

diff --git a/yjit/bindgen/src/main.rs b/yjit/bindgen/src/main.rs
index 33f81366c2..7049688e1f 100644
--- a/yjit/bindgen/src/main.rs
+++ b/yjit/bindgen/src/main.rs
@@ -210,6 +210,7 @@ fn main() {
         .opaque_type("rb_execution_context_.*")
         .blocklist_type("rb_control_frame_struct")
         .opaque_type("rb_control_frame_struct")
+        .allowlist_function("rb_vm_bh_to_procval")
 
         // From yjit.c
         .allowlist_function("rb_iseq_(get|set)_yjit_payload")
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 60d8551a57..473d464d52 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -5573,6 +5573,98 @@ fn gen_getblockparamproxy(
     KeepCompiling
 }
 
+fn gen_getblockparam(
+    jit: &mut JITState,
+    ctx: &mut Context,
+    cb: &mut CodeBlock,
+    ocb: &mut OutlinedCb,
+) -> CodegenStatus {
+    // EP level
+    let level = jit_get_arg(jit, 1).as_u32();
+
+    // Save the PC and SP because we might allocate
+    jit_prepare_routine_call(jit, ctx, cb, REG0);
+
+    // A mirror of the interpreter code. Checking for the case
+    // where it's pushing rb_block_param_proxy.
+    let side_exit = get_side_exit(jit, ocb, ctx);
+
+    // Load environment pointer EP from CFP
+    gen_get_ep(cb, REG1, level);
+
+    // Bail when VM_ENV_FLAGS(ep, VM_FRAME_FLAG_MODIFIED_BLOCK_PARAM) is non zero
+    let flag_check = mem_opnd(
+        64,
+        REG1,
+        (SIZEOF_VALUE as i32) * (VM_ENV_DATA_INDEX_FLAGS as i32),
+    );
+    // FIXME: This is testing bits in the same place that the WB check is testing.
+    // We should combine these at some point
+    test(
+        cb,
+        flag_check,
+        uimm_opnd(VM_FRAME_FLAG_MODIFIED_BLOCK_PARAM.into()),
+    );
+
+    // If the frame flag has been modified, then the actual proc value is
+    // already in the EP and we should just use the value.
+    let frame_flag_modified = cb.new_label("frame_flag_modified".to_string());
+    jnz_label(cb, frame_flag_modified);
+
+    // This instruction writes the block handler to the EP.  If we need to
+    // fire a write barrier for the write, then exit (we'll let the
+    // interpreter handle it so it can fire the write barrier).
+    // flags & VM_ENV_FLAG_WB_REQUIRED
+    let flags_opnd = mem_opnd(
+        64,
+        REG1,
+        SIZEOF_VALUE as i32 * VM_ENV_DATA_INDEX_FLAGS as i32,
+    );
+    test(cb, flags_opnd, imm_opnd(VM_ENV_FLAG_WB_REQUIRED.into()));
+
+    // if (flags & VM_ENV_FLAG_WB_REQUIRED) != 0
+    jnz_ptr(cb, side_exit);
+
+    // Load the block handler for the current frame
+    // note, VM_ASSERT(VM_ENV_LOCAL_P(ep))
+    mov(
+        cb,
+        C_ARG_REGS[1],
+        mem_opnd(
+            64,
+            REG1,
+            (SIZEOF_VALUE as i32) * (VM_ENV_DATA_INDEX_SPECVAL as i32),
+        ),
+    );
+
+    // Convert the block handler in to a proc
+    // call rb_vm_bh_to_procval(const rb_execution_context_t *ec, VALUE block_handler)
+    mov(cb, C_ARG_REGS[0], REG_EC);
+    call_ptr(cb, REG0, rb_vm_bh_to_procval as *const u8);
+
+    // Load environment pointer EP from CFP (again)
+    gen_get_ep(cb, REG1, level);
+
+    // Set the frame modified flag
+    or(cb, flag_check, uimm_opnd(VM_FRAME_FLAG_MODIFIED_BLOCK_PARAM.into()));
+
+    // Write the value at the environment pointer
+    let idx = jit_get_arg(jit, 0).as_i32();
+    let offs = -(SIZEOF_VALUE as i32 * idx);
+    mov(cb, mem_opnd(64, REG1, offs), RAX);
+
+    cb.write_label(frame_flag_modified);
+
+    // Push the proc on the stack
+    let stack_ret = ctx.stack_push(Type::Unknown);
+    mov(cb, RAX, mem_opnd(64, REG1, offs));
+    mov(cb, stack_ret, RAX);
+
+    cb.link_labels();
+
+    KeepCompiling
+}
+
 fn gen_invokebuiltin(
     jit: &mut JITState,
     ctx: &mut Context,
@@ -5743,6 +5835,7 @@ fn get_gen_fn(opcode: VALUE) -> Option<InsnGenFn> {
         OP_JUMP => Some(gen_jump),
 
         OP_GETBLOCKPARAMPROXY => Some(gen_getblockparamproxy),
+        OP_GETBLOCKPARAM => Some(gen_getblockparam),
         OP_OPT_SEND_WITHOUT_BLOCK => Some(gen_opt_send_without_block),
         OP_SEND => Some(gen_send),
         OP_INVOKESUPER => Some(gen_invokesuper),
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index ac54ba4446..d1bba429d0 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@@ -589,6 +589,9 @@ pub const VM_ENV_FLAG_WB_REQUIRED: vm_frame_env_flags = 8;
 pub const VM_ENV_FLAG_ISOLATED: vm_frame_env_flags = 16;
 pub type vm_frame_env_flags = u32;
 extern "C" {
+    pub fn rb_vm_bh_to_procval(ec: *const rb_execution_context_t, block_handler: VALUE) -> VALUE;
+}
+extern "C" {
     pub fn rb_vm_frame_method_entry(
         cfp: *const rb_control_frame_t,
     ) -> *const rb_callable_method_entry_t;