ZJIT: Check if shape is too complex before reading ivar by index (ruby#15478)

This fixes a crash when the new shape after a transition  is too complex;
we need to check that it's not complex before trying to read by index.

Author: tekknolagi

zjit/src/hir.rs

@@ -3155,8 +3155,6 @@ impl Function {
                             // We're only looking at T_OBJECT so ignore all of the imemo stuff.
                             assert!(recv_type.flags().is_t_object());
                             next_shape_id = ShapeId(unsafe { rb_shape_transition_add_ivar_no_warnings(class, current_shape_id.0, id) });
-                            let ivar_result = unsafe { rb_shape_get_iv_index(next_shape_id.0, id, &mut ivar_index) };
-                            assert!(ivar_result, "New shape must have the ivar index");
                             // If the VM ran out of shapes, or this class generated too many leaf,
                             // it may be de-optimized into OBJ_TOO_COMPLEX_SHAPE (hash-table).
                             let new_shape_too_complex = unsafe { rb_jit_shape_too_complex_p(next_shape_id.0) };
@@ -3165,6 +3163,8 @@ impl Function {
                                 self.push_insn(block, Insn::IncrCounter(Counter::setivar_fallback_new_shape_too_complex));
                                 self.push_insn_id(block, insn_id); continue;
                             }
+                            let ivar_result = unsafe { rb_shape_get_iv_index(next_shape_id.0, id, &mut ivar_index) };
+                            assert!(ivar_result, "New shape must have the ivar index");
                             let current_capacity = unsafe { rb_jit_shape_capacity(current_shape_id.0) };
                             let next_capacity = unsafe { rb_jit_shape_capacity(next_shape_id.0) };
                             // If the new shape has a different capacity, or is TOO_COMPLEX, we'll have to

zjit/src/hir/opt_tests.rs

@@ -3932,6 +3932,38 @@ mod hir_opt_tests {
         ");
     }
 
+    #[test]
+    fn test_dont_specialize_setivar_when_next_shape_is_too_complex() {
+        eval(r#"
+            class AboutToBeTooComplex
+              def test = @abc = 5
+            end
+            SHAPE_MAX_VARIATIONS = 8  # see shape.h
+            SHAPE_MAX_VARIATIONS.times do
+              AboutToBeTooComplex.new.instance_variable_set(:"@a#{_1}", 1)
+            end
+            AboutToBeTooComplex.new.test
+            TEST = AboutToBeTooComplex.instance_method(:test)
+        "#);
+        assert_snapshot!(hir_string_proc("TEST"), @r"
+        fn test@<compiled>:3:
+        bb0():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          Jump bb2(v1)
+        bb1(v4:BasicObject):
+          EntryPoint JIT(0)
+          Jump bb2(v4)
+        bb2(v6:BasicObject):
+          v10:Fixnum[5] = Const Value(5)
+          PatchPoint SingleRactorMode
+          IncrCounter setivar_fallback_new_shape_too_complex
+          SetIvar v6, :@abc, v10
+          CheckInterrupts
+          Return v10
+        ");
+    }
+
     #[test]
     fn test_elide_freeze_with_frozen_hash() {
         eval("