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, Stmt};
+use proc_macro::TokenStream;
+use quote::{format_ident, quote};
+use syn::{parse_macro_input, ItemFn, Stmt};
 
 pub(crate) fn requires(attr: TokenStream, item: TokenStream) -> TokenStream {
     rewrite_attr(attr, item, "requires")
@@ -21,7 +21,8 @@ fn rewrite_stmt_attr(attr: TokenStream, stmt_stream: TokenStream, name: &str) ->
     quote!(
         #[kani_core::#attribute(#args)]
         #stmt
-    ).into()
+    )
+    .into()
 }
 
 fn rewrite_attr(attr: TokenStream, item: TokenStream, name: &str) -> TokenStream {
@@ -31,5 +32,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::{format_ident, quote, quote_spanned};
+use syn::{
+    parse_macro_input, parse_quote, spanned::Spanned, Data, DataEnum, DeriveInput, Fields,
+    GenericParam, Generics, Ident, Index, ItemStruct,
+};
 
 #[cfg(kani_host)]
 #[path = "kani.rs"]
@@ -12,6 +17,135 @@ 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 and enums.
+///
+/// # 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 enums, the body of `is_safe` matches on the variant and calls `is_safe` on its fields,
+/// # Example
+///
+/// ```ignore
+/// #[derive(Invariant)]
+/// enum MyEnum {
+///     OptionOne(u32, u32),
+///     OptionTwo(Square),
+///     OptionThree
+/// }
+/// ```
+///
+/// expands to:
+/// ```ignore
+/// impl core::ub_checks::Invariant for MyEnum {
+///   fn is_safe(&self) -> bool {
+///     match self {
+///       MyEnum::OptionOne(field1, field2) => field1.is_safe() && field2.is_safe(),
+///       MyEnum::OptionTwo(field1) => field1.is_safe(),
+///       MyEnum::OptionThree => true,
+///     }
+///    }
+/// }
+/// ```
+/// 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;
+    let safe_body = match derive_item.data {
+        Data::Struct(struct_data) => {
+            safe_body(&struct_data.fields)
+        },
+        Data::Enum(enum_data) => {
+            let variant_checks = variant_checks(enum_data, item_name);
+
+            quote! {
+                match self {
+                    #(#variant_checks),*
+                }
+            }
+        },
+        Data::Union(..) => unimplemented!("Attempted to derive Invariant on a union; Invariant can only be derived for structs and enums."),
+    };
+
+    // 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)
+}
+
 #[proc_macro_error]
 #[proc_macro_attribute]
 pub fn requires(attr: TokenStream, item: TokenStream) -> TokenStream {
@@ -29,3 +163,96 @@ pub fn ensures(attr: TokenStream, item: TokenStream) -> TokenStream {
 pub fn loop_invariant(attr: TokenStream, stmt_stream: TokenStream) -> TokenStream {
     tool::loop_invariant(attr, stmt_stream)
 }
+
+/// 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 safety checks for each variant of an enum
+fn variant_checks(enum_data: DataEnum, item_name: &Ident) -> Vec<proc_macro2::TokenStream> {
+    enum_data
+        .variants
+        .iter()
+        .map(|variant| {
+            let variant_name = &variant.ident;
+            match &variant.fields {
+                Fields::Unnamed(fields) => {
+                    let field_names: Vec<_> = fields
+                        .unnamed
+                        .iter()
+                        .enumerate()
+                        .map(|(i, _)| format_ident!("field{}", i + 1))
+                        .collect();
+
+                    let field_checks: Vec<_> = field_names
+                        .iter()
+                        .map(|field_name| {
+                            quote! { #field_name.is_safe() }
+                        })
+                        .collect();
+
+                    quote! {
+                        #item_name::#variant_name(#(#field_names),*) => #(#field_checks)&&*
+                    }
+                }
+                Fields::Unit => {
+                    quote! {
+                        #item_name::#variant_name => true
+                    }
+                }
+                Fields::Named(_) => unreachable!("Enums do not have named fields"),
+            }
+        })
+        .collect()
+}
+
+/// 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::{ensures, 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,
@@ -132,6 +137,7 @@ impl Layout {
     #[inline]
     #[rustc_allow_const_fn_unstable(ptr_alignment_type)]
     #[requires(Layout::from_size_align(size, align).is_ok())]
+    #[ensures(|result| result.is_safe())]
     #[ensures(|result| result.size() == size)]
     #[ensures(|result| result.align() == align)]
     pub const unsafe fn from_size_align_unchecked(size: usize, align: usize) -> Self {

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
         }
     }