Invariant trait

library/contracts/safety/src/kani.rs

@@ -1,4 +1,4 @@
-use proc_macro::{TokenStream};
+use proc_macro::TokenStream;
 use quote::{quote, format_ident};
 use syn::{ItemFn, parse_macro_input};
 

library/contracts/safety/src/lib.rs

@@ -3,6 +3,8 @@
 
 use proc_macro::TokenStream;
 use proc_macro_error::proc_macro_error;
+use syn::{ItemStruct, parse_macro_input};
+use quote::quote;
 
 #[cfg(kani_host)]
 #[path = "kani.rs"]
@@ -12,6 +14,27 @@ mod tool;
 #[path = "runtime.rs"]
 mod tool;
 
+#[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)
+}
+
 #[proc_macro_error]
 #[proc_macro_attribute]
 pub fn requires(attr: TokenStream, item: TokenStream) -> TokenStream {

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"]
+#[invariant(self.align.is_safe())]
 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.
@@ -391,7 +395,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 ret = unsafe { mem::transmute::<usize, Alignment>(align) };
+            kani::assert(ret.is_safe(), "Alignment is safe");
+            ret
         }
     }
 

library/core/src/ub_checks.rs

@@ -209,3 +209,23 @@ 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;
+}