ZJIT: Handle setivar with shape transition

tekknolagi · tekknolagi · commit 57f2cc27e55e · 2025-12-03T20:16:17.000-05:00
diff --git a/jit.c b/jit.c
@@ -767,3 +767,9 @@ rb_yarv_str_eql_internal(VALUE str1, VALUE str2)
 }
 
 void rb_jit_str_concat_codepoint(VALUE str, VALUE codepoint);
+
+attr_index_t
+rb_jit_shape_capacity(shape_id_t shape_id)
+{
+    return RSHAPE_CAPACITY(shape_id);
+}
diff --git a/zjit/bindgen/src/main.rs b/zjit/bindgen/src/main.rs
@@ -100,6 +100,7 @@ fn main() {
         .allowlist_function("rb_shape_id_offset")
         .allowlist_function("rb_shape_get_iv_index")
         .allowlist_function("rb_shape_transition_add_ivar_no_warnings")
+        .allowlist_function("rb_jit_shape_capacity")
         .allowlist_var("rb_invalid_shape_id")
         .allowlist_type("shape_id_fl_type")
         .allowlist_var("VM_KW_SPECIFIED_BITS_MAX")
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
@@ -350,6 +350,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         &Insn::Const { val: Const::CPtr(val) } => gen_const_cptr(val),
         &Insn::Const { val: Const::CInt64(val) } => gen_const_long(val),
         &Insn::Const { val: Const::CUInt16(val) } => gen_const_uint16(val),
+        &Insn::Const { val: Const::CUInt32(val) } => gen_const_uint32(val),
         Insn::Const { .. } => panic!("Unexpected Const in gen_insn: {insn}"),
         Insn::NewArray { elements, state } => gen_new_array(asm, opnds!(elements), &function.frame_state(*state)),
         Insn::NewHash { elements, state } => gen_new_hash(jit, asm, opnds!(elements), &function.frame_state(*state)),
@@ -475,7 +476,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::LoadEC => gen_load_ec(),
         Insn::LoadSelf => gen_load_self(),
         &Insn::LoadField { recv, id, offset, return_type: _ } => gen_load_field(asm, opnd!(recv), id, offset),
-        &Insn::StoreField { recv, id, offset, val } => no_output!(gen_store_field(asm, opnd!(recv), id, offset, opnd!(val))),
+        &Insn::StoreField { recv, id, offset, val } => no_output!(gen_store_field(asm, opnd!(recv), id, offset, opnd!(val), function.type_of(val))),
         &Insn::WriteBarrier { recv, val } => no_output!(gen_write_barrier(asm, opnd!(recv), opnd!(val), function.type_of(val))),
         &Insn::IsBlockGiven => gen_is_block_given(jit, asm),
         Insn::ArrayInclude { elements, target, state } => gen_array_include(jit, asm, opnds!(elements), opnd!(target), &function.frame_state(*state)),
@@ -1099,10 +1100,10 @@ fn gen_load_field(asm: &mut Assembler, recv: Opnd, id: ID, offset: i32) -> Opnd
     asm.load(Opnd::mem(64, recv, offset))
 }
 
