ZJIT: Specialize setinstancevariable when ivar is already in shape (#15290)

Don't support shape transitions for now.

insns.def

@@ -224,6 +224,7 @@ setinstancevariable
 (VALUE val)
 ()
 // attr bool leaf = false; /* has rb_check_frozen() */
+// attr bool zjit_profile = true;
 {
     vm_setinstancevariable(GET_ISEQ(), GET_SELF(), id, val, ic);
 }

zjit/bindgen/src/main.rs

@@ -101,6 +101,7 @@ fn main() {
         .allowlist_function("rb_shape_get_iv_index")
         .allowlist_function("rb_shape_transition_add_ivar_no_warnings")
         .allowlist_var("rb_invalid_shape_id")
+        .allowlist_type("shape_id_fl_type")
         .allowlist_var("VM_KW_SPECIFIED_BITS_MAX")
         .allowlist_var("SHAPE_ID_NUM_BITS")
         .allowlist_function("rb_obj_is_kind_of")

zjit/src/cruby.rs

@@ -263,6 +263,10 @@ impl ShapeId {
     pub fn is_too_complex(self) -> bool {
         unsafe { rb_jit_shape_too_complex_p(self.0) }
     }
+
+    pub fn is_frozen(self) -> bool {
+        (self.0 & SHAPE_ID_FL_FROZEN) != 0
+    }
 }
 
 // Given an ISEQ pointer, convert PC to insn_idx

zjit/src/cruby_bindings.inc.rs

@@ -1462,6 +1462,13 @@ pub const VM_ENV_FLAG_ISOLATED: vm_frame_env_flags = 16;
 pub type vm_frame_env_flags = u32;
 pub type attr_index_t = u16;
 pub type shape_id_t = u32;
+pub const SHAPE_ID_HEAP_INDEX_MASK: shape_id_fl_type = 29360128;
+pub const SHAPE_ID_FL_FROZEN: shape_id_fl_type = 33554432;
+pub const SHAPE_ID_FL_HAS_OBJECT_ID: shape_id_fl_type = 67108864;
+pub const SHAPE_ID_FL_TOO_COMPLEX: shape_id_fl_type = 134217728;
+pub const SHAPE_ID_FL_NON_CANONICAL_MASK: shape_id_fl_type = 100663296;
+pub const SHAPE_ID_FLAGS_MASK: shape_id_fl_type = 264241152;
+pub type shape_id_fl_type = u32;
 #[repr(C)]
 pub struct rb_cvar_class_tbl_entry {
     pub index: u32,
@@ -1744,35 +1751,36 @@ pub const YARVINSN_trace_setlocal_WC_1: ruby_vminsn_type = 215;
 pub const YARVINSN_trace_putobject_INT2FIX_0_: ruby_vminsn_type = 216;
 pub const YARVINSN_trace_putobject_INT2FIX_1_: ruby_vminsn_type = 217;
 pub const YARVINSN_zjit_getinstancevariable: ruby_vminsn_type = 218;
-pub const YARVINSN_zjit_definedivar: ruby_vminsn_type = 219;
-pub const YARVINSN_zjit_send: ruby_vminsn_type = 220;
-pub const YARVINSN_zjit_opt_send_without_block: ruby_vminsn_type = 221;
-pub const YARVINSN_zjit_objtostring: ruby_vminsn_type = 222;
-pub const YARVINSN_zjit_opt_nil_p: ruby_vminsn_type = 223;
-pub const YARVINSN_zjit_invokeblock: ruby_vminsn_type = 224;
-pub const YARVINSN_zjit_opt_plus: ruby_vminsn_type = 225;
-pub const YARVINSN_zjit_opt_minus: ruby_vminsn_type = 226;
-pub const YARVINSN_zjit_opt_mult: ruby_vminsn_type = 227;
-pub const YARVINSN_zjit_opt_div: ruby_vminsn_type = 228;
-pub const YARVINSN_zjit_opt_mod: ruby_vminsn_type = 229;
-pub const YARVINSN_zjit_opt_eq: ruby_vminsn_type = 230;
-pub const YARVINSN_zjit_opt_neq: ruby_vminsn_type = 231;
-pub const YARVINSN_zjit_opt_lt: ruby_vminsn_type = 232;
-pub const YARVINSN_zjit_opt_le: ruby_vminsn_type = 233;
-pub const YARVINSN_zjit_opt_gt: ruby_vminsn_type = 234;
-pub const YARVINSN_zjit_opt_ge: ruby_vminsn_type = 235;
-pub const YARVINSN_zjit_opt_ltlt: ruby_vminsn_type = 236;
-pub const YARVINSN_zjit_opt_and: ruby_vminsn_type = 237;
-pub const YARVINSN_zjit_opt_or: ruby_vminsn_type = 238;
-pub const YARVINSN_zjit_opt_aref: ruby_vminsn_type = 239;
-pub const YARVINSN_zjit_opt_aset: ruby_vminsn_type = 240;
-pub const YARVINSN_zjit_opt_length: ruby_vminsn_type = 241;
-pub const YARVINSN_zjit_opt_size: ruby_vminsn_type = 242;
-pub const YARVINSN_zjit_opt_empty_p: ruby_vminsn_type = 243;
-pub const YARVINSN_zjit_opt_succ: ruby_vminsn_type = 244;
-pub const YARVINSN_zjit_opt_not: ruby_vminsn_type = 245;
-pub const YARVINSN_zjit_opt_regexpmatch2: ruby_vminsn_type = 246;
-pub const VM_INSTRUCTION_SIZE: ruby_vminsn_type = 247;
+pub const YARVINSN_zjit_setinstancevariable: ruby_vminsn_type = 219;
+pub const YARVINSN_zjit_definedivar: ruby_vminsn_type = 220;
+pub const YARVINSN_zjit_send: ruby_vminsn_type = 221;
+pub const YARVINSN_zjit_opt_send_without_block: ruby_vminsn_type = 222;
+pub const YARVINSN_zjit_objtostring: ruby_vminsn_type = 223;
+pub const YARVINSN_zjit_opt_nil_p: ruby_vminsn_type = 224;
+pub const YARVINSN_zjit_invokeblock: ruby_vminsn_type = 225;
+pub const YARVINSN_zjit_opt_plus: ruby_vminsn_type = 226;
+pub const YARVINSN_zjit_opt_minus: ruby_vminsn_type = 227;
+pub const YARVINSN_zjit_opt_mult: ruby_vminsn_type = 228;
+pub const YARVINSN_zjit_opt_div: ruby_vminsn_type = 229;
+pub const YARVINSN_zjit_opt_mod: ruby_vminsn_type = 230;
+pub const YARVINSN_zjit_opt_eq: ruby_vminsn_type = 231;
+pub const YARVINSN_zjit_opt_neq: ruby_vminsn_type = 232;
+pub const YARVINSN_zjit_opt_lt: ruby_vminsn_type = 233;
+pub const YARVINSN_zjit_opt_le: ruby_vminsn_type = 234;
+pub const YARVINSN_zjit_opt_gt: ruby_vminsn_type = 235;
+pub const YARVINSN_zjit_opt_ge: ruby_vminsn_type = 236;
+pub const YARVINSN_zjit_opt_ltlt: ruby_vminsn_type = 237;
+pub const YARVINSN_zjit_opt_and: ruby_vminsn_type = 238;
+pub const YARVINSN_zjit_opt_or: ruby_vminsn_type = 239;
+pub const YARVINSN_zjit_opt_aref: ruby_vminsn_type = 240;
+pub const YARVINSN_zjit_opt_aset: ruby_vminsn_type = 241;
+pub const YARVINSN_zjit_opt_length: ruby_vminsn_type = 242;
+pub const YARVINSN_zjit_opt_size: ruby_vminsn_type = 243;
+pub const YARVINSN_zjit_opt_empty_p: ruby_vminsn_type = 244;
+pub const YARVINSN_zjit_opt_succ: ruby_vminsn_type = 245;
+pub const YARVINSN_zjit_opt_not: ruby_vminsn_type = 246;
+pub const YARVINSN_zjit_opt_regexpmatch2: ruby_vminsn_type = 247;
+pub const VM_INSTRUCTION_SIZE: ruby_vminsn_type = 248;
 pub type ruby_vminsn_type = u32;
 pub type rb_iseq_callback = ::std::option::Option<
     unsafe extern "C" fn(arg1: *const rb_iseq_t, arg2: *mut ::std::os::raw::c_void),

zjit/src/hir.rs

@@ -2873,6 +2873,50 @@ impl Function {
                         };
                         self.make_equal_to(insn_id, replacement);
                     }
+                    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
+                            self.push_insn_id(block, insn_id); continue;
+                        };
+                        if recv_type.flags().is_immediate() {
+                            // Instance variable lookups on immediate values are always nil
+                            self.push_insn_id(block, insn_id); continue;
+                        }
+                        assert!(recv_type.shape().is_valid());
+                        if !recv_type.flags().is_t_object() {
+                            // Check if the receiver is a T_OBJECT
+                            self.push_insn_id(block, insn_id); continue;
+                        }
+                        if recv_type.shape().is_too_complex() {
+                            // too-complex shapes can't use index access
+                            self.push_insn_id(block, insn_id); continue;
+                        }
+                        if recv_type.shape().is_frozen() {
+                            // Can't set ivars on frozen objects
+                            self.push_insn_id(block, insn_id); continue;
+                        }
+                        let mut ivar_index: u16 = 0;
+                        if !unsafe { rb_shape_get_iv_index(recv_type.shape().0, id, &mut ivar_index) } {
+                            // TODO(max): Shape transition
+                            self.push_insn_id(block, insn_id); continue;
+                        }
+                        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 });
+                        // 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
+                            let offset = ROBJECT_OFFSET_AS_ARY as i32 + (SIZEOF_VALUE * ivar_index.to_usize()) as i32;
+                            (self_val, offset)
+                        } else {
+                            let as_heap = self.push_insn(block, Insn::LoadField { recv: self_val, id: ID!(_as_heap), offset: ROBJECT_OFFSET_AS_HEAP_FIELDS as i32, return_type: types::CPtr });
+                            let offset = SIZEOF_VALUE_I32 * ivar_index as i32;
+                            (as_heap, offset)
+                        };
+                        let replacement = self.push_insn(block, Insn::StoreField { recv: ivar_storage, id, offset, val });
+                        self.push_insn(block, Insn::WriteBarrier { recv: self_val, val });
+                        self.make_equal_to(insn_id, replacement);
+                    }
                     _ => { self.push_insn_id(block, insn_id); }
                 }
             }
