ZJIT: Fold defined?(yield) for inlined methods (#17861)

* ZJIT: Fold defined?(yield) when inlining method with blockiseq

We know this statically.

* ZJIT: Fold negative case of defined?(yield) too

Leave handling non-blockiseq arguments for later (right now they should
never appear in SendDirect instructions).

zjit/src/hir.rs

@@ -5103,6 +5103,9 @@ impl Function {
                 let SendDirectData { recv, cme, iseq, args, kw_bits, jit_entry_idx, block: call_block, state, .. } = *data;
                 // SendDirect invariant: block is either None or BlockIseq.
                 // BlockArg is rejected upstream during type specialization.
+                // TODO(max): If we accept BlockArg here, we need to change the folding of Defined
+                // in HIR construction for the defined opcode to check the send flags of the method
+                // being inlined, too.
                 let blockiseq: Option<IseqPtr> = call_block.map(|bh| match bh {
                     BlockHandler::BlockIseq(bi) => bi,
                     BlockHandler::BlockArg => unreachable!("BlockArg in SendDirect"),
@@ -8239,18 +8242,31 @@ fn add_iseq_to_hir(
                         // Similar to gen_is_block_given
                         Insn::Const { val: Const::Value(Qnil) }
                     } else {
-                        // For DEFINED_YIELD, codegen materializes the local EP inline (similar to
-                        // gen_is_block_given) to check for a block handler. Precompute the lexical
-                        // distance from this iseq up to local_iseq so codegen does not have to
-                        // walk the parent chain. Any DEFINED_YIELD reaching this branch has a
-                        // method local_iseq by construction -- the above branch has already
-                        // diverted the non-method case to Qnil.
-                        let lep_level = if op_type == DEFINED_YIELD as usize {
-                            get_lvar_level(iseq)
+                        if op_type == DEFINED_YIELD as usize && matches!(mode, AddIseqMode::Inlined { .. }) {
+                            // If we are inlining a method that has a blockiseq handler, we can fold Defined(DEFINED_YIELD).
+                            // TODO(max): If we handle non-blockiseq block arguments such as
+                            // &:symbol or just &block forwarding, we need to revisit this and
+                            // check flags.
+                            let has_block = matches!(mode, AddIseqMode::Inlined { blockiseq: Some(_), .. });
+                            if has_block {
+                                Insn::Const { val: Const::Value(pushval) }
+                            } else {
+                                Insn::Const { val: Const::Value(Qnil) }
+                            }
                         } else {
-                            0
-                        };
-                        Insn::Defined { op_type, obj, pushval, v, lep_level, state: exit_id }
+                            // For DEFINED_YIELD, codegen materializes the local EP inline (similar to
+                            // gen_is_block_given) to check for a block handler. Precompute the lexical
+                            // distance from this iseq up to local_iseq so codegen does not have to
+                            // walk the parent chain. Any DEFINED_YIELD reaching this branch has a
+                            // method local_iseq by construction -- the above branch has already
+                            // diverted the non-method case to Qnil.
+                            let lep_level = if op_type == DEFINED_YIELD as usize {
+                                get_lvar_level(iseq)
+                            } else {
+                                0
+                            };
+                            Insn::Defined { op_type, obj, pushval, v, lep_level, state: exit_id }
+                        }
                     };
                     state.stack_push(fun.push_insn(block, insn));
                 }

zjit/src/hir/opt_tests.rs

@@ -17218,6 +17218,112 @@ mod hir_opt_tests {
         ");
     }
 
+    #[test]
+    fn test_inline_with_block_folds_defined_yield() {
+        eval(r"
+            def foo = defined?(yield)
+            def test = foo { |x| x }
+            test
+        ");
+        assert_snapshot!(hir_string("test"), @"
+        fn test@<compiled>:3:
+        bb1():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          Jump bb3(v1)
+        bb2():
+          EntryPoint JIT(0)
+          v4:BasicObject = LoadArg :self@0
+          Jump bb3(v4)
+        bb3(v6:BasicObject):
+          PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
+          v18:ObjectSubclass[class_exact*:Object@VALUE(0x1000)] = GuardType v6, ObjectSubclass[class_exact*:Object@VALUE(0x1000)] recompile
+          PushInlineFrame v18 (0x1038)
+          v27:StringExact[VALUE(0x1040)] = Const Value(VALUE(0x1040))
+          CheckInterrupts
+          PopInlineFrame
+          Return v27
+        ");
+    }
+
+    #[test]
+    fn test_inline_without_block_folds_defined_yield() {
+        eval(r"
+            def foo = defined?(yield)
+            def test = foo
+            test
+        ");
+        assert_snapshot!(hir_string("test"), @"
+        fn test@<compiled>:3:
+        bb1():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          Jump bb3(v1)
+        bb2():
+          EntryPoint JIT(0)
+          v4:BasicObject = LoadArg :self@0
+          Jump bb3(v4)
+        bb3(v6:BasicObject):
+          PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
+          v18:ObjectSubclass[class_exact*:Object@VALUE(0x1000)] = GuardType v6, ObjectSubclass[class_exact*:Object@VALUE(0x1000)] recompile
+          PushInlineFrame v18 (0x1038)
+          v27:NilClass = Const Value(nil)
+          CheckInterrupts
+          PopInlineFrame
+          Return v27
+        ");
+    }
+
+    #[test]
+    fn test_inline_array_each_with_block_folds_defined_yield() {
+        set_inline_threshold(100);
+        eval(r"
+            def test = [1, 2, 3].each { |x| x }
+            test
+        ");
+        assert_snapshot!(hir_string("test"), @"
+        fn test@<compiled>:2:
+        bb1():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          Jump bb3(v1)
+        bb2():
+          EntryPoint JIT(0)
+          v4:BasicObject = LoadArg :self@0
+          Jump bb3(v4)
+        bb3(v6:BasicObject):
+          v10:ArrayExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
+          v11:ArrayExact = ArrayDup v10
+          PatchPoint NoSingletonClass(Array@0x1008)
+          PatchPoint MethodRedefined(Array@0x1008, each@0x1010, cme:0x1018)
+          PushInlineFrame v11 (0x1040)
+          v51:Fixnum[0] = Const Value(0)
+          Jump bb10(v11, v51)
+        bb10(v64:ArrayExact, v65:Fixnum):
+          v69:CInt64 = ArrayLength v64
+          v70:Fixnum = BoxFixnum v69
+          v71:BoolExact = FixnumGe v65, v70
+          v73:CBool = Test v71
+          CondBranch v73, bb13(), bb9(v64, v65)
+        bb13():
+          CheckInterrupts
+          PopInlineFrame
+          Return v64
+        bb9(v86:ArrayExact, v87:Fixnum):
+          v92:CInt64 = UnboxFixnum v87
+          v93:BasicObject = ArrayAref v86, v92
+          v95:CPtr = GetEP 0
+          v96:CInt64 = LoadField v95, :VM_ENV_DATA_INDEX_SPECVAL@0x1048
+          v97:CInt64[-4] = Const CInt64(-4)
+          v98:CInt64 = IntAnd v96, v97
+          v99:BasicObject = InvokeBlockIseqDirect (0x1050), v98, v93
+          v103:Fixnum[1] = Const Value(1)
+          v104:Fixnum = FixnumAdd v87, v103
+          PatchPoint NoEPEscape(each)
+          Jump bb10(v86, v104)
+        ");
+    }
+
     #[test]
     fn test_delete_duplicate_store() {
         eval("