Invariant trait

library/contracts/safety/src/kani.rs

@@ -1,6 +1,6 @@
-use proc_macro::{TokenStream};
-use quote::{quote, format_ident};
-use syn::{ItemFn, parse_macro_input};
+use proc_macro::TokenStream;
+use quote::{format_ident, quote};
+use syn::{parse_macro_input, ItemFn};
 
 pub(crate) fn requires(attr: TokenStream, item: TokenStream) -> TokenStream {
     rewrite_attr(attr, item, "requires")
@@ -17,5 +17,6 @@ fn rewrite_attr(attr: TokenStream, item: TokenStream, name: &str) -> TokenStream
     quote!(
         #[kani_core::#attribute(#args)]
         #fn_item
-    ).into()
+    )
+    .into()
 }

library/contracts/safety/src/lib.rs

@@ -3,6 +3,11 @@
 
 use proc_macro::TokenStream;
 use proc_macro_error::proc_macro_error;
+use quote::{quote, quote_spanned};
+use syn::{
+    parse_macro_input, parse_quote, spanned::Spanned, Data, DeriveInput, Fields, GenericParam,
+    Generics, Index, ItemStruct,
+};
 
 #[cfg(kani_host)]
 #[path = "kani.rs"]
@@ -12,6 +17,101 @@ mod tool;
 #[path = "runtime.rs"]
 mod tool;
 
+/// Expands the `#[invariant(...)]` attribute macro.
+/// The macro expands to an implementation of the `is_safe` method for the `Invariant` trait.
+/// This attribute is only supported for structs.
+///
+/// # Example
+///
+/// ```ignore
+/// #[invariant(self.width == self.height)]
+/// struct Square {
+///     width: u32,
+///     height: u32,
+/// }
+/// ```
+/// 
+/// expands to:
+/// ```ignore
+/// impl core::ub_checks::Invariant for Square {
+///   fn is_safe(&self) -> bool {
+///     self.width == self.height
+///   }
+/// }
+/// ```
+/// For more information on the Invariant trait, see its documentation in core::ub_checks.
+#[proc_macro_error]
+#[proc_macro_attribute]
+pub fn invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
+    let safe_body = proc_macro2::TokenStream::from(attr);
+    let item = parse_macro_input!(item as ItemStruct);
+    let item_name = &item.ident;
+    let (impl_generics, ty_generics, where_clause) = item.generics.split_for_impl();
+
+    let expanded = quote! {
+        #item
+        #[unstable(feature="invariant", issue="none")]
+        impl #impl_generics core::ub_checks::Invariant for #item_name #ty_generics #where_clause {
+            fn is_safe(&self) -> bool {
+                #safe_body
+            }
+        }
+    };
+
+    proc_macro::TokenStream::from(expanded)
+}
+
+/// Expands the derive macro for the Invariant trait.
+/// The macro expands to an implementation of the `is_safe` method for the `Invariant` trait.
+/// This macro is only supported for structs.
+///
+/// # Example
+///
+/// ```ignore
+/// #[derive(Invariant)]
+/// struct Square {
+///     width: u32,
+///     height: u32,
+/// }
+/// ```
+/// 
+/// expands to:
+/// ```ignore
+/// impl core::ub_checks::Invariant for Square {
+///   fn is_safe(&self) -> bool {
+///     self.width.is_safe() && self.height.is_safe()
+///   }
+/// }
+/// ```
+/// For more information on the Invariant trait, see its documentation in core::ub_checks.
+#[proc_macro_error]
+#[proc_macro_derive(Invariant)]
+pub fn derive_invariant(item: TokenStream) -> TokenStream {
+    let derive_item = parse_macro_input!(item as DeriveInput);
+    let item_name = &derive_item.ident;
+    if let Data::Struct(struct_data) = derive_item.data {
+        let safe_body = safe_body(&struct_data.fields);
+
+        // Add a bound `T: Invariant` to every type parameter T.
+        let generics = add_trait_bound_invariant(derive_item.generics);
+        // Generate an expression to sum up the heap size of each field.
+        let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
+
+        let expanded = quote! {
+            // The generated implementation.
+            #[unstable(feature="invariant", issue="none")]
+            impl #impl_generics core::ub_checks::Invariant for #item_name #ty_generics #where_clause {
+                fn is_safe(&self) -> bool {
+                    #safe_body
+                }
+            }
+        };
+        proc_macro::TokenStream::from(expanded)
+    } else {
+        panic!("Attempted to derive the Invariant trait on a non-struct type.")
+    }
+}
+
 #[proc_macro_error]
 #[proc_macro_attribute]
 pub fn requires(attr: TokenStream, item: TokenStream) -> TokenStream {
@@ -23,3 +123,58 @@ pub fn requires(attr: TokenStream, item: TokenStream) -> TokenStream {
 pub fn ensures(attr: TokenStream, item: TokenStream) -> TokenStream {
     tool::ensures(attr, item)
 }
+
+/// Add a bound `T: Invariant` to every type parameter T.
+fn add_trait_bound_invariant(mut generics: Generics) -> Generics {
+    generics.params.iter_mut().for_each(|param| {
+        if let GenericParam::Type(type_param) = param {
+            type_param
+                .bounds
+                .push(parse_quote!(core::ub_checks::Invariant));
+        }
+    });
+    generics
+}
+
+/// Generate the body for the `is_safe` method.
+/// For each field of the type, enforce that it is safe.
+fn safe_body(fields: &Fields) -> proc_macro2::TokenStream {
+    match fields {
+        Fields::Named(ref fields) => {
+            let field_safe_calls: Vec<proc_macro2::TokenStream> = fields
+                .named
+                .iter()
+                .map(|field| {
+                    let name = &field.ident;
+                    quote_spanned! {field.span()=>
+                        self.#name.is_safe()
+                    }
+                })
+                .collect();
+            if !field_safe_calls.is_empty() {
+                quote! { #( #field_safe_calls )&&* }
+            } else {
+                quote! { true }
+            }
+        }
+        Fields::Unnamed(ref fields) => {
+            let field_safe_calls: Vec<proc_macro2::TokenStream> = fields
+                .unnamed
+                .iter()
+                .enumerate()
+                .map(|(idx, field)| {
+                    let field_idx = Index::from(idx);
+                    quote_spanned! {field.span()=>
+                        self.#field_idx.is_safe()
+                    }
+                })
+                .collect();
+            if !field_safe_calls.is_empty() {
+                quote! { #( #field_safe_calls )&&* }
+            } else {
+                quote! { true }
+            }
+        }
+        Fields::Unit => quote! { true },
+    }
+}

library/core/src/alloc/layout.rs

@@ -4,14 +4,18 @@
 // collections, resulting in having to optimize down excess IR multiple times.
 // Your performance intuition is useless. Run perf.
 
-use safety::requires;
+use safety::{ensures, Invariant, requires};
 use crate::error::Error;
 use crate::ptr::{Alignment, NonNull};
 use crate::{assert_unsafe_precondition, cmp, fmt, mem};
 
 #[cfg(kani)]
 use crate::kani;
 
+// Used only for contract verification.
+#[allow(unused_imports)]
+use crate::ub_checks::Invariant;
+
 // While this function is used in one place and its implementation
 // could be inlined, the previous attempts to do so made rustc
 // slower:
@@ -39,6 +43,7 @@ const fn size_align<T>() -> (usize, usize) {
 #[stable(feature = "alloc_layout", since = "1.28.0")]
 #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
 #[lang = "alloc_layout"]
+#[derive(Invariant)]
 pub struct Layout {
     // size of the requested block of memory, measured in bytes.
     size: usize,
@@ -128,6 +133,7 @@ impl Layout {
     #[inline]
     #[rustc_allow_const_fn_unstable(ptr_alignment_type)]
     #[requires(Layout::from_size_align(size, align).is_ok())]
+    #[ensures(|layout| layout.is_safe())]
     pub const unsafe fn from_size_align_unchecked(size: usize, align: usize) -> Self {
         assert_unsafe_precondition!(
             check_library_ub,

library/core/src/ptr/alignment.rs

@@ -1,11 +1,14 @@
-use safety::{ensures, requires};
+use safety::{ensures, invariant, requires};
 use crate::num::NonZero;
 use crate::ub_checks::assert_unsafe_precondition;
 use crate::{cmp, fmt, hash, mem, num};
 
 #[cfg(kani)]
 use crate::kani;
 
+#[cfg(kani)]
+use crate::ub_checks::Invariant;
+
 /// A type storing a `usize` which is a power of two, and thus
 /// represents a possible alignment in the Rust abstract machine.
 ///
@@ -14,6 +17,7 @@ use crate::kani;
 #[unstable(feature = "ptr_alignment_type", issue = "102070")]
 #[derive(Copy, Clone, PartialEq, Eq)]
 #[repr(transparent)]
+#[invariant(self.as_usize().is_power_of_two())]
 pub struct Alignment(AlignmentEnum);
 
 // Alignment is `repr(usize)`, but via extra steps.
@@ -256,6 +260,7 @@ impl Default for Alignment {
 #[cfg(target_pointer_width = "16")]
 #[derive(Copy, Clone, PartialEq, Eq)]
 #[repr(u16)]
+#[cfg_attr(kani, derive(kani::Arbitrary))]
 enum AlignmentEnum {
     _Align1Shl0 = 1 << 0,
     _Align1Shl1 = 1 << 1,
@@ -278,6 +283,7 @@ enum AlignmentEnum {
 #[cfg(target_pointer_width = "32")]
 #[derive(Copy, Clone, PartialEq, Eq)]
 #[repr(u32)]
+#[cfg_attr(kani, derive(kani::Arbitrary))]
 enum AlignmentEnum {
     _Align1Shl0 = 1 << 0,
     _Align1Shl1 = 1 << 1,
@@ -316,6 +322,7 @@ enum AlignmentEnum {
 #[cfg(target_pointer_width = "64")]
 #[derive(Copy, Clone, PartialEq, Eq)]
 #[repr(u64)]
+#[cfg_attr(kani, derive(kani::Arbitrary))]
 enum AlignmentEnum {
     _Align1Shl0 = 1 << 0,
     _Align1Shl1 = 1 << 1,
@@ -390,8 +397,9 @@ mod verify {
 
     impl kani::Arbitrary for Alignment {
         fn any() -> Self {
-            let align = kani::any_where(|a: &usize| a.is_power_of_two());
-            unsafe { mem::transmute::<usize, Alignment>(align) }
+            let obj = Self { 0: kani::any() };
+            kani::assume(obj.is_safe());
+            obj
         }
     }
 

library/core/src/ub_checks.rs

@@ -209,3 +209,61 @@ mod predicates {
 mod predicates {
     pub use crate::kani::mem::{can_dereference, can_write, can_read_unaligned, can_write_unaligned};
 }
+
+/// This trait should be used to specify and check type safety invariants for a
+/// type. For type invariants, we refer to the definitions in the Rust's Unsafe
+/// Code Guidelines Reference:
+/// <https://rust-lang.github.io/unsafe-code-guidelines/glossary.html#validity-and-safety-invariant>
+///
+/// In summary, the reference distinguishes two kinds of type invariants:
+///  - *Validity invariant*: An invariant that all data must uphold any time
+///    it's accessed or copied in a typed manner. This invariant is exploited by
+///    the compiler to perform optimizations.
+///  - *Safety invariant*: An invariant that safe code may assume all data to
+///    uphold. This invariant can be temporarily violated by unsafe code, but
+///    must always be upheld when interfacing with unknown safe code.
+///
+/// Therefore, validity invariants must be upheld at all times, while safety
+/// invariants only need to be upheld at the boundaries to safe code.
+pub trait Invariant {
+    /// Specify the type's safety invariants 
+    fn is_safe(&self) -> bool;
+}
+
+/// Any value is considered safe for the type
+macro_rules! trivial_invariant {
+    ( $type: ty ) => {
+        impl Invariant for $type {
+            #[inline(always)]
+            fn is_safe(&self) -> bool {
+                true
+            }
+        }
+    };
+}
+
+trivial_invariant!(u8);
+trivial_invariant!(u16);
+trivial_invariant!(u32);
+trivial_invariant!(u64);
+trivial_invariant!(u128);
+trivial_invariant!(usize);
+
+trivial_invariant!(i8);
+trivial_invariant!(i16);
+trivial_invariant!(i32);
+trivial_invariant!(i64);
+trivial_invariant!(i128);
+trivial_invariant!(isize);
+
+trivial_invariant!(());
+trivial_invariant!(bool);
+trivial_invariant!(char);
+
+// We do not constrain the safety invariant for floating point types.
+// Users can create a newtype wrapping the floating point type and define an
+// invariant that checks for NaN, infinite, or subnormal values.
+trivial_invariant!(f16);
+trivial_invariant!(f32);
+trivial_invariant!(f64);
+trivial_invariant!(f128);