@@ -4906,6 +4950,8 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
             // profiled cfp->self.
             if opcode == YARVINSN_getinstancevariable || opcode == YARVINSN_trace_getinstancevariable {
                 profiles.profile_self(&exit_state, self_param);
+            } else if opcode == YARVINSN_setinstancevariable || opcode == YARVINSN_trace_setinstancevariable {
+                profiles.profile_self(&exit_state, self_param);
             } else if opcode == YARVINSN_definedivar || opcode == YARVINSN_trace_definedivar {
                 profiles.profile_self(&exit_state, self_param);
             } else if opcode == YARVINSN_invokeblock || opcode == YARVINSN_trace_invokeblock {

zjit/src/hir/opt_tests.rs

@@ -3479,6 +3479,151 @@ mod hir_opt_tests {
         ");
     }
 
+    #[test]
+    fn test_specialize_monomorphic_setivar_already_in_shape() {
+        eval("
+            @foo = 4
+            def test = @foo = 5
+            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[5] = Const Value(5)
+          PatchPoint SingleRactorMode
+          v19:HeapBasicObject = GuardType v6, HeapBasicObject
+          v20:HeapBasicObject = GuardShape v19, 0x1000
+          StoreField v20, :@foo@0x1001, v10
+          WriteBarrier v20, v10
+          CheckInterrupts
+          Return v10
+        ");
+    }
+
+    #[test]
+    fn test_dont_specialize_monomorphic_setivar_with_shape_transition() {
+        eval("
+            def test = @foo = 5
+            test
+        ");
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:2:
+        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
+          SetIvar v6, :@foo, v10
+          CheckInterrupts
+          Return v10
+        ");
+    }
+
+    #[test]
+    fn test_dont_specialize_setivar_with_t_data() {
+        eval("
+            class C < Range
+              def test = @a = 5
+            end
+            obj = C.new 0, 1
+            obj.instance_variable_set(:@a, 1)
+            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):
+          v10:Fixnum[5] = Const Value(5)
+          PatchPoint SingleRactorMode
+          SetIvar v6, :@a, v10
+          CheckInterrupts
+          Return v10
+        ");
+    }
+
+    #[test]
+    fn test_dont_specialize_polymorphic_setivar() {
+        set_call_threshold(3);
+        eval("
+            class C
+              def test = @a = 5
+            end
+            obj = C.new
+            obj.instance_variable_set(:@a, 1)
+            obj.test
+            obj = C.new
+            obj.instance_variable_set(:@b, 1)
+            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):
+          v10:Fixnum[5] = Const Value(5)
+          PatchPoint SingleRactorMode
+          SetIvar v6, :@a, v10
+          CheckInterrupts
+          Return v10
+        ");
+    }
+
+    #[test]
+    fn test_dont_specialize_complex_shape_setivar() {
+        eval(r#"
+            class C
+              def test = @a = 5
+            end
+            obj = C.new
+            (0..1000).each do |i|
+              obj.instance_variable_set(:"@v#{i}", 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):
+          v10:Fixnum[5] = Const Value(5)
+          PatchPoint SingleRactorMode
+          SetIvar v6, :@a, v10
+          CheckInterrupts
+          Return v10
+        ");
+    }
+
     #[test]
     fn test_elide_freeze_with_frozen_hash() {
         eval("

zjit/src/profile.rs

@@ -82,6 +82,7 @@ fn profile_insn(bare_opcode: ruby_vminsn_type, ec: EcPtr) {
         YARVINSN_opt_aset  => profile_operands(profiler, profile, 3),
         YARVINSN_opt_not   => profile_operands(profiler, profile, 1),
         YARVINSN_getinstancevariable => profile_self(profiler, profile),
+        YARVINSN_setinstancevariable => profile_self(profiler, profile),
         YARVINSN_definedivar   => profile_self(profiler, profile),
         YARVINSN_opt_regexpmatch2    => profile_operands(profiler, profile, 2),
         YARVINSN_objtostring   => profile_operands(profiler, profile, 1),