-fn gen_store_field(asm: &mut Assembler, recv: Opnd, id: ID, offset: i32, val: Opnd) {
+fn gen_store_field(asm: &mut Assembler, recv: Opnd, id: ID, offset: i32, val: Opnd, val_type: Type) {
     asm_comment!(asm, "Store field id={} offset={}", id.contents_lossy(), offset);
     let recv = asm.load(recv);
-    asm.store(Opnd::mem(64, recv, offset), val);
+    asm.store(Opnd::mem(val_type.num_bits(), recv, offset), val);
 }
 
 fn gen_write_barrier(asm: &mut Assembler, recv: Opnd, val: Opnd, val_type: Type) {
@@ -1159,6 +1160,10 @@ fn gen_const_uint16(val: u16) -> lir::Opnd {
     Opnd::UImm(val as u64)
 }
 
+fn gen_const_uint32(val: u32) -> lir::Opnd {
+    Opnd::UImm(val as u64)
+}
+
 /// Compile a basic block argument
 fn gen_param(asm: &mut Assembler, idx: usize) -> lir::Opnd {
     // Allocate a register or a stack slot
diff --git a/zjit/src/cruby.rs b/zjit/src/cruby.rs
@@ -999,8 +999,9 @@ mod manual_defs {
     use super::*;
 
     pub const SIZEOF_VALUE: usize = 8;
+    pub const BITS_PER_BYTE: usize = 8;
     pub const SIZEOF_VALUE_I32: i32 = SIZEOF_VALUE as i32;
-    pub const VALUE_BITS: u8 = 8 * SIZEOF_VALUE as u8;
+    pub const VALUE_BITS: u8 = BITS_PER_BYTE as u8 * SIZEOF_VALUE as u8;
 
     pub const RUBY_LONG_MIN: isize = std::os::raw::c_long::MIN as isize;
     pub const RUBY_LONG_MAX: isize = std::os::raw::c_long::MAX as isize;
@@ -1382,6 +1383,7 @@ pub(crate) mod ids {
         name: thread_ptr
         name: self_              content: b"self"
         name: rb_ivar_get_at_no_ractor_check
+        name: _shape_id
     }
 
     /// Get an CRuby `ID` to an interned string, e.g. a particular method name.
diff --git a/zjit/src/cruby_bindings.inc.rs b/zjit/src/cruby_bindings.inc.rs
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
@@ -2947,10 +2947,34 @@ impl Function {
                             self.push_insn_id(block, insn_id); continue;
                         }
                         let mut ivar_index: u16 = 0;
+                        let mut next_shape_id = recv_type.shape();
                         if !unsafe { rb_shape_get_iv_index(recv_type.shape().0, id, &mut ivar_index) } {
-                            // TODO(max): Shape transition
-                            self.push_insn(block, Insn::IncrCounter(Counter::setivar_fallback_shape_transition));
-                            self.push_insn_id(block, insn_id); continue;
+                            // Current shape does not contain this ivar; do a shape transition.
+                            let current_shape_id = recv_type.shape();
+                            let class = recv_type.class();
+                            // 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) };
+                            // TODO(max): Is it OK to bail out here after making a shape transition?
+                            if new_shape_too_complex {
+                                self.push_insn(block, Insn::IncrCounter(Counter::setivar_fallback_new_shape_too_complex));
+                                self.push_insn_id(block, insn_id); continue;
+                            }
+                            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
+                            // reallocate it.
+                            let needs_extension = next_capacity != current_capacity;
+                            if needs_extension {
+                                self.push_insn(block, Insn::IncrCounter(Counter::setivar_fallback_new_shape_needs_extension));
+                                self.push_insn_id(block, insn_id); continue;
+                            }
+                            // Fall through to emitting the ivar write
                         }
                         let self_val = self.push_insn(block, Insn::GuardType { val: self_val, guard_type: types::HeapBasicObject, state });
                         let self_val = self.push_insn(block, Insn::GuardShape { val: self_val, shape: recv_type.shape(), state });
@@ -2966,6 +2990,13 @@ impl Function {
                         };
                         self.push_insn(block, Insn::StoreField { recv: ivar_storage, id, offset, val });
                         self.push_insn(block, Insn::WriteBarrier { recv: self_val, val });
+                        if next_shape_id != recv_type.shape() {
+                            // Write the new shape ID
+                            assert_eq!(SHAPE_ID_NUM_BITS, 32);
+                            let shape_id = self.push_insn(block, Insn::Const { val: Const::CUInt32(next_shape_id.0) });
+                            let shape_id_offset = unsafe { rb_shape_id_offset() };
+                            self.push_insn(block, Insn::StoreField { recv: self_val, id: ID!(_shape_id), offset: shape_id_offset, val: shape_id });
+                        }
                     }
                     _ => { self.push_insn_id(block, insn_id); }
                 }
diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs
@@ -3483,6 +3483,39 @@ mod hir_opt_tests {
         ");
     }
 
+    #[test]
+    fn test_dont_specialize_complex_shape_definedivar() {
+        eval(r#"
+            class C
+              def test = defined?(@a)
+            end
+            obj = C.new
+            (0..1000).each do |i|
+              obj.instance_variable_set(:"@v#{i}", i)
+            end
+            (0..1000).each do |i|
+              obj.remove_instance_variable(:"@v#{i}")
+            end
+            obj.test
+            TEST = C.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):
+          IncrCounter definedivar_fallback_too_complex
+          v10:StringExact|NilClass = DefinedIvar v6, :@a
+          CheckInterrupts
+          Return v10
+        ");
+    }
+
     #[test]
     fn test_specialize_monomorphic_setivar_already_in_shape() {
         eval("
@@ -3512,7 +3545,7 @@ mod hir_opt_tests {
     }
 
     #[test]
-    fn test_dont_specialize_monomorphic_setivar_with_shape_transition() {
+    fn test_specialize_monomorphic_setivar_with_shape_transition() {
         eval("
             def test = @foo = 5
             test
@@ -3529,13 +3562,57 @@ mod hir_opt_tests {
         bb2(v6:BasicObject):
           v10:Fixnum[5] = Const Value(5)
           PatchPoint SingleRactorMode
-          IncrCounter setivar_fallback_shape_transition
-          SetIvar v6, :@foo, v10
+          v19:HeapBasicObject = GuardType v6, HeapBasicObject
+          v20:HeapBasicObject = GuardShape v19, 0x1000
+          StoreField v20, :@foo@0x1001, v10
+          WriteBarrier v20, v10
+          v23:CUInt32[4194311] = Const CUInt32(4194311)
+          StoreField v20, :_shape_id@0x1002, v23
           CheckInterrupts
           Return v10
         ");
     }
 
+    #[test]
+    fn test_specialize_multiple_monomorphic_setivar_with_shape_transition() {
+        eval("
+            def test
+              @foo = 1
+              @bar = 2
+            end
+            test
+        ");
+        assert_snapshot!(hir_string("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[1] = Const Value(1)
+          PatchPoint SingleRactorMode
+          v25:HeapBasicObject = GuardType v6, HeapBasicObject
+          v26:HeapBasicObject = GuardShape v25, 0x1000
+          StoreField v26, :@foo@0x1001, v10
+          WriteBarrier v26, v10
+          v29:CUInt32[4194311] = Const CUInt32(4194311)
+          StoreField v26, :_shape_id@0x1002, v29
+          v16:Fixnum[2] = Const Value(2)
+          PatchPoint SingleRactorMode
+          v31:HeapBasicObject = GuardType v6, HeapBasicObject
+          v32:HeapBasicObject = GuardShape v31, 0x1003
+          StoreField v32, :@bar@0x1004, v16
+          WriteBarrier v32, v16
+          v35:CUInt32[4194312] = Const CUInt32(4194312)
+          StoreField v32, :_shape_id@0x1002, v35
+          CheckInterrupts
+          Return v16
+        ");
+    }
+
     #[test]
     fn test_dont_specialize_setivar_with_t_data() {
         eval("
@@ -3610,6 +3687,9 @@ mod hir_opt_tests {
             (0..1000).each do |i|
               obj.instance_variable_set(:"@v#{i}", i)
             end
+            (0..1000).each do |i|
+              obj.remove_instance_variable(:"@v#{i}")
+            end
             obj.test
             TEST = C.instance_method(:test)
         "#);
@@ -3625,7 +3705,7 @@ mod hir_opt_tests {
         bb2(v6:BasicObject):
           v10:Fixnum[5] = Const Value(5)
           PatchPoint SingleRactorMode
-          IncrCounter setivar_fallback_shape_transition
+          IncrCounter setivar_fallback_too_complex
           SetIvar v6, :@a, v10
           CheckInterrupts
           Return v10
@@ -5416,6 +5496,43 @@ mod hir_opt_tests {
         ");
     }
 
+    #[test]
+    fn test_dont_optimize_getivar_with_too_complex_shape() {
+        eval(r#"
+            class C
+              attr_accessor :foo
+            end
+            obj = C.new
+            (0..1000).each do |i|
+              obj.instance_variable_set(:"@v#{i}", i)
+            end
+            (0..1000).each do |i|
+              obj.remove_instance_variable(:"@v#{i}")
+            end
+            def test(o) = o.foo
+            test obj
+        "#);
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:12:
+        bb0():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          v2:BasicObject = GetLocal l0, SP@4
+          Jump bb2(v1, v2)
+        bb1(v5:BasicObject, v6:BasicObject):
+          EntryPoint JIT(0)
+          Jump bb2(v5, v6)
+        bb2(v8:BasicObject, v9:BasicObject):
+          PatchPoint MethodRedefined(C@0x1000, foo@0x1008, cme:0x1010)
+          PatchPoint NoSingletonClass(C@0x1000)
+          v21:HeapObject[class_exact:C] = GuardType v9, HeapObject[class_exact:C]
+          IncrCounter getivar_fallback_too_complex
+          v22:BasicObject = GetIvar v21, :@foo
+          CheckInterrupts
+          Return v22
+        ");
+    }
+
     #[test]
     fn test_optimize_send_with_block() {
         eval(r#"
@@ -5696,8 +5813,11 @@ mod hir_opt_tests {
           v16:Fixnum[5] = Const Value(5)
           PatchPoint MethodRedefined(C@0x1000, foo=@0x1008, cme:0x1010)
           v26:HeapObject[class_exact:C] = GuardType v9, HeapObject[class_exact:C]
-          IncrCounter setivar_fallback_shape_transition
-          SetIvar v26, :@foo, v16
+          v29:HeapObject[class_exact:C] = GuardShape v26, 0x1038
+          StoreField v29, :@foo@0x1039, v16
+          WriteBarrier v29, v16
+          v32:CUInt32[4194311] = Const CUInt32(4194311)
+          StoreField v29, :_shape_id@0x103a, v32
           CheckInterrupts
           Return v16
         ");
@@ -5728,8 +5848,11 @@ mod hir_opt_tests {
           v16:Fixnum[5] = Const Value(5)
           PatchPoint MethodRedefined(C@0x1000, foo=@0x1008, cme:0x1010)
           v26:HeapObject[class_exact:C] = GuardType v9, HeapObject[class_exact:C]
-          IncrCounter setivar_fallback_shape_transition
-          SetIvar v26, :@foo, v16
+          v29:HeapObject[class_exact:C] = GuardShape v26, 0x1038
+          StoreField v29, :@foo@0x1039, v16
+          WriteBarrier v29, v16
+          v32:CUInt32[4194311] = Const CUInt32(4194311)
+          StoreField v29, :_shape_id@0x103a, v32
           CheckInterrupts
           Return v16
         ");
@@ -9026,18 +9149,22 @@ mod hir_opt_tests {
          SetLocal l0, EP@3, v27
          v39:BasicObject = GetLocal l0, EP@3
          PatchPoint SingleRactorMode
-         IncrCounter setivar_fallback_shape_transition
-         SetIvar v14, :@formatted, v39
-         v45:Class[VMFrozenCore] = Const Value(VALUE(0x1000))
-         PatchPoint MethodRedefined(Class@0x1008, lambda@0x1010, cme:0x1018)
-         PatchPoint NoSingletonClass(Class@0x1008)
-         v60:BasicObject = CCallWithFrame v45, :RubyVM::FrozenCore.lambda@0x1040, block=0x1048
+         v56:HeapBasicObject = GuardType v14, HeapBasicObject
+         v57:HeapBasicObject = GuardShape v56, 0x1000
+         StoreField v57, :@formatted@0x1001, v39
+         WriteBarrier v57, v39
+         v60:CUInt32[4194311] = Const CUInt32(4194311)
+         StoreField v57, :_shape_id@0x1002, v60
+         v45:Class[VMFrozenCore] = Const Value(VALUE(0x1008))
+         PatchPoint MethodRedefined(Class@0x1010, lambda@0x1018, cme:0x1020)
+         PatchPoint NoSingletonClass(Class@0x1010)
+         v65:BasicObject = CCallWithFrame v45, :RubyVM::FrozenCore.lambda@0x1048, block=0x1050
          v48:BasicObject = GetLocal l0, EP@6
          v49:BasicObject = GetLocal l0, EP@5
          v50:BasicObject = GetLocal l0, EP@4
          v51:BasicObject = GetLocal l0, EP@3
          CheckInterrupts
-         Return v60
+         Return v65
        ");
     }
 }
diff --git a/zjit/src/hir_type/mod.rs b/zjit/src/hir_type/mod.rs
diff --git a/zjit/src/stats.rs b/zjit/src/stats.rs

Original file line number	Diff line number	Diff line change
`@@ -767,3 +767,9 @@ rb_yarv_str_eql_internal(VALUE str1, VALUE str2)`
`767`	`767`	`}`
`768`	`768`
`769`	`769`	`void rb_jit_str_concat_codepoint(VALUE str, VALUE codepoint);`
	`770`	`+`
	`771`	`+attr_index_t`
	`772`	`+rb_jit_shape_capacity(shape_id_t shape_id)`
	`773`	`+{`
	`774`	`+ return RSHAPE_CAPACITY(shape_id);`
	`775`	`+}`