Calculate layout of Invariant as needed, use enum for size

5225225 · 5225225 · commit 19e0f6f4683e · 2022-06-18T15:04:18.000+01:00
diff --git a/compiler/rustc_const_eval/src/interpret/intrinsics/validity_invariants_of.rs b/compiler/rustc_const_eval/src/interpret/intrinsics/validity_invariants_of.rs
@@ -1,29 +1,48 @@
-use rustc_middle::mir::interpret::{Allocation, ConstAllocation};
-use rustc_middle::mir::Mutability;
+use rustc_hir::lang_items::LangItem;
+use rustc_middle::mir::interpret::{AllocRange, Allocation, ConstAllocation, Scalar as MirScalar};
 use rustc_middle::ty::layout::LayoutCx;
 use rustc_middle::ty::{ParamEnv, ParamEnvAnd};
 use rustc_middle::ty::{Ty, TyCtxt};
 use rustc_target::abi::{
-    Abi, Align, Endian, FieldsShape, HasDataLayout, Scalar, Size, WrappingRange,
+    Abi, FieldsShape, HasDataLayout, Integer, Primitive, Scalar, Size, TyAndLayout, WrappingRange,
 };
 
+#[derive(Debug, Clone, Copy)]
+enum InvariantSize {
+    U8,
+    U16,
+    U32,
+    U64,
+    U128,
+    Pointer,
+}
+
 #[derive(Debug, Clone, Copy)]
 struct Invariant {
     offset: Size,
-    size: Size,
-    start: u128,
-    end: u128,
+    size: InvariantSize,
+    valid_range_start: u128,
+    valid_range_end: u128,
 }
 
-// TODO: Don't add duplicate invariants (maybe use a HashMap?)
+// FIXME: Don't add duplicate invariants (maybe use a HashMap?)
 fn add_invariants<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>, invs: &mut Vec<Invariant>, offset: Size) {
     let x = tcx.layout_of(ParamEnvAnd { param_env: ParamEnv::reveal_all(), value: ty });
 
     if let Ok(layout) = x {
         if let Abi::Scalar(Scalar::Initialized { value, valid_range }) = layout.layout.abi() {
-            let size = value.size(&tcx);
+            let size = match value {
+                Primitive::Int(Integer::I8, _) => InvariantSize::U8,
+                Primitive::Int(Integer::I16, _) => InvariantSize::U16,
+                Primitive::Int(Integer::I32, _) => InvariantSize::U32,
+                Primitive::Int(Integer::I64, _) => InvariantSize::U64,
+                Primitive::Int(Integer::I128, _) => InvariantSize::U128,
+                Primitive::F32 => InvariantSize::U32,
+                Primitive::F64 => InvariantSize::U64,
+                Primitive::Pointer => InvariantSize::Pointer,
+            };
             let WrappingRange { start, end } = valid_range;
-            invs.push(Invariant { offset, size, start, end })
+            invs.push(Invariant { offset, size, valid_range_start: start, valid_range_end: end })
         }
 
         let param_env = ParamEnv::reveal_all();
@@ -33,8 +52,8 @@ fn add_invariants<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>, invs: &mut Vec<Invarian
             FieldsShape::Primitive => {}
             FieldsShape::Union(_) => {}
             FieldsShape::Array { stride, count } => {
-                // TODO: should we just bail if we're making a Too Large type?
-                // (Like [bool; 1_000_000])
+                // We may wish to bail out if we're generating too many invariants.
+                // That would lead to false negatives, though.
                 for idx in 0..*count {
                     let off = offset + *stride * idx;
                     let f = layout.field(&unwrap, idx as usize);
@@ -57,22 +76,13 @@ fn add_invariants<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>, invs: &mut Vec<Invarian
     }
 }
 
-fn extend_encoded_int(to: &mut Vec<u8>, endian: Endian, ptr_size: PointerSize, value: Size) {
-    match (endian, ptr_size) {
-        (Endian::Little, PointerSize::Bits16) => to.extend((value.bytes() as u16).to_le_bytes()),
-        (Endian::Little, PointerSize::Bits32) => to.extend((value.bytes() as u32).to_le_bytes()),
-        (Endian::Little, PointerSize::Bits64) => to.extend((value.bytes()).to_le_bytes()),
-        (Endian::Big, PointerSize::Bits16) => to.extend((value.bytes() as u16).to_be_bytes()),
-        (Endian::Big, PointerSize::Bits32) => to.extend((value.bytes() as u32).to_be_bytes()),
-        (Endian::Big, PointerSize::Bits64) => to.extend((value.bytes()).to_be_bytes()),
-    }
-}
-
-#[derive(Clone, Copy)]
-enum PointerSize {
-    Bits16,
-    Bits32,
-    Bits64,
+fn get_layout_of_invariant<'tcx>(tcx: TyCtxt<'tcx>) -> TyAndLayout<'tcx, Ty<'tcx>> {
+    let item = tcx.require_lang_item(LangItem::ValidityInvariant, None);
+    let ty = tcx.type_of(item);
+    let layout = tcx
+        .layout_of(ParamEnv::reveal_all().and(ty))
+        .expect("invalid layout for ValidityInvariant lang item");
+    layout
 }
 
 /// Directly returns a `ConstAllocation` containing a list of validity invariants of the given type.
@@ -83,38 +93,54 @@ pub(crate) fn alloc_validity_invariants_of<'tcx>(
     let mut invs: Vec<Invariant> = Vec::new();
 
     let layout = tcx.data_layout();
-
-    let ptr_size = match layout.pointer_size.bits() {
-        16 => PointerSize::Bits16,
-        32 => PointerSize::Bits32,
-        64 => PointerSize::Bits64,
-        _ => {
-            // Not sure if this can happen, but just return an empty slice?
-            let alloc =
-                Allocation::from_bytes(Vec::new(), Align::from_bytes(8).unwrap(), Mutability::Not);
-            return tcx.intern_const_alloc(alloc);
-        }
-    };
+    let validity_invariant = get_layout_of_invariant(tcx);
 
     add_invariants(tcx, ty, &mut invs, Size::ZERO);
 
-    let encode_range = match layout.endian {
-        Endian::Little => |r: u128| r.to_le_bytes(),
-        Endian::Big => |r: u128| r.to_be_bytes(),
-    };
+    let allocation_size = validity_invariant.layout.size() * invs.len() as u64;
+    let mut alloc =
+        Allocation::uninit(allocation_size, validity_invariant.layout.align().abi, true).unwrap();
+
+    let offset_off = validity_invariant.layout.fields().offset(0);
+    let size_off = validity_invariant.layout.fields().offset(1);
+    let start_off = validity_invariant.layout.fields().offset(2);
+    let end_off = validity_invariant.layout.fields().offset(3);
+
+    for (idx, invariant) in invs.iter().enumerate() {
+        let offset = idx as u64 * validity_invariant.layout.size();
+
+        let offset_range = AllocRange { start: offset + offset_off, size: layout.pointer_size };
+        alloc
+            .write_scalar(
+                &tcx,
+                offset_range,
+                MirScalar::from_machine_usize(invariant.offset.bytes(), &tcx).into(),
+            )
+            .unwrap();
+
+        let size_range = AllocRange { start: offset + size_off, size: Size::from_bytes(1) };
+        alloc
+            .write_scalar(&tcx, size_range, MirScalar::from_u8(invariant.size as u8).into())
+            .unwrap();
 
-    let mut encoded = Vec::new();
+        let offset_range = AllocRange { start: offset + start_off, size: Size::from_bytes(16) };
+        alloc
+            .write_scalar(
+                &tcx,
+                offset_range,
+                MirScalar::from_u128(invariant.valid_range_start).into(),
+            )
+            .unwrap();
 
-    // TODO: this needs to match the layout of `Invariant` in core/src/intrinsics.rs
-    // how do we ensure that?
-    for inv in invs {
-        extend_encoded_int(&mut encoded, layout.endian, ptr_size, inv.offset);
-        extend_encoded_int(&mut encoded, layout.endian, ptr_size, inv.size);
-        encoded.extend(encode_range(inv.start));
-        encoded.extend(encode_range(inv.end));
+        let offset_range = AllocRange { start: offset + end_off, size: Size::from_bytes(16) };
+        alloc
+            .write_scalar(
+                &tcx,
+                offset_range,
+                MirScalar::from_u128(invariant.valid_range_end).into(),
+            )
+            .unwrap();
     }
 
-    // TODO: The alignment here should be calculated from the struct definition, I guess?
-    let alloc = Allocation::from_bytes(encoded, Align::from_bytes(8).unwrap(), Mutability::Not);
     tcx.intern_const_alloc(alloc)
 }
diff --git a/compiler/rustc_hir/src/lang_items.rs b/compiler/rustc_hir/src/lang_items.rs
@@ -327,6 +327,7 @@ language_item_table! {
     Range,                   sym::Range,               range_struct,               Target::Struct,         GenericRequirement::None;
     RangeToInclusive,        sym::RangeToInclusive,    range_to_inclusive_struct,  Target::Struct,         GenericRequirement::None;
     RangeTo,                 sym::RangeTo,             range_to_struct,            Target::Struct,         GenericRequirement::None;
+    ValidityInvariant,       sym::ValidityInvariant,   validity_invariant_struct,  Target::Struct,         GenericRequirement::None;
 }
 
 pub enum GenericRequirement {
diff --git a/compiler/rustc_middle/src/mir/interpret/value.rs b/compiler/rustc_middle/src/mir/interpret/value.rs
@@ -258,6 +258,11 @@ impl<Tag> Scalar<Tag> {
         Scalar::Int(i.into())
     }
 
+    #[inline]
+    pub fn from_u128(i: u128) -> Self {
+        Scalar::Int(i.into())
+    }
+
     #[inline]
     pub fn from_machine_usize(i: u64, cx: &impl HasDataLayout) -> Self {
         Self::from_uint(i, cx.data_layout().pointer_size)
diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs
@@ -277,6 +277,7 @@ symbols! {
         TyKind,
         Unknown,
         UnsafeArg,
+        ValidityInvariant,
         Vec,
         VecDeque,
         Wrapper,
diff --git a/library/core/src/intrinsics.rs b/library/core/src/intrinsics.rs
@@ -2138,18 +2138,42 @@ pub const unsafe fn copy_nonoverlapping<T>(src: *const T, dst: *mut T, count: us
     }
 }
 
-#[repr(C)]
 #[derive(Debug, Clone, Copy)]
-pub struct Invariant {
-    offset: u64,
-    size: u64,
-    start: u128,
-    end: u128,
+#[repr(u8)]
+enum InvariantSize {
+    U8 = 0,
+    U16 = 1,
+    U32 = 2,
+    U64 = 3,
+    U128 = 4,
+    Pointer = 5,
+}
+
+/// An Invariant is a field and its valid range. The valid range may be full, in which case it
+/// should still be branched on, to allow tools like memory sanitizer to verify the memory is not
+/// uninit.
+// Be sure to keep the fields in order (offset, size, start, end), validity_invariants_of.rs needs
+// that.
+#[cfg(not(bootstrap))]
+#[derive(Debug, Clone, Copy)]
+#[lang = "ValidityInvariant"]
+struct Invariant {
+    /// The offset in bytes from the start of the struct
+    offset: usize,
+    /// The size/type of the invariant.
+    size: InvariantSize,
+    /// The start point of the valid range of this field. Is allowed to be > valid_range_end, see
+    /// <https://doc.rust-lang.org/nightly/nightly-rustc/rustc_target/abi/struct.WrappingRange.html>
+    /// which this follows the semantics of.
+    valid_range_start: u128,
+
+    /// The end point of the range.
+    valid_range_end: u128,
 }
 
 #[cfg(not(bootstrap))]
 /// Returns a list of all validity invariants of the type.
-pub const fn validity_invariants_of<T>() -> &'static [Invariant] {
+const fn validity_invariants_of<T>() -> &'static [Invariant] {
     extern "rust-intrinsic" {
         #[rustc_const_unstable(feature = "validity_invariants_of", issue = "none")]
         pub fn validity_invariants_of<T>() -> &'static [u8];
@@ -2435,34 +2459,35 @@ pub(crate) unsafe fn assert_validity_of<T>(value: *const T) -> bool {
     struct Unaligned<T>(T);
 
     // SAFETY: The pointer dereferences here are valid if `value` is valid.
-    // though TODO: introduce a new size for "pointer", since reading a pointer as an int *is* UB.
     unsafe {
         let invariants = validity_invariants_of::<T>();
         for invariant in invariants {
             let off = invariant.offset as usize;
-            let start = invariant.start;
-            let end = invariant.end;
+            let start = invariant.valid_range_start;
+            let end = invariant.valid_range_end;
 
-            // TODO: Maybe replace this with an enum?
             let (value, max): (u128, u128) = match invariant.size {
-                1 => ((*(value.cast::<u8>().add(off))).into(), u8::MAX.into()),
-                2 => (
+                InvariantSize::U8 => ((*(value.cast::<u8>().add(off))).into(), u8::MAX.into()),
+                InvariantSize::U16 => (
                     (*value.cast::<u8>().add(off).cast::<Unaligned<u16>>()).0.into(),
                     u16::MAX.into(),
                 ),
-                4 => (
+                InvariantSize::U32 => (
                     (*value.cast::<u8>().add(off).cast::<Unaligned<u32>>()).0.into(),
                     u32::MAX.into(),
                 ),
-                8 => (
+                InvariantSize::U64 => (
                     (*value.cast::<u8>().add(off).cast::<Unaligned<u64>>()).0.into(),
                     u64::MAX.into(),
                 ),
-                16 => (
+                InvariantSize::U128 => (
                     (*value.cast::<u8>().add(off).cast::<Unaligned<u128>>()).0.into(),
                     u128::MAX.into(),
                 ),
-                s => panic!("unexpected size {s}"),
+                InvariantSize::Pointer => (
+                    (*value.cast::<u8>().add(off).cast::<Unaligned<*const ()>>()).0.addr() as u128,
+                    usize::MAX as u128,
+                ),
             };
 
             if start > end {

Original file line number	Diff line number	Diff line change
`@@ -327,6 +327,7 @@ language_item_table! {`
`327`	`327`	`Range, sym::Range, range_struct, Target::Struct, GenericRequirement::None;`
`328`	`328`	`RangeToInclusive, sym::RangeToInclusive, range_to_inclusive_struct, Target::Struct, GenericRequirement::None;`
`329`	`329`	`RangeTo, sym::RangeTo, range_to_struct, Target::Struct, GenericRequirement::None;`
	`330`	`+ ValidityInvariant, sym::ValidityInvariant, validity_invariant_struct, Target::Struct, GenericRequirement::None;`
`330`	`331`	`}`
`331`	`332`
`332`	`333`	`pub enum GenericRequirement {`