ZJIT: Add missing guard on ivar access on T_{DATA,CLASS,MODULE}

T_DATA, T_MODULE, and T_CLASS objects can share the exact same shape.
The shape on these objects give an index off of the fields array to
get at the ivar. When two objects share the same shape, but differ
in the T_* builtin type, however, the way to get to the fields array
differ.

Previously, we did not guard the builtin type, so the guard allowed
using say, loading `t_string[RCLASS_OFFSET_PRIME_FIELDS_OBJ]`. A classic
type confusion situation that crashed.

Guard the builtin type, in addition to the shape. Note that this is not
necessary for T_OBJECTs since they never have the same shape as other
builtin types.

Author: XrXr

zjit/src/codegen.rs

@@ -2549,7 +2549,7 @@ fn gen_guard_type(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, guard
 
         asm.cmp(klass, Opnd::Value(expected_class));
         asm.jne(jit, side_exit);
-    } else if guard_type.is_subtype(types::String) {
+    } else if guard_type.is_subtype(types::TypedTData) {
         let side = side_exit(jit, state, GuardType(guard_type));
 
         // Check special constant
@@ -2560,13 +2560,15 @@ fn gen_guard_type(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, guard
         asm.cmp(val, Qfalse.into());
         asm.je(jit, side.clone());
 
+        // Check the builtin type and RUBY_TYPED_FL_IS_TYPED_DATA with mask and compare
         let val = asm.load_mem(val);
-
         let flags = asm.load(Opnd::mem(VALUE_BITS, val, RUBY_OFFSET_RBASIC_FLAGS));
-        let tag   = asm.and(flags, Opnd::UImm(RUBY_T_MASK as u64));
-        asm.cmp(tag, Opnd::UImm(RUBY_T_STRING as u64));
+        let mask     = RUBY_T_MASK.to_usize() | RUBY_TYPED_FL_IS_TYPED_DATA.to_usize();
+        let expected = RUBY_T_DATA.to_usize() | RUBY_TYPED_FL_IS_TYPED_DATA.to_usize();
+        let masked = asm.and(flags, mask.into());
+        asm.cmp(masked, expected.into());
         asm.jne(jit, side);
-    } else if guard_type.is_subtype(types::Array) {
+    } else if let Some(builtin_type) = guard_type.builtin_type_equivalent() {
         let side = side_exit(jit, state, GuardType(guard_type));
 
         // Check special constant
@@ -2577,11 +2579,11 @@ fn gen_guard_type(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, guard
         asm.cmp(val, Qfalse.into());
         asm.je(jit, side.clone());
 
+        // Mask and check the builtin type
         let val = asm.load_mem(val);
-
         let flags = asm.load(Opnd::mem(VALUE_BITS, val, RUBY_OFFSET_RBASIC_FLAGS));
         let tag   = asm.and(flags, Opnd::UImm(RUBY_T_MASK as u64));
-        asm.cmp(tag, Opnd::UImm(RUBY_T_ARRAY as u64));
+        asm.cmp(tag, Opnd::UImm(builtin_type as u64));
         asm.jne(jit, side);
     } else if guard_type.bit_equal(types::HeapBasicObject) {
         let side_exit = side_exit(jit, state, GuardType(guard_type));

zjit/src/codegen_tests.rs

@@ -3623,6 +3623,97 @@ fn test_attr_accessor_getivar() {
     assert_snapshot!(assert_compiles("c = C.new; [test(c), test(c)]"), @"[4, 4]");
 }
 
+#[test]
+fn test_getivar_t_data_then_string() {
+    // This is a regression test for a type confusion miscomp where
+    // we end up reading the fields object using an offset off of a
+    // string, assuming that it has a the same layout as a T_DATA object.
+    // At the time of writing the fields object of strings are stored
+    // in a global table, out-of-line of each string.
+    // The string and the thread end up sharing one shape ID.
+    set_call_threshold(2);
+    eval(r#"
+      module GetThousand
+        def test = @var1000
+      end
+      class Thread
+        include GetThousand
+      end
+      class String
+        include GetThousand
+      end
+      OBJ = Thread.new { }
+      OBJ.join
+      STR = +''
+      (0..1000).each do |i|
+        ivar_name = :"@var#{i}"
+        OBJ.instance_variable_set(ivar_name, i)
+        STR.instance_variable_set(ivar_name, i)
+      end
+      OBJ.test; OBJ.test # profile and compile for Thread (T_DATA)
+    "#);
+    assert_snapshot!(assert_compiles("[STR.test, STR.test]"), @"[1000, 1000]");
+}
+
+#[test]
+fn test_getivar_t_object_then_string() {
+    // This test construct an object and a string that have the same set of ivars.
+    // They wouldn't share the same shape ID, though, and we rely on this fact in
+    // our guards.
+    set_call_threshold(2);
+    eval(r#"
+      module GetThousand
+        def test = @var1000
+      end
+      class MyObject
+        include GetThousand
+      end
+      class String
+        include GetThousand
+      end
+      OBJ = MyObject.new
+      STR = +''
+      (0..1000).each do |i|
+        ivar_name = :"@var#{i}"
+        OBJ.instance_variable_set(ivar_name, i)
+        STR.instance_variable_set(ivar_name, i)
+      end
+      OBJ.test; OBJ.test # profile and compile for MyObject
+    "#);
+    assert_snapshot!(assert_compiles("[STR.test, STR.test]"), @"[1000, 1000]");
+}
+
+#[test]
+fn test_getivar_t_class_then_string() {
+    // This is a regression test for a type confusion miscomp where
+    // we end up reading the fields object using an offset off of a
+    // string, assuming that it has a the same layout as a T_CLASS object.
+    // At the time of writing the fields object of strings are stored
+    // in a global table, out-of-line of each string.
+    // The string and the class end up sharing one shape ID.
+    set_call_threshold(2);
+    eval(r#"
+      module GetThousand
+        def test = @var1000
+      end
+      class MyClass
+        extend GetThousand
+      end
+      class String
+        include GetThousand
+      end
+      STR = +''
+      (0..1000).each do |i|
+        ivar_name = :"@var#{i}"
+        MyClass.instance_variable_set(ivar_name, i)
+        STR.instance_variable_set(ivar_name, i)
+      end
+      p MyClass.test; p MyClass.test # profile and compile for MyClass
+      p STR.test
+    "#);
+    assert_snapshot!(assert_compiles("[STR.test, STR.test]"), @"[1000, 1000]");
+}
+
 #[test]
 fn test_attr_accessor_setivar() {
     eval("

zjit/src/cruby.rs

@@ -604,6 +604,7 @@ impl VALUE {
         unsafe { rb_jit_class_fields_embedded_p(self) }
     }
 
+    /// Typed `T_DATA` made from `TypedData_Make_Struct()` (e.g. Thread, ARGF)
     pub fn typed_data_p(self) -> bool {
         !self.special_const_p() &&
             self.builtin_type() == RUBY_T_DATA &&

zjit/src/hir.rs

@@ -4431,6 +4431,23 @@ impl Function {
         }
     }
 
+    /// This puts a guard that establishes the preconditon for [Self::load_ivar]
+    fn load_ivar_guard_type(&mut self, block: BlockId, recv: InsnId, recv_type: ProfiledType, state: InsnId) -> InsnId {
+        if recv_type.class().is_subclass_of(unsafe { rb_cClass }) == ClassRelationship::Subclass {
+            // Check class first since `Class < Module`
+            self.push_insn(block, Insn::GuardType { val: recv, guard_type: types::Class, state })
+        } else if recv_type.class().is_subclass_of(unsafe { rb_cModule }) == ClassRelationship::Subclass {
+            self.push_insn(block, Insn::GuardType { val: recv, guard_type: types::Module, state })
+        } else if recv_type.flags().is_typed_data() {
+            self.push_insn(block, Insn::GuardType { val: recv, guard_type: types::TypedTData, state })
+        } else {
+            // HeapBasicObject is wider than T_OBJECT, but shapes for T_OBJECTs are in a pool of
+            // its own and are guaranteed to be different from shapes of any other T_* types. So
+            // the shape check that follows already covers checking for T_OBJECT.
+            self.push_insn(block, Insn::GuardType { val: recv, guard_type: types::HeapBasicObject, state })
+        }
+    }
+
     fn load_ivar(&mut self, block: BlockId, self_val: InsnId, recv_type: ProfiledType, id: ID, state: InsnId) -> InsnId {
         // Too-complex shapes use hash tables; rb_shape_get_iv_index doesn't support them.
         // Callers must filter these out before calling load_ivar.
@@ -4500,7 +4517,7 @@ impl Function {
                             self.count(block, Counter::getivar_fallback_too_complex);
                             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.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);
                         let replacement = self.load_ivar(block, self_val, recv_type, id, state);
@@ -4529,7 +4546,7 @@ impl Function {
                             self.count(block, Counter::definedivar_fallback_too_complex);
                             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.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);
                         let mut ivar_index: u16 = 0;
@@ -4601,7 +4618,7 @@ impl Function {
                             }
                             // 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.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);
                         // Current shape contains this ivar

zjit/src/hir/opt_tests.rs

@@ -7346,7 +7346,7 @@ mod hir_opt_tests {
           Jump bb3(v4)
         bb3(v6:BasicObject):
           PatchPoint SingleRactorMode
-          v17:HeapBasicObject = GuardType v6, HeapBasicObject
+          v17:Module = GuardType v6, Module
           v18:CShape = LoadField v17, :_shape_id@0x1000
           v19:CShape[0x1001] = GuardBitEquals v18, CShape(0x1001)
           PatchPoint RootBoxOnly
@@ -7381,7 +7381,7 @@ mod hir_opt_tests {
           Jump bb3(v4)
         bb3(v6:BasicObject):
           PatchPoint SingleRactorMode
-          v17:HeapBasicObject = GuardType v6, HeapBasicObject
+          v17:Module = GuardType v6, Module
           v18:CShape = LoadField v17, :_shape_id@0x1000
           v19:CShape[0x1001] = GuardBitEquals v18, CShape(0x1001)
           PatchPoint RootBoxOnly
@@ -7414,7 +7414,7 @@ mod hir_opt_tests {
           Jump bb3(v4)
         bb3(v6:BasicObject):
           PatchPoint SingleRactorMode
-          v17:HeapBasicObject = GuardType v6, HeapBasicObject
+          v17:Class = GuardType v6, Class
           v18:CShape = LoadField v17, :_shape_id@0x1000
           v19:CShape[0x1001] = GuardBitEquals v18, CShape(0x1001)
           PatchPoint RootBoxOnly
@@ -7449,7 +7449,7 @@ mod hir_opt_tests {
           Jump bb3(v4)
         bb3(v6:BasicObject):
           PatchPoint SingleRactorMode
-          v17:HeapBasicObject = GuardType v6, HeapBasicObject
+          v17:Class = GuardType v6, Class
           v18:CShape = LoadField v17, :_shape_id@0x1000
           v19:CShape[0x1001] = GuardBitEquals v18, CShape(0x1001)
           PatchPoint RootBoxOnly
@@ -7520,7 +7520,7 @@ mod hir_opt_tests {
           Jump bb3(v4)
         bb3(v6:BasicObject):
           PatchPoint SingleRactorMode
-          v17:HeapBasicObject = GuardType v6, HeapBasicObject
+          v17:TypedTData = GuardType v6, TypedTData
           v18:CShape = LoadField v17, :_shape_id@0x1000
           v19:CShape[0x1001] = GuardBitEquals v18, CShape(0x1001)
           v20:RubyValue = LoadField v17, :_fields_obj@0x1002
@@ -7556,7 +7556,7 @@ mod hir_opt_tests {
           Jump bb3(v4)
         bb3(v6:BasicObject):
           PatchPoint SingleRactorMode
-          v17:HeapBasicObject = GuardType v6, HeapBasicObject
+          v17:TypedTData = GuardType v6, TypedTData
           v18:CShape = LoadField v17, :_shape_id@0x1000
           v19:CShape[0x1001] = GuardBitEquals v18, CShape(0x1001)
           v20:RubyValue = LoadField v17, :_fields_obj@0x1002

zjit/src/hir_type/gen_hir_type.rb

@@ -132,6 +132,11 @@ def final_type name, base: $object, c_name: nil
 true_exact = final_type "TrueClass", c_name: "rb_cTrueClass"
 false_exact = final_type "FalseClass", c_name: "rb_cFalseClass"
 
+# Typed T_DATA objects (RTYPEDDATA_P). These have a distinct memory layout
+# for field access (fields_obj at a fixed offset in RTypedData). These
+# don't have a common class ancestor below BasicObject.
+basic_object.subtype "TypedTData"
+
 # Build the cvalue object universe. This is for C-level types that may be
 # passed around when calling into the Ruby VM or after some strength reduction
 # of HIR.

zjit/src/hir_type/hir_type.inc.rs

@@ -1,6 +1,7 @@
 //! High-level intermediate representation types.
 
 #![allow(non_upper_case_globals)]
+use crate::cruby;
 use crate::cruby::{rb_block_param_proxy, Qfalse, Qnil, Qtrue, RUBY_T_ARRAY, RUBY_T_CLASS, RUBY_T_HASH, RUBY_T_MODULE, RUBY_T_STRING, VALUE};
 use crate::cruby::{rb_cInteger, rb_cFloat, rb_cArray, rb_cHash, rb_cString, rb_cSymbol, rb_cRange, rb_cModule, rb_zjit_singleton_class_p};
 use crate::cruby::ClassRelationship;
@@ -213,6 +214,7 @@ impl Type {
             else if val.class_of() == unsafe { rb_cSymbol } { bits::DynamicSymbol }
             else if let Some(bits) = Self::bits_from_exact_class(val.class_of()) { bits }
             else if let Some(bits) = Self::bits_from_subclass(val.class_of()) { bits }
+            else if val.typed_data_p() { bits::TypedTData }
             else {
                 unreachable!("Class {} is not a subclass of BasicObject! Don't know what to do.",
                              get_class_name(val.class_of()))
@@ -429,6 +431,24 @@ impl Type {
         matches!(self.spec, Specialization::Object(_))
     }
 
+    /// Find a `T_*` type that is exactly as wide as `self`.
+    pub fn builtin_type_equivalent(&self) -> Option<cruby::ruby_value_type> {
+        if self.bit_equal(types::Array) {
+            Some(cruby::RUBY_T_ARRAY)
+        } else if self.bit_equal(types::Class) {
+            Some(cruby::RUBY_T_CLASS)
+        } else if self.bit_equal(types::Module) {
+            Some(cruby::RUBY_T_MODULE)
+        } else if self.bit_equal(types::String) {
+            Some(cruby::RUBY_T_STRING)
+        } else if self.bit_equal(types::Hash) {
+            Some(cruby::RUBY_T_HASH)
+        } else {
+            // Note that types::TypedTData is narrower than T_DATA, so not here.
+            None
+        }
+    }
+
     fn is_builtin(class: VALUE) -> bool {
         types::ExactBitsAndClass
             .iter()