From c49757ed6d562c152275df6e60baf6f07fbea7f4 Mon Sep 17 00:00:00 2001
From: "Celina G. Val" <celinval@amazon.com>
Date: Tue, 11 Jun 2024 20:50:31 -0700
Subject: [PATCH] Add a few more contract and harness examples

---
 library/core/src/intrinsics.rs | 51 ++++++++++++++++++++++++++++++++++
 library/core/src/mem/mod.rs    | 40 ++++++++++++++++++++++++++
 2 files changed, 91 insertions(+)

diff --git a/library/core/src/intrinsics.rs b/library/core/src/intrinsics.rs
index 5a2a4c5ae6ebe..ce17fd31fad21 100644
--- a/library/core/src/intrinsics.rs
+++ b/library/core/src/intrinsics.rs
@@ -63,12 +63,16 @@
 )]
 #![allow(missing_docs)]
 
+use safety::requires;
 use crate::marker::DiscriminantKind;
 use crate::marker::Tuple;
 use crate::mem::align_of;
 use crate::ptr;
 use crate::ub_checks;
 
+#[cfg(kani)]
+use crate::kani;
+
 pub mod mir;
 pub mod simd;
 
@@ -2709,6 +2713,13 @@ pub const fn is_val_statically_known<T: Copy>(_arg: T) -> bool {
 #[rustc_intrinsic]
 // This has fallback `const fn` MIR, so shouldn't need stability, see #122652
 #[rustc_const_unstable(feature = "const_typed_swap", issue = "none")]
+#[cfg_attr(kani, kani::modifies(x))]
+#[cfg_attr(kani, kani::modifies(y))]
+#[requires(ub_checks::can_dereference(x) && ub_checks::can_write(x))]
+#[requires(ub_checks::can_dereference(y) && ub_checks::can_write(y))]
+#[requires(x.addr() != y.addr() || core::mem::size_of::<T>() == 0)]
+#[requires(x.addr() > y.addr() || (y.addr() >= x.addr() + core::mem::size_of::<T>()))]
+#[requires(x.addr() < y.addr() || (x.addr() >= y.addr() + core::mem::size_of::<T>()))]
 pub const unsafe fn typed_swap<T>(x: *mut T, y: *mut T) {
     // SAFETY: The caller provided single non-overlapping items behind
     // pointers, so swapping them with `count: 1` is fine.
@@ -3142,3 +3153,43 @@ pub(crate) const fn miri_promise_symbolic_alignment(ptr: *const (), align: usize
 
     const_eval_select((ptr, align), compiletime, runtime);
 }
+
+#[cfg(kani)]
+#[unstable(feature="kani", issue="none")]
+mod verify {
+    use core::{cmp, fmt};
+    use super::*;
+    use crate::kani;
+
+    #[kani::proof_for_contract(typed_swap)]
+    pub fn check_typed_swap_u8() {
+        check_swap::<u8>()
+    }
+
+    #[kani::proof_for_contract(typed_swap)]
+    pub fn check_typed_swap_char() {
+        check_swap::<char>()
+    }
+
+    #[kani::proof_for_contract(typed_swap)]
+    pub fn check_typed_swap_non_zero() {
+        check_swap::<core::num::NonZeroI32>()
+    }
+
+    pub fn check_swap<T: kani::Arbitrary + Copy + cmp::PartialEq + fmt::Debug>() {
+        let mut x = kani::any::<T>();
+        let old_x = x;
+        let x_addr = (&x as *const T).addr();
+
+        let mut y = kani::any::<T>();
+        let old_y = y;
+        let y_addr = (&y as *const T).addr();
+
+        unsafe { typed_swap(&mut x, &mut y) };
+
+        assert_eq!(y, old_x);
+        assert_eq!(x, old_y);
+        assert_eq!(y_addr, (&y as *const T).addr());
+        assert_eq!(x_addr, (&x as *const T).addr());
+    }
+}
diff --git a/library/core/src/mem/mod.rs b/library/core/src/mem/mod.rs
index 9054ade2d7968..cf3f819cd5c15 100644
--- a/library/core/src/mem/mod.rs
+++ b/library/core/src/mem/mod.rs
@@ -13,6 +13,9 @@ use crate::intrinsics;
 use crate::marker::DiscriminantKind;
 use crate::ptr;
 
+#[cfg(kani)]
+use crate::kani;
+
 mod manually_drop;
 #[stable(feature = "manually_drop", since = "1.20.0")]
 pub use manually_drop::ManuallyDrop;
@@ -725,6 +728,8 @@ pub unsafe fn uninitialized<T>() -> T {
 #[stable(feature = "rust1", since = "1.0.0")]
 #[rustc_const_unstable(feature = "const_swap", issue = "83163")]
 #[rustc_diagnostic_item = "mem_swap"]
+#[cfg_attr(kani, crate::kani::modifies(x))]
+#[cfg_attr(kani, crate::kani::modifies(y))]
 pub const fn swap<T>(x: &mut T, y: &mut T) {
     // SAFETY: `&mut` guarantees these are typed readable and writable
     // as well as non-overlapping.
@@ -1345,3 +1350,38 @@ pub macro offset_of($Container:ty, $($fields:expr)+ $(,)?) {
     // The `{}` is for better error messages
     {builtin # offset_of($Container, $($fields)+)}
 }
+
+#[cfg(kani)]
+#[unstable(feature="kani", issue="none")]
+mod verify {
+    use super::*;
+    use crate::kani;
+
+    /// Use this type to ensure that mem swap does not drop the value.
+    #[derive(kani::Arbitrary)]
+    struct CannotDrop<T: kani::Arbitrary> {
+        inner: T,
+    }
+
+    impl<T: kani::Arbitrary> Drop for CannotDrop<T> {
+        fn drop(&mut self) {
+            unreachable!("Cannot drop")
+        }
+    }
+
+    #[kani::proof_for_contract(swap)]
+    pub fn check_swap_primitive() {
+        let mut x: u8 = kani::any();
+        let mut y: u8 = kani::any();
+        swap(&mut x, &mut y)
+    }
+
+    #[kani::proof_for_contract(swap)]
+    pub fn check_swap_adt_no_drop() {
+        let mut x: CannotDrop<char> = kani::any();
+        let mut y: CannotDrop<char> = kani::any();
+        swap(&mut x, &mut y);
+        forget(x);
+        forget(y);
+    }
+}