ZJIT: Compile modified path in getblockparamproxy (#16573)

ZJIT: Avoid side exits on modified getblockparamproxy

Author: nozomemein

zjit/src/codegen_tests.rs

@@ -420,6 +420,33 @@ fn test_getblockparamproxy() {
     assert_snapshot!(assert_compiles("test { 1 }"), @"1");
 }
 
+#[test]
+fn test_getblockparamproxy_modified() {
+    eval("
+        def test(&block)
+          b = block
+          0.then(&block)
+        end
+        test { 1 }
+    ");
+    assert_contains_opcode("test", YARVINSN_getblockparamproxy);
+    assert_snapshot!(inspect("test { 1 }"), @"1");
+}
+
+#[test]
+fn test_getblockparamproxy_modified_nested_block() {
+    eval("
+        def test(&block)
+          proc do
+            b = block
+            0.then(&block)
+          end
+        end
+        test { 1 }.call
+    ");
+    assert_snapshot!(inspect("test { 1 }.call"), @"1");
+}
+
 #[test]
 fn test_getblockparam() {
     eval("
@@ -462,9 +489,7 @@ fn test_setblockparam_nested_block() {
 }
 
 #[test]
-fn test_setblockparam_side_exit() {
-    // This pattern side exits because `block.call` goes through
-    // getblockparamproxy's modified-block-parameter case.
+fn test_getblockparamproxy_after_setblockparam() {
     eval("
         def test(&block)
           block = proc { 3 }
@@ -473,21 +498,7 @@ fn test_setblockparam_side_exit() {
         test { 1 }
     ");
     assert_contains_opcode("test", YARVINSN_setblockparam);
-    assert_snapshot!(inspect("test { 1 }"), @"3");
-}
-
-#[test]
-fn test_getblockparam_proxy_side_exit_restores_block_local() {
-    eval("
-        def test(&block)
-          b = block
-          raise \"test\" unless block
-          b ? 2 : 3
-        end
-        test {}
-    ");
-    assert_contains_opcode("test", YARVINSN_getblockparam);
-    assert_snapshot!(assert_compiles("test {}"), @"2");
+    assert_snapshot!(assert_compiles("test { 1 }"), @"3");
 }
 
 #[test]

zjit/src/hir.rs

@@ -524,7 +524,6 @@ pub enum SideExitReason {
     CalleeSideExit,
     ObjToStringFallback,
     Interrupt,
-    BlockParamProxyModified,
     BlockParamProxyNotIseqOrIfunc,
     BlockParamProxyNotNil,
     BlockParamWbRequired,
@@ -2611,6 +2610,24 @@ impl Function {
         id
     }
 
+    fn new_branch_block(
+        &mut self,
+        insn_idx: u32,
+        exit_state: &FrameState,
+        locals_count: usize,
+        stack_count: usize,
+    ) -> (BlockId, InsnId, FrameState, InsnId) {
+        let block = self.new_block(insn_idx);
+        let self_param = self.push_insn(block, Insn::Param);
+        let mut state = exit_state.clone();
+        state.locals.clear();
+        state.stack.clear();
+        state.locals.extend((0..locals_count).map(|_| self.push_insn(block, Insn::Param)));
+        state.stack.extend((0..stack_count).map(|_| self.push_insn(block, Insn::Param)));
+        let snapshot = self.push_insn(block, Insn::Snapshot { state: state.clone() });
+        (block, self_param, state, snapshot)
+    }
+
     fn remove_block(&mut self, block_id: BlockId) {
         if BlockId(self.blocks.len() - 1) != block_id {
             panic!("Can only remove the last block");
@@ -7578,7 +7595,9 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     });
                 }
                 YARVINSN_getblockparamproxy => {
+                    let ep_offset = get_arg(pc, 0).as_u32();
                     let level = get_arg(pc, 1).as_u32();
+                    let branch_insn_idx = exit_state.insn_idx as u32;
 
                     let profiled_block_type = if let Some([block_handler_distribution]) = profiles.payload.profile.get_operand_types(exit_state.insn_idx) {
                         let summary = TypeDistributionSummary::new(block_handler_distribution);
@@ -7587,16 +7606,44 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                         None
                     };
 
+                    let locals_count = state.locals.len();
+                    let stack_count = state.stack.len();
+                    let entry_args = state.as_args(self_param);
+
+                    // `getblockparamproxy` has two semantic paths:
+                    // - modified: return the already-materialized block local from EP
+                    // - unmodified: inspect the block handler and produce proxy/nil
+                    let (modified_block, modified_self_param, mut modified_state, ..) =
+                        fun.new_branch_block(branch_insn_idx, &exit_state, locals_count, stack_count);
+                    let (unmodified_block, unmodified_self_param, mut unmodified_state, unmodified_exit_id) =
+                        fun.new_branch_block(branch_insn_idx, &exit_state, locals_count, stack_count);
+                    let join_block = insn_idx_to_block.get(&insn_idx).copied().unwrap_or_else(|| fun.new_block(insn_idx));
+
                     let ep = fun.push_insn(block, Insn::GetEP { level });
-                    let flags = fun.push_insn(block, Insn::LoadField { recv: ep, id: ID!(_env_data_index_flags), offset: SIZEOF_VALUE_I32 * (VM_ENV_DATA_INDEX_FLAGS as i32), return_type: types::CInt64 });
-                    fun.push_insn(block, Insn::GuardNoBitsSet { val: flags, mask: Const::CUInt64(VM_FRAME_FLAG_MODIFIED_BLOCK_PARAM.into()), mask_name: Some(ID!(VM_FRAME_FLAG_MODIFIED_BLOCK_PARAM)), reason: SideExitReason::BlockParamProxyModified, state: exit_id });
+                    let is_modified = fun.push_insn(block, Insn::IsBlockParamModified { ep });
+
+                    fun.push_insn(block, Insn::IfTrue { val: is_modified, target: BranchEdge { target: modified_block, args: entry_args.clone() }});
+                    fun.push_insn(block, Insn::Jump(BranchEdge { target: unmodified_block, args: entry_args }));
+
+                    // Push modified block: load the block local via EP.
+                    let ep = fun.push_insn(modified_block, Insn::GetEP { level });
+                    let modified_val = fun.get_local_from_ep(modified_block, ep, ep_offset, level, types::BasicObject);
+                    if level == 0 {
+                        modified_state.setlocal(ep_offset, modified_val);
+                    }
+                    modified_state.stack_push(modified_val);
+                    fun.push_insn(modified_block, Insn::Jump(BranchEdge { target: join_block, args: modified_state.as_args(modified_self_param) }));
 
-                    let block_handler = fun.push_insn(block, Insn::LoadField { recv: ep, id: ID!(_env_data_index_specval), offset: SIZEOF_VALUE_I32 * VM_ENV_DATA_INDEX_SPECVAL, return_type: types::CInt64 });
+                    // Push unmodified block: inspect the current block handler to
+                    // decide whether this path returns `nil` or `BlockParamProxy`.
+                    let ep = fun.push_insn(unmodified_block, Insn::GetEP { level });
+                    let block_handler = fun.push_insn(unmodified_block, Insn::LoadField { recv: ep, id: ID!(_env_data_index_specval), offset: SIZEOF_VALUE_I32 * VM_ENV_DATA_INDEX_SPECVAL, return_type: types::CInt64 });
 
                     match profiled_block_type {
                         Some(ty) if ty.nil_p() => {
-                            fun.push_insn(block, Insn::GuardBitEquals { val: block_handler, expected: Const::CInt64(VM_BLOCK_HANDLER_NONE.into()), reason: SideExitReason::BlockParamProxyNotNil, state: exit_id });
-                            state.stack_push(fun.push_insn(block, Insn::Const { val: Const::Value(Qnil) }));
+                            fun.push_insn(unmodified_block, Insn::GuardBitEquals { val: block_handler, expected: Const::CInt64(VM_BLOCK_HANDLER_NONE.into()), reason: SideExitReason::BlockParamProxyNotNil, state: unmodified_exit_id });
+                            unmodified_state.stack_push(fun.push_insn(unmodified_block, Insn::Const { val: Const::Value(Qnil) }));
+                            fun.push_insn(unmodified_block, Insn::Jump(BranchEdge { target: join_block, args: unmodified_state.as_args(unmodified_self_param) }));
                         }
                         _ => {
                             // This handles two cases which are nearly identical
@@ -7607,31 +7654,19 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                             const _: () = assert!(RUBY_SYMBOL_FLAG & 1 == 0, "guard below rejects symbol block handlers");
 
                             // Bail out if the block handler is neither ISEQ nor ifunc
-                            fun.push_insn(block, Insn::GuardAnyBitSet { val: block_handler, mask: Const::CUInt64(0x1), mask_name: None, reason: SideExitReason::BlockParamProxyNotIseqOrIfunc, state: exit_id });
+                            fun.push_insn(unmodified_block, Insn::GuardAnyBitSet { val: block_handler, mask: Const::CUInt64(0x1), mask_name: None, reason: SideExitReason::BlockParamProxyNotIseqOrIfunc, state: unmodified_exit_id });
                             // TODO(Shopify/ruby#753): GC root, so we should be able to avoid unnecessary GC tracing
-                            state.stack_push(fun.push_insn(block, Insn::Const { val: Const::Value(unsafe { rb_block_param_proxy }) }));
+                            unmodified_state.stack_push(fun.push_insn(unmodified_block, Insn::Const { val: Const::Value(unsafe { rb_block_param_proxy }) }));
+                            fun.push_insn(unmodified_block, Insn::Jump(BranchEdge { target: join_block, args: unmodified_state.as_args(unmodified_self_param) }));
                         }
                     }
+
+                    // Continue compilation from the merged continuation block at the next
+                    // instruction.
+                    queue.push_back((unmodified_state, join_block, insn_idx, local_inval));
+                    break;
                 }
                 YARVINSN_getblockparam => {
-                    fn new_branch_block(
-                        fun: &mut Function,
-                        insn_idx: u32,
-                        exit_state: &FrameState,
-                        locals_count: usize,
-                        stack_count: usize,
-                    ) -> (BlockId, InsnId, FrameState, InsnId) {
-                        let block = fun.new_block(insn_idx);
-                        let self_param = fun.push_insn(block, Insn::Param);
-                        let mut state = exit_state.clone();
-                        state.locals.clear();
-                        state.stack.clear();
-                        state.locals.extend((0..locals_count).map(|_| fun.push_insn(block, Insn::Param)));
-                        state.stack.extend((0..stack_count).map(|_| fun.push_insn(block, Insn::Param)));
-                        let snapshot = fun.push_insn(block, Insn::Snapshot { state: state.clone() });
-                        (block, self_param, state, snapshot)
-                    }
-
                     fn finish_getblockparam_branch(
                         fun: &mut Function,
                         block: BlockId,
@@ -7667,9 +7702,9 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
 
                     // Set up branch and join blocks.
                     let (modified_block, modified_self_param, mut modified_state, ..) =
-                    new_branch_block(&mut fun, branch_insn_idx, &exit_state, locals_count, stack_count);
+                    fun.new_branch_block(branch_insn_idx, &exit_state, locals_count, stack_count);
                     let (unmodified_block, unmodified_self_param, mut unmodified_state, unmodified_exit_id) =
-                    new_branch_block(&mut fun, branch_insn_idx, &exit_state, locals_count, stack_count);
+                    fun.new_branch_block(branch_insn_idx, &exit_state, locals_count, stack_count);
                     let join_block = insn_idx_to_block.get(&insn_idx).copied().unwrap_or_else(|| fun.new_block(insn_idx));
 
                     fun.push_insn(block, Insn::IfTrue {