ZJIT: Guard T_* in addition to shape in polymorphic getivar

This is a8f3c34 ("ZJIT: Add missing guard on ivar access on
T_{DATA,CLASS,MODULE}") but for the polymorphic implementation in HIR
build.

Author: XrXr

jit.c

@@ -40,6 +40,9 @@ enum jit_bindgen_constants {
     // Field offsets for the RString struct
     RUBY_OFFSET_RSTRING_LEN = offsetof(struct RString, len),
 
+    // Shape constant related to RBasic::flags. (See RBASIC_SET_SHAPE_ID())
+    RB_SHAPE_FLAG_SHIFT = SHAPE_FLAG_SHIFT,
+
     // Field offsets for rb_execution_context_t
     RUBY_OFFSET_EC_CFP = offsetof(rb_execution_context_t, cfp),
     RUBY_OFFSET_EC_INTERRUPT_FLAG = offsetof(rb_execution_context_t, interrupt_flag),

yjit/src/cruby_bindings.inc.rs

@@ -617,6 +617,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         &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 { val: Const::CUInt64(val) } => Opnd::UImm(val),
         &Insn::Const { val: Const::CShape(val) } => {
             assert_eq!(SHAPE_ID_NUM_BITS, 32);
             gen_const_uint32(val.0)

zjit/src/codegen.rs

@@ -3714,6 +3714,7 @@ fn test_getivar_t_class_then_string() {
     assert_snapshot!(assert_compiles("[STR.test, STR.test]"), @"[1000, 1000]");
 }
 
+
 #[test]
 fn test_attr_accessor_setivar() {
     eval("

zjit/src/codegen_tests.rs

@@ -6342,8 +6342,16 @@ impl Function {
             Insn::IntAnd { left, right }
             | Insn::IntOr { left, right } => {
                 // TODO: Expand this to other matching C integer sizes when we need them.
-                self.assert_subtype(insn_id, left, types::CInt64)?;
-                self.assert_subtype(insn_id, right, types::CInt64)
+                let left_type = self.type_of(left);
+                if left_type.is_subtype(types::CInt64) {
+                    self.assert_subtype(insn_id, right, types::CInt64)
+                } else if left_type.is_subtype(types::CUInt64) {
+                    self.assert_subtype(insn_id, right, types::CUInt64)
+                } else {
+                    let all_ints = types::CInt64.union(types::CUInt64);
+                    self.assert_subtype(insn_id, left, all_ints)?;
+                    self.assert_subtype(insn_id, right, all_ints)
+                }
             }
             Insn::BoxBool { val }
             | Insn::IfTrue { val, .. }
@@ -8185,31 +8193,37 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     }
                     if let Some(summary) = fun.polymorphic_summary(&profiles, self_param, exit_state.insn_idx) {
                         self_param = fun.push_insn(block, Insn::GuardType { val: self_param, guard_type: types::HeapBasicObject, state: exit_id });
-                        let shape = fun.load_shape(block, self_param);
+                        let rbasic_flags = fun.load_rbasic_flags(block, self_param);
                         let join_block = insn_idx_to_block.get(&insn_idx).copied().unwrap_or_else(|| fun.new_block(insn_idx));
                         let join_param = fun.push_insn(join_block, Insn::Param);
                         // Dedup by expected shape so objects with different classes but the same shape can share code
                         // TODO(max): De-duplicate further by checking ivar offsets to allow
                         // different shapes with the same ivar layout to share code
-                        let mut seen_shapes = Vec::with_capacity(summary.buckets().len());
+                        let mut seen_shape_and_flags = Vec::with_capacity(summary.buckets().len());
                         for &profiled_type in summary.buckets() {
                             // End of the buckets
                             if profiled_type.is_empty() { break; }
                             // Instance variable lookups on immediate values are always nil; don't bother
                             if profiled_type.flags().is_immediate() { continue; }
                             let expected_shape = profiled_type.shape();
+                            let (expected_rbasic_flags, rbasic_flags_mask) = profiled_type.rbasic_flags_and_mask();
                             assert!(expected_shape.is_valid());
                             // Too-complex shapes use hash tables for ivars;
                             // rb_shape_get_iv_index doesn't work for them.
                             // Let the fallthrough GetIvar handle these.
                             if expected_shape.is_too_complex() { continue; }
-                            if seen_shapes.contains(&expected_shape) { continue; }
-                            seen_shapes.push(expected_shape);
-                            let expected_shape_const = fun.push_insn(block, Insn::Const { val: Const::CShape(expected_shape) });
-                            let has_shape = fun.push_insn(block, Insn::IsBitEqual { left: shape, right: expected_shape_const });
+                            if seen_shape_and_flags.contains(&expected_rbasic_flags) { continue; }
+                            seen_shape_and_flags.push(expected_rbasic_flags);
+                            let rbasic_flags_mask = fun.push_insn(block, Insn::Const { val: Const::CUInt64(rbasic_flags_mask) });
+                            // The expected shape can change over run, so we put it
+                            // as a pointer to keep it stable in snapshot tests.
+                            let expected_rbasic_flags = fun.push_insn(block, Insn::Const { val: Const::CPtr(ptr::without_provenance(expected_rbasic_flags.to_usize())) });
+                            let expected_rbasic_flags = fun.push_insn(block, Insn::RefineType { val: expected_rbasic_flags, new_type: types::CUInt64 });
+                            let masked = fun.push_insn(block, Insn::IntAnd { left: rbasic_flags, right: rbasic_flags_mask});
+                            let has_shape_and_type = fun.push_insn(block, Insn::IsBitEqual { left: masked, right: expected_rbasic_flags });
                             let iftrue_block = fun.new_block(insn_idx);
                             let target = BranchEdge { target: iftrue_block, args: vec![] };
-                            fun.push_insn(block, Insn::IfTrue { val: has_shape, target });
+                            fun.push_insn(block, Insn::IfTrue { val: has_shape_and_type, target });
                             let result = fun.load_ivar(iftrue_block, self_param, profiled_type, id, exit_id);
                             fun.push_insn(iftrue_block, Insn::Jump(BranchEdge { target: join_block, args: vec![result] }));
                         }

zjit/src/cruby_bindings.inc.rs

@@ -7712,29 +7712,35 @@ mod hir_opt_tests {
         bb3(v6:BasicObject):
           PatchPoint SingleRactorMode
           v11:HeapBasicObject = GuardType v6, HeapBasicObject
-          v12:CShape = LoadField v11, :_shape_id@0x1000
-          v14:CShape[0x1001] = Const CShape(0x1001)
-          v15:CBool = IsBitEqual v12, v14
-          IfTrue v15, bb5()
-          v20:CShape[0x1002] = Const CShape(0x1002)
-          v21:CBool = IsBitEqual v12, v20
-          IfTrue v21, bb6()
-          v25:BasicObject = GetIvar v11, :@foo
-          Jump bb4(v25)
+          v12:CUInt64 = LoadField v11, :_rbasic_flags@0x1000
+          v14:CUInt64[18446744069414584351] = Const CUInt64(18446744069414584351)
+          v15:CPtr[CPtr(0x1008)] = Const CPtr(0x1010)
+          v16 = RefineType v15, CUInt64
+          v17:CInt64 = IntAnd v12, v14
+          v18:CBool = IsBitEqual v17, v16
+          IfTrue v18, bb5()
+          v23:CUInt64[18446744069414584351] = Const CUInt64(18446744069414584351)
+          v24:CPtr[CPtr(0x1008)] = Const CPtr(0x1010)
+          v25 = RefineType v24, CUInt64
+          v26:CInt64 = IntAnd v12, v23
+          v27:CBool = IsBitEqual v26, v25
+          IfTrue v27, bb6()
+          v31:BasicObject = GetIvar v11, :@foo
+          Jump bb4(v31)
         bb5():
-          v17:CPtr = LoadField v11, :_as_heap@0x1003
-          v18:BasicObject = LoadField v17, :@foo@0x1004
-          Jump bb4(v18)
+          v20:CPtr = LoadField v11, :_as_heap@0x1018
+          v21:BasicObject = LoadField v20, :@foo@0x1019
+          Jump bb4(v21)
         bb6():
-          v23:BasicObject = LoadField v11, :@foo@0x1003
-          Jump bb4(v23)
+          v29:BasicObject = LoadField v11, :@foo@0x1018
+          Jump bb4(v29)
         bb4(v13:BasicObject):
-          v28:Fixnum[1] = Const Value(1)
-          PatchPoint MethodRedefined(Integer@0x1008, +@0x1010, cme:0x1018)
-          v39:Fixnum = GuardType v13, Fixnum
-          v40:Fixnum = FixnumAdd v39, v28
+          v34:Fixnum[1] = Const Value(1)
+          PatchPoint MethodRedefined(Integer@0x1020, +@0x1028, cme:0x1030)
+          v45:Fixnum = GuardType v13, Fixnum
+          v46:Fixnum = FixnumAdd v45, v34
           CheckInterrupts
-          Return v40
+          Return v46
         ");
     }
 
@@ -7782,31 +7788,37 @@ mod hir_opt_tests {
         bb3(v6:BasicObject):
           PatchPoint SingleRactorMode
           v11:HeapBasicObject = GuardType v6, HeapBasicObject
-          v12:CShape = LoadField v11, :_shape_id@0x1000
-          v14:CShape[0x1001] = Const CShape(0x1001)
-          v15:CBool = IsBitEqual v12, v14
-          IfTrue v15, bb5()
-          v19:CShape[0x1002] = Const CShape(0x1002)
-          v20:CBool = IsBitEqual v12, v19
-          IfTrue v20, bb6()
-          v38:CShape = LoadField v11, :_shape_id@0x1000
-          v39:CShape[0x1001] = GuardBitEquals v38, CShape(0x1001)
-          v40:BasicObject = LoadField v11, :@foo@0x1003
-          Jump bb4(v40)
+          v12:CUInt64 = LoadField v11, :_rbasic_flags@0x1000
+          v14:CUInt64[18446744069414584351] = Const CUInt64(18446744069414584351)
+          v15:CPtr[CPtr(0x1008)] = Const CPtr(0x1010)
+          v16 = RefineType v15, CUInt64
+          v17:CInt64 = IntAnd v12, v14
+          v18:CBool = IsBitEqual v17, v16
+          IfTrue v18, bb5()
+          v22:CUInt64[18446744069414584351] = Const CUInt64(18446744069414584351)
+          v23:CPtr[CPtr(0x1008)] = Const CPtr(0x1010)
+          v24 = RefineType v23, CUInt64
+          v25:CInt64 = IntAnd v12, v22
+          v26:CBool = IsBitEqual v25, v24
+          IfTrue v26, bb6()
+          v44:CShape = LoadField v11, :_shape_id@0x1018
+          v45:CShape[0x1019] = GuardBitEquals v44, CShape(0x1019)
+          v46:BasicObject = LoadField v11, :@foo@0x101a
+          Jump bb4(v46)
         bb5():
-          v17:BasicObject = LoadField v11, :@foo@0x1003
-          Jump bb4(v17)
+          v20:BasicObject = LoadField v11, :@foo@0x101a
+          Jump bb4(v20)
         bb6():
-          v22:CPtr = LoadField v11, :_as_heap@0x1003
-          v23:BasicObject = LoadField v22, :@foo@0x1004
-          Jump bb4(v23)
+          v28:CPtr = LoadField v11, :_as_heap@0x101a
+          v29:BasicObject = LoadField v28, :@foo@0x101b
+          Jump bb4(v29)
         bb4(v13:BasicObject):
-          v28:Fixnum[1] = Const Value(1)
-          PatchPoint MethodRedefined(Integer@0x1008, +@0x1010, cme:0x1018)
-          v43:Fixnum = GuardType v13, Fixnum
-          v44:Fixnum = FixnumAdd v43, v28
+          v34:Fixnum[1] = Const Value(1)
+          PatchPoint MethodRedefined(Integer@0x1020, +@0x1028, cme:0x1030)
+          v49:Fixnum = GuardType v13, Fixnum
+          v50:Fixnum = FixnumAdd v49, v34
           CheckInterrupts
-          Return v44
+          Return v50
         ");
     }
 
@@ -7847,28 +7859,99 @@ mod hir_opt_tests {
         bb3(v6:BasicObject):
           PatchPoint SingleRactorMode
           v11:HeapBasicObject = GuardType v6, HeapBasicObject
-          v12:CShape = LoadField v11, :_shape_id@0x1000
-          v14:CShape[0x1001] = Const CShape(0x1001)
-          v15:CBool = IsBitEqual v12, v14
-          IfTrue v15, bb5()
-          v19:CShape[0x1002] = Const CShape(0x1002)
-          v20:CBool = IsBitEqual v12, v19
-          IfTrue v20, bb6()
-          v24:BasicObject = GetIvar v11, :@foo
-          Jump bb4(v24)
+          v12:CUInt64 = LoadField v11, :_rbasic_flags@0x1000
+          v14:CUInt64[18446744069414584351] = Const CUInt64(18446744069414584351)
+          v15:CPtr[CPtr(0x1008)] = Const CPtr(0x1010)
+          v16 = RefineType v15, CUInt64
+          v17:CInt64 = IntAnd v12, v14
+          v18:CBool = IsBitEqual v17, v16
+          IfTrue v18, bb5()
+          v22:CUInt64[18446744069414584351] = Const CUInt64(18446744069414584351)
+          v23:CPtr[CPtr(0x1008)] = Const CPtr(0x1010)
+          v24 = RefineType v23, CUInt64
+          v25:CInt64 = IntAnd v12, v22
+          v26:CBool = IsBitEqual v25, v24
+          IfTrue v26, bb6()
+          v30:BasicObject = GetIvar v11, :@foo
+          Jump bb4(v30)
         bb5():
-          v17:BasicObject = LoadField v11, :@foo@0x1003
-          Jump bb4(v17)
+          v20:BasicObject = LoadField v11, :@foo@0x1018
+          Jump bb4(v20)
         bb6():
-          v22:BasicObject = LoadField v11, :@foo@0x1003
-          Jump bb4(v22)
+          v28:BasicObject = LoadField v11, :@foo@0x1018
+          Jump bb4(v28)
         bb4(v13:BasicObject):
-          v27:Fixnum[1] = Const Value(1)
-          PatchPoint MethodRedefined(Integer@0x1008, +@0x1010, cme:0x1018)
-          v38:Fixnum = GuardType v13, Fixnum
-          v39:Fixnum = FixnumAdd v38, v27
+          v33:Fixnum[1] = Const Value(1)
+          PatchPoint MethodRedefined(Integer@0x1020, +@0x1028, cme:0x1030)
+          v44:Fixnum = GuardType v13, Fixnum
+          v45:Fixnum = FixnumAdd v44, v33
           CheckInterrupts
-          Return v39
+          Return v45
+        ");
+    }
+
+    #[test]
+    fn test_getivar_polymorphic_t_class_and_typed_data() {
+        set_call_threshold(3);
+        eval(r#"
+          module Reader
+            def test = @a
+          end
+
+          class A
+            extend Reader
+            @a = 0
+          end
+
+          ARGF.instance_eval do
+            extend Reader
+            @a = :a
+          end
+
+          A.test
+          ARGF.test
+        "#);
+        assert_snapshot!(assert_compiles("[A.test, ARGF.test]"), @"[0, :a]");
+        assert_snapshot!(hir_string_proc("Reader.instance_method(: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 SingleRactorMode
+          v11:HeapBasicObject = GuardType v6, HeapBasicObject
+          v12:CUInt64 = LoadField v11, :_rbasic_flags@0x1000
+          v14:CUInt64[0xffffffff0000005f] = Const CUInt64(0xffffffff0000005f)
+          v15:CPtr[CPtr(0x1001)] = Const CPtr(0x1001)
+          v16 = RefineType v15, CUInt64
+          v17:CInt64 = IntAnd v12, v14
+          v18:CBool = IsBitEqual v17, v16
+          IfTrue v18, bb5()
+          v23:CUInt64[0xffffffff0000001f] = Const CUInt64(0xffffffff0000001f)
+          v24:CPtr[CPtr(0x1002)] = Const CPtr(0x1002)
+          v25 = RefineType v24, CUInt64
+          v26:CInt64 = IntAnd v12, v23
+          v27:CBool = IsBitEqual v26, v25
+          IfTrue v27, bb6()
+          v33:BasicObject = GetIvar v11, :@a
+          Jump bb4(v33)
+        bb5():
+          v20:RubyValue = LoadField v11, :_fields_obj@0x1003
+          v21:BasicObject = LoadField v20, :@a@0x1003
+          Jump bb4(v21)
+        bb6():
+          PatchPoint RootBoxOnly
+          v30:RubyValue = LoadField v11, :_fields_obj@0x1004
+          v31:BasicObject = LoadField v30, :@a@0x1003
+          Jump bb4(v31)
+        bb4(v13:BasicObject):
+          CheckInterrupts
+          Return v13
         ");
     }
 

zjit/src/hir.rs

@@ -333,6 +333,25 @@ impl ProfiledType {
         self.flags
     }
 
+    /// For ivar access
+    pub fn rbasic_flags_and_mask(&self) -> (u64, u64) {
+        let shape_flag_shift = u64::from(RB_SHAPE_FLAG_SHIFT);
+        let (shape, shape_mask) = (u64::from(self.shape().0) << shape_flag_shift, !0 << shape_flag_shift);
+        let (builtin_type, type_mask) = if self.flags().is_t_object() {
+            (RUBY_T_OBJECT, RUBY_T_MASK)
+        } else if self.class().is_subclass_of(unsafe { rb_cClass }) == ClassRelationship::Subclass {
+            // Check class first since `Class < Module`
+            (RUBY_T_CLASS, RUBY_T_MASK)
+        } else if self.class().is_subclass_of(unsafe { rb_cModule }) == ClassRelationship::Subclass {
+            (RUBY_T_MODULE, RUBY_T_MASK)
+        } else if self.flags().is_typed_data() {
+            (RUBY_T_DATA | RUBY_TYPED_FL_IS_TYPED_DATA, RUBY_T_MASK | RUBY_TYPED_FL_IS_TYPED_DATA)
+        } else {
+            (0, 0)
+        };
+        (shape | u64::from(builtin_type), shape_mask | u64::from(type_mask))
+    }
+
     pub fn is_fixnum(&self) -> bool {
         self.class == unsafe { rb_cInteger } && self.flags.is_immediate()
     }