ZJIT: Implement getblockparamproxy (ruby#14483)

Co-authored-by: Alan Wu <[email protected]>

Author: k0kubun

test/ruby/test_zjit.rb

@@ -308,6 +308,15 @@ def test_setlocal_on_eval
     }
   end
 
+  def test_getblockparamproxy
+    assert_compiles '1', %q{
+      def test(&block)
+        0.then(&block)
+      end
+      test { 1 }
+    }, insns: [:getblockparamproxy]
+  end
+
   def test_call_a_forwardable_method
     assert_runs '[]', %q{
       def test_root = forwardable

zjit/src/codegen.rs

@@ -398,6 +398,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::GetGlobal { id, state } => gen_getglobal(jit, asm, *id, &function.frame_state(*state)),
         &Insn::GetLocal { ep_offset, level } => gen_getlocal_with_ep(asm, ep_offset, level),
         &Insn::SetLocal { val, ep_offset, level } => no_output!(gen_setlocal_with_ep(asm, opnd!(val), function.type_of(val), ep_offset, level)),
+        &Insn::GetBlockParamProxy { level, state } => gen_get_block_param_proxy(jit, asm, level, &function.frame_state(state)),
         Insn::GetConstantPath { ic, state } => gen_get_constant_path(jit, asm, *ic, &function.frame_state(*state)),
         Insn::SetIvar { self_val, id, val, state: _ } => no_output!(gen_setivar(asm, opnd!(self_val), *id, opnd!(val))),
         Insn::SideExit { state, reason } => no_output!(gen_side_exit(jit, asm, reason, &function.frame_state(*state))),
@@ -548,6 +549,29 @@ fn gen_setlocal_with_ep(asm: &mut Assembler, val: Opnd, val_type: Type, local_ep
     }
 }
 
+fn gen_get_block_param_proxy(jit: &JITState, asm: &mut Assembler, level: u32, state: &FrameState) -> lir::Opnd {
+    // Bail out if the `&block` local variable has been modified
+    let ep = gen_get_ep(asm, level);
+    let flags = Opnd::mem(64, ep, SIZEOF_VALUE_I32 * (VM_ENV_DATA_INDEX_FLAGS as i32));
+    asm.test(flags, VM_FRAME_FLAG_MODIFIED_BLOCK_PARAM.into());
+    asm.jnz(side_exit(jit, state, SideExitReason::BlockParamProxyModified));
+
+    // This handles two cases which are nearly identical
+    // Block handler is a tagged pointer. Look at the tag.
+    //   VM_BH_ISEQ_BLOCK_P(): block_handler & 0x03 == 0x01
+    //   VM_BH_IFUNC_P():      block_handler & 0x03 == 0x03
+    // So to check for either of those cases we can use: val & 0x1 == 0x1
+    const _: () = assert!(RUBY_SYMBOL_FLAG & 1 == 0, "guard below rejects symbol block handlers");
+
+    // Bail ouf if the block handler is neither ISEQ nor ifunc
+    let block_handler = asm.load(Opnd::mem(64, ep, SIZEOF_VALUE_I32 * VM_ENV_DATA_INDEX_SPECVAL));
+    asm.test(block_handler, 0x1.into());
+    asm.jz(side_exit(jit, state, SideExitReason::BlockParamProxyNotIseqOrIfunc));
+
+    // Return the rb_block_param_proxy instance (GC root, so put as a number to avoid unnecessary GC tracing)
+    unsafe { rb_block_param_proxy }.as_u64().into()
+}
+
 fn gen_get_constant_path(jit: &JITState, asm: &mut Assembler, ic: *const iseq_inline_constant_cache, state: &FrameState) -> Opnd {
     unsafe extern "C" {
         fn rb_vm_opt_getconstant_path(ec: EcPtr, cfp: CfpPtr, ic: *const iseq_inline_constant_cache) -> VALUE;
@@ -1623,12 +1647,12 @@ fn compile_iseq(iseq: IseqPtr) -> Result<Function, CompileError> {
 }
 
 /// Build a Target::SideExit for non-PatchPoint instructions
-fn side_exit(jit: &mut JITState, state: &FrameState, reason: SideExitReason) -> Target {
+fn side_exit(jit: &JITState, state: &FrameState, reason: SideExitReason) -> Target {
     build_side_exit(jit, state, reason, None)
 }
 
 /// Build a Target::SideExit out of a FrameState
-fn build_side_exit(jit: &mut JITState, state: &FrameState, reason: SideExitReason, label: Option<Label>) -> Target {
+fn build_side_exit(jit: &JITState, state: &FrameState, reason: SideExitReason, label: Option<Label>) -> Target {
     let mut stack = Vec::new();
     for &insn_id in state.stack() {
         stack.push(jit.get_opnd(insn_id));

zjit/src/hir.rs

@@ -462,6 +462,8 @@ pub enum SideExitReason {
     CalleeSideExit,
     ObjToStringFallback,
     Interrupt,
+    BlockParamProxyModified,
+    BlockParamProxyNotIseqOrIfunc,
 }
 
 impl std::fmt::Display for SideExitReason {
@@ -556,6 +558,9 @@ pub enum Insn {
     GetLocal { level: u32, ep_offset: u32 },
     /// Set a local variable in a higher scope or the heap
     SetLocal { level: u32, ep_offset: u32, val: InsnId },
+    /// Get a special singleton instance `rb_block_param_proxy` if the block
+    /// handler for the EP specified by `level` is an ISEQ or an ifunc.
+    GetBlockParamProxy { level: u32, state: InsnId },
     GetSpecialSymbol { symbol_type: SpecialBackrefSymbol, state: InsnId },
     GetSpecialNumber { nth: u64, state: InsnId },
 
@@ -900,6 +905,7 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
             Insn::SetGlobal { id, val, .. } => write!(f, "SetGlobal :{}, {val}", id.contents_lossy()),
             Insn::GetLocal { level, ep_offset } => write!(f, "GetLocal l{level}, EP@{ep_offset}"),
             Insn::SetLocal { val, level, ep_offset } => write!(f, "SetLocal l{level}, EP@{ep_offset}, {val}"),
+            Insn::GetBlockParamProxy { level, .. } => write!(f, "GetBlockParamProxy l{level}"),
             Insn::GetSpecialSymbol { symbol_type, .. } => write!(f, "GetSpecialSymbol {symbol_type:?}"),
             Insn::GetSpecialNumber { nth, .. } => write!(f, "GetSpecialNumber {nth}"),
             Insn::ToArray { val, .. } => write!(f, "ToArray {val}"),
@@ -1361,6 +1367,7 @@ impl Function {
             &NewRange { low, high, flag, state } => NewRange { low: find!(low), high: find!(high), flag, state: find!(state) },
             &NewRangeFixnum { low, high, flag, state } => NewRangeFixnum { low: find!(low), high: find!(high), flag, state: find!(state) },
             &ArrayMax { ref elements, state } => ArrayMax { elements: find_vec!(elements), state: find!(state) },
+            &GetBlockParamProxy { level, state } => GetBlockParamProxy { level, state: find!(state) },
             &SetGlobal { id, val, state } => SetGlobal { id, val: find!(val), state },
             &GetIvar { self_val, id, state } => GetIvar { self_val: find!(self_val), id, state },
             &LoadIvarEmbedded { self_val, id, index } => LoadIvarEmbedded { self_val: find!(self_val), id, index },
@@ -1472,6 +1479,7 @@ impl Function {
             Insn::ObjToString { .. } => types::BasicObject,
             Insn::AnyToString { .. } => types::String,
             Insn::GetLocal { .. } => types::BasicObject,
+            Insn::GetBlockParamProxy { .. } => types::BasicObject,
             // The type of Snapshot doesn't really matter; it's never materialized. It's used only
             // as a reference for FrameState, which we use to generate side-exit code.
             Insn::Snapshot { .. } => types::Any,
@@ -2300,6 +2308,7 @@ impl Function {
             | &Insn::LoadIvarExtended { self_val, .. } => {
                 worklist.push_back(self_val);
             }
+            &Insn::GetBlockParamProxy { state, .. } |
             &Insn::GetGlobal { state, .. } |
             &Insn::GetSpecialSymbol { state, .. } |
             &Insn::GetSpecialNumber { state, .. } |
@@ -3466,6 +3475,11 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     let level = get_arg(pc, 1).as_u32();
                     fun.push_insn(block, Insn::SetLocal { val: state.stack_pop()?, ep_offset, level });
                 }
+                YARVINSN_getblockparamproxy => {
+                    let level = get_arg(pc, 1).as_u32();
+                    let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
+                    state.stack_push(fun.push_insn(block, Insn::GetBlockParamProxy { level, state: exit_id }));
+                }
                 YARVINSN_pop => { state.stack_pop()?; }
                 YARVINSN_dup => { state.stack_push(state.stack_top()?); }
                 YARVINSN_dupn => {
@@ -5816,7 +5830,16 @@ mod tests {
           v5:NilClass = Const Value(nil)
           v10:BasicObject = InvokeBuiltin dir_s_open, v0, v1, v2
           PatchPoint NoEPEscape(open)
-          SideExit UnhandledYARVInsn(getblockparamproxy)
+          v16:BasicObject = GetBlockParamProxy l0
+          CheckInterrupts
+          v19:CBool = Test v16
+          IfFalse v19, bb1(v0, v1, v2, v3, v4, v10)
+          PatchPoint NoEPEscape(open)
+          SideExit UnhandledYARVInsn(invokeblock)
+        bb1(v27:BasicObject, v28:BasicObject, v29:BasicObject, v30:BasicObject, v31:BasicObject, v32:BasicObject):
+          PatchPoint NoEPEscape(open)
+          CheckInterrupts
+          Return v32
         ");
     }
 
@@ -7980,6 +8003,19 @@ mod opt_tests {
         ");
     }
 
+    #[test]
+    fn test_getblockparamproxy() {
+        eval("
+            def test(&block) = tap(&block)
+        ");
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:2:
+        bb0(v0:BasicObject, v1:BasicObject):
+          v6:BasicObject = GetBlockParamProxy l0
+          SideExit UnhandledCallType(BlockArg)
+        ");
+    }
+
     #[test]
     fn test_getinstancevariable() {
         eval("

zjit/src/stats.rs

@@ -103,6 +103,8 @@ make_counters! {
         exit_obj_to_string_fallback,
         exit_interrupt,
         exit_optional_arguments,
+        exit_block_param_proxy_modified,
+        exit_block_param_proxy_not_iseq_or_ifunc,
     }
 
     // unhanded_call_: Unhandled call types
@@ -219,22 +221,24 @@ pub fn exit_counter_ptr(reason: crate::hir::SideExitReason) -> *mut u64 {
     use crate::hir::SideExitReason::*;
     use crate::stats::Counter::*;
     let counter = match reason {
-        UnknownNewarraySend(_)    => exit_unknown_newarray_send,
-        UnhandledCallType(_)      => exit_unhandled_call_type,
-        UnknownSpecialVariable(_) => exit_unknown_special_variable,
-        UnhandledHIRInsn(_)       => exit_unhandled_hir_insn,
-        UnhandledYARVInsn(_)      => exit_unhandled_yarv_insn,
-        FixnumAddOverflow         => exit_fixnum_add_overflow,
-        FixnumSubOverflow         => exit_fixnum_sub_overflow,
-        FixnumMultOverflow        => exit_fixnum_mult_overflow,
-        GuardType(_)              => exit_guard_type_failure,
-        GuardTypeNot(_)           => exit_guard_type_not_failure,
-        GuardBitEquals(_)         => exit_guard_bit_equals_failure,
-        GuardShape(_)             => exit_guard_shape_failure,
-        PatchPoint(_)             => exit_patchpoint,
-        CalleeSideExit            => exit_callee_side_exit,
-        ObjToStringFallback       => exit_obj_to_string_fallback,
-        Interrupt                 => exit_interrupt,
+        UnknownNewarraySend(_)        => exit_unknown_newarray_send,
+        UnhandledCallType(_)          => exit_unhandled_call_type,
+        UnknownSpecialVariable(_)     => exit_unknown_special_variable,
+        UnhandledHIRInsn(_)           => exit_unhandled_hir_insn,
+        UnhandledYARVInsn(_)          => exit_unhandled_yarv_insn,
+        FixnumAddOverflow             => exit_fixnum_add_overflow,
+        FixnumSubOverflow             => exit_fixnum_sub_overflow,
+        FixnumMultOverflow            => exit_fixnum_mult_overflow,
+        GuardType(_)                  => exit_guard_type_failure,
+        GuardTypeNot(_)               => exit_guard_type_not_failure,
+        GuardBitEquals(_)             => exit_guard_bit_equals_failure,
+        GuardShape(_)                 => exit_guard_shape_failure,
+        PatchPoint(_)                 => exit_patchpoint,
+        CalleeSideExit                => exit_callee_side_exit,
+        ObjToStringFallback           => exit_obj_to_string_fallback,
+        Interrupt                     => exit_interrupt,
+        BlockParamProxyModified       => exit_block_param_proxy_modified,
+        BlockParamProxyNotIseqOrIfunc => exit_block_param_proxy_not_iseq_or_ifunc,
     };
     counter_ptr(counter)
 }