ZJIT: Recompile ISEQ on guard_shape_failure exits for setivar/definedivar (#16881)

* ZJIT: Recompile ISEQ on guard_shape_failure exits for setivar/definedivar

Profile self when setivar and definedivar shape guards fail, matching
getivar's recompilation behavior.

On the final ISEQ version, keep the generic SetIvar/DefinedIvar fallback
instead of emitting another shape guard that would keep side exiting.

* ZJIT: Clarify attr_writer SetIvar reprofile TODO

`ProfileSend { argc: 1 }` can identify the attr_writer receiver on the
stack, but the existing ProfileSend path is for no-profile sends. With
normal profiling settings, the send instruction has already finished
profiling by the time the lowered SetIvar shape guard fails, so
exit_recompile returns early and never recompiles into the generic
SetIvar fallback.

Update the TODO to make it clear that attr_writer SetIvar needs a
separate reprofile strategy for profiling the receiver operand after
send profiling has already finished.

* ZJIT: Add TODO comments for DefinedIvar/SetIvar follow-ups

We want to support polymorphic paths for both DefinedIvar and SetIvar in the near future.

2 files changed, 152 insertions, 12 deletions

diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index 2e72737ca1..0cced1281c 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -4412,9 +4412,15 @@ impl Function {
                             self.count(block, Counter::definedivar_fallback_complex);
                             self.push_insn_id(block, insn_id); continue;
                         }
+                        if self.policy.no_side_exits {
+                            // TODO: Support polymorphic DefinedIvar shape-specialized paths.
+                            // https://github.com/Shopify/ruby/issues/980
+                            // On the final version, keep the DefinedIvar fallback instead of another shape guard.
+                            self.push_insn_id(block, insn_id); continue;
+                        }
                         let self_val = self.load_ivar_guard_type(block, self_val, recv_type, state);
                         let shape = self.load_shape(block, self_val);
-                        self.guard_shape(block, shape, recv_type.shape(), state, None);
+                        self.guard_shape(block, shape, recv_type.shape(), state, Some(Recompile::ProfileSelf));
                         let mut ivar_index: attr_index_t = 0;
                         let replacement = if unsafe { rb_shape_get_iv_index(recv_type.shape().0, id, &mut ivar_index) } {
                             self.push_insn(block, Insn::Const { val: Const::Value(pushval) })
@@ -4426,7 +4432,7 @@ impl Function {
                         };
                         self.make_equal_to(insn_id, replacement);
                     }
-                    Insn::SetIvar { self_val, id, val, state, ic: _ } => {
+                    Insn::SetIvar { self_val, id, val, state, ic } => {
                         let frame_state = self.frame_state(state);
                         let Some(recv_type) = self.profiled_type_of_at(self_val, frame_state.insn_idx) else {
                             // No (monomorphic/skewed polymorphic) profile info
@@ -4454,6 +4460,12 @@ impl Function {
                             self.count(block, Counter::setivar_fallback_frozen);
                             self.push_insn_id(block, insn_id); continue;
                         }
+                        if self.policy.no_side_exits {
+                            // TODO: Support polymorphic SetIvar shape-specialized paths.
+                            // https://github.com/Shopify/ruby/issues/927
+                            // On the final version, keep the SetIvar fallback instead of another shape guard.
+                            self.push_insn_id(block, insn_id); continue;
+                        }
                         let mut ivar_index: attr_index_t = 0;
                         let mut next_shape_id = recv_type.shape();
                         if !unsafe { rb_shape_get_iv_index(recv_type.shape().0, id, &mut ivar_index) } {
@@ -4486,7 +4498,11 @@ impl Function {
                         }
                         let self_val = self.load_ivar_guard_type(block, self_val, recv_type, state);
                         let shape = self.load_shape(block, self_val);
-                        self.guard_shape(block, shape, recv_type.shape(), state, None);
+                        // TODO: attr_writer SetIvar has a null inline cache and may target a receiver
+                        // operand other than CFP self. Support it with a reprofile strategy that
+                        // profiles the receiver operand even after the send insn has finished profiling.
+                        let recompile = if ic.is_null() { None } else { Some(Recompile::ProfileSelf) };
+                        self.guard_shape(block, shape, recv_type.shape(), state, recompile);
                         // Current shape contains this ivar
                         let (ivar_storage, offset) = if recv_type.flags().is_embedded() {
                             // See ROBJECT_FIELDS() from include/ruby/internal/core/robject.h
diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs
index 0976ae659e..f1f7ceb76f 100644
--- a/zjit/src/hir/opt_tests.rs
+++ b/zjit/src/hir/opt_tests.rs
@@ -7625,6 +7625,136 @@ mod hir_opt_tests {
     }
 
     #[test]
+    fn test_definedivar_shape_guard_recompile() {
+        // Call with one shape to compile, then call with a different shape to
+        // trigger shape guard exits and recompilation. On the recompiled version,
+        // DefinedIvar stays as a C call fallback to avoid more shape guard exits.
+        eval("
+            class C
+              def initialize(extra = false)
+                @bar = 0 if extra  # changes the shape
+                @foo = 42
+              end
+              def has_foo = defined?(@foo)
+            end
+
+            c = C.new
+            c.has_foo  # profile
+            c.has_foo  # compile (version 1 with shape guard)
+            d = C.new(true)  # same class, different shape
+            100.times { d.has_foo }  # trigger shape guard exits -> recompile
+            100.times { c.has_foo }  # run recompiled version (version 2)
+        ");
+        assert_snapshot!(hir_string_proc("C.new.method(:has_foo)"), @"
+        fn has_foo@<compiled>:7:
+        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:StringExact|NilClass = DefinedIvar v6, :@foo
+          CheckInterrupts
+          Return v10
+        ");
+    }
+
+    #[test]
+    fn test_setivar_shape_guard_recompile() {
+        // Call with one shape to compile, then call with a different shape to
+        // trigger shape guard exits and recompilation. On the recompiled version,
+        // SetIvar stays as a C call fallback to avoid more shape guard exits.
+        eval("
+            class C
+              def initialize(extra = false)
+                @bar = 0 if extra  # changes the shape
+                @foo = 42
+              end
+              def foo = @foo = 5
+            end
+
+            c = C.new
+            c.foo  # profile
+            c.foo  # compile (version 1 with shape guard)
+            d = C.new(true)  # same class, different shape
+            100.times { d.foo }  # trigger shape guard exits -> recompile
+            100.times { c.foo }  # run recompiled version (version 2)
+        ");
+        assert_snapshot!(hir_string_proc("C.new.method(:foo)"), @"
+        fn foo@<compiled>:7:
+        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:Fixnum[5] = Const Value(5)
+          PatchPoint SingleRactorMode
+          SetIvar v6, :@foo, v10
+          CheckInterrupts
+          Return v10
+        ");
+    }
+
+    #[test]
+    fn test_setivar_shape_guard_attr_writer_no_recompile() {
+        // attr_writer SetIvar has no inline cache and may target a receiver
+        // operand other than CFP self, so don't recompile here yet.
+        eval("
+            class C
+              attr_writer :foo
+              def initialize(extra = false)
+                @bar = 0 if extra  # changes the shape
+                @foo = 42
+              end
+            end
+
+            class D
+              def write(obj)
+                obj.foo = 5
+              end
+            end
+
+            c = C.new
+            d = D.new
+            d.write(c)  # profile
+            d.write(c)  # compile (version 1 with shape guard)
+            e = C.new(true)  # same class, different shape
+            100.times { d.write(e) }  # shape guard exits, but no recompile
+        ");
+        assert_snapshot!(hir_string_proc("D.new.method(:write)"), @"
+        fn write@<compiled>:12:
+        bb1():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          v2:CPtr = LoadSP
+          v3:BasicObject = LoadField v2, :obj@0x1000
+          Jump bb3(v1, v3)
+        bb2():
+          EntryPoint JIT(0)
+          v6:BasicObject = LoadArg :self@0
+          v7:BasicObject = LoadArg :obj@1
+          Jump bb3(v6, v7)
+        bb3(v9:BasicObject, v10:BasicObject):
+          v17:Fixnum[5] = Const Value(5)
+          PatchPoint MethodRedefined(C@0x1008, foo=@0x1010, cme:0x1018)
+          v28:ObjectSubclass[class_exact:C] = GuardType v10, ObjectSubclass[class_exact:C]
+          v31:CShape = LoadField v28, :_shape_id@0x1040
+          v32:CShape[0x1041] = GuardBitEquals v31, CShape(0x1041)
+          StoreField v28, :@foo@0x1042, v17
+          WriteBarrier v28, v17
+          CheckInterrupts
+          Return v17
+        ");
+    }
+
+    #[test]
     fn test_optimize_getivar_on_module_embedded() {
         eval("
             module M
@@ -13546,23 +13676,17 @@ mod hir_opt_tests {
          CheckInterrupts
          SetLocal :formatted, l0, EP@3, v16
          PatchPoint SingleRactorMode
-         v60:HeapBasicObject = GuardType v15, HeapBasicObject
-         v61:CShape = LoadField v60, :_shape_id@0x1003
-         v62:CShape[0x1004] = GuardBitEquals v61, CShape(0x1004)
-         StoreField v60, :@formatted@0x1005, v16
-         WriteBarrier v60, v16
-         v65:CShape[0x1006] = Const CShape(0x1006)
-         StoreField v60, :_shape_id@0x1003, v65
+         SetIvar v15, :@formatted, v16
          v47:Class[VMFrozenCore] = Const Value(VALUE(0x1008))
          PatchPoint MethodRedefined(Class@0x1010, lambda@0x1018, cme:0x1020)
-         v69:BasicObject = CCallWithFrame v47, :RubyVM::FrozenCore.lambda@0x1048, block=0x1050
+         v62:BasicObject = CCallWithFrame v47, :RubyVM::FrozenCore.lambda@0x1048, block=0x1050
          v50:CPtr = GetEP 0
          v51:BasicObject = LoadField v50, :a@0x1001
          v52:BasicObject = LoadField v50, :_b@0x1002
          v53:BasicObject = LoadField v50, :_c@0x1058
          v54:BasicObject = LoadField v50, :formatted@0x1059
          CheckInterrupts
-         Return v69
+         Return v62
        ");
     }