From 746b5dede0bfd53a89f2859de2b96db101c574b9 Mon Sep 17 00:00:00 2001
From: Qinheping Hu <qinhh@amazon.com>
Date: Tue, 22 Oct 2024 01:38:39 -0500
Subject: [PATCH 1/2] Add loop contracts and harness for BinaryHeap::sift_up

---
 library/alloc/Cargo.toml                      |  1 +
 .../alloc/src/collections/binary_heap/mod.rs  | 66 +++++++++++++++++++
 library/alloc/src/lib.rs                      |  2 +
 scripts/check_kani.sh                         |  3 +-
 4 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/library/alloc/Cargo.toml b/library/alloc/Cargo.toml
index 1bd4434d4f7e9..af76ac7e23458 100644
--- a/library/alloc/Cargo.toml
+++ b/library/alloc/Cargo.toml
@@ -10,6 +10,7 @@ edition = "2021"
 
 [dependencies]
 core = { path = "../core" }
+safety = { path = "../contracts/safety" }
 compiler_builtins = { version = "0.1.123", features = ['rustc-dep-of-std'] }
 
 [dev-dependencies]
diff --git a/library/alloc/src/collections/binary_heap/mod.rs b/library/alloc/src/collections/binary_heap/mod.rs
index fe9f1010d327c..1de0214520e05 100644
--- a/library/alloc/src/collections/binary_heap/mod.rs
+++ b/library/alloc/src/collections/binary_heap/mod.rs
@@ -155,6 +155,12 @@ use crate::collections::TryReserveError;
 use crate::slice;
 use crate::vec::{self, AsVecIntoIter, Vec};
 
+use safety::{ensures, requires};
+
+#[cfg(kani)]
+#[unstable(feature="kani", issue="none")]
+use core::kani;
+
 #[cfg(test)]
 mod tests;
 
@@ -672,11 +678,13 @@ impl<T: Ord, A: Allocator> BinaryHeap<T, A> {
     /// # Safety
     ///
     /// The caller must guarantee that `pos < self.len()`.
+    #[requires(pos < self.len() && start <= pos)]
     unsafe fn sift_up(&mut self, start: usize, pos: usize) -> usize {
         // Take out the value at `pos` and create a hole.
         // SAFETY: The caller guarantees that pos < self.len()
         let mut hole = unsafe { Hole::new(&mut self.data, pos) };
 
+        #[cfg_attr(kani, kani::loop_invariant(hole.pos() <= pos))]
         while hole.pos() > start {
             let parent = (hole.pos() - 1) / 2;
 
@@ -1897,3 +1905,61 @@ impl<'a, T: 'a + Ord + Copy, A: Allocator> Extend<&'a T> for BinaryHeap<T, A> {
         self.reserve(additional);
     }
 }
+
+#[cfg(kani)]
+#[unstable(feature = "kani", issue = "none")]
+mod verify {
+    use super::*;
+    use crate::boxed::Box;
+
+    /// Generates an arbitrary vector whose length is at most MAX_LENGTH.
+    pub fn any_vec<T, const MAX_LENGTH: usize>() -> Vec<T>
+    where
+        T: kani::Arbitrary,
+    {
+        let real_length: usize = kani::any_where(|sz| *sz <= MAX_LENGTH);
+        match real_length {
+            0 => vec![],
+            exact if exact == MAX_LENGTH => exact_vec::<T, MAX_LENGTH>(),
+            _ => {
+                let mut any_vec = exact_vec::<T, MAX_LENGTH>();
+                any_vec.truncate(real_length);
+                any_vec.shrink_to_fit();
+                assert!(any_vec.capacity() == any_vec.len());
+                any_vec
+            }
+        }
+    }
+
+    /// Generates an arbitrary vector that is exactly EXACT_LENGTH long.
+    pub fn exact_vec<T, const EXACT_LENGTH: usize>() -> Vec<T>
+    where
+        T: kani::Arbitrary,
+    {
+        let boxed_array: Box<[T; EXACT_LENGTH]> = Box::new(kani::any());
+        <[T]>::into_vec(boxed_array)
+    }
+
+    fn any_binary_heap<T: kani::Arbitrary, const MAX_SIZE: usize>() -> BinaryHeap<T> {
+        let data: Vec<T> = any_vec::<T, MAX_SIZE>();
+        BinaryHeap { data }
+    }
+
+    #[kani::proof]
+    pub fn check_any_binary_heap() {
+        let h: BinaryHeap<u8> = any_binary_heap::<u8, 8>();
+        assert!(h.len() <= 8);
+    }
+
+    #[kani::proof]
+    pub fn check_sift_up() {
+        let _ = Box::new(0);
+        const MAX_SIZE: usize = 1000;
+        let mut h: BinaryHeap<u8> = any_binary_heap::<u8, MAX_SIZE>();
+        let pos: usize = kani::any_where(|x| *x < h.len());
+        let start: usize = kani::any_where(|x| *x <= pos);
+        unsafe {
+            h.sift_up(start, pos);
+        }
+    }
+}
diff --git a/library/alloc/src/lib.rs b/library/alloc/src/lib.rs
index 7aaa4e73df72c..ffcf54f630965 100644
--- a/library/alloc/src/lib.rs
+++ b/library/alloc/src/lib.rs
@@ -192,6 +192,8 @@
 #![feature(unboxed_closures)]
 #![feature(unsized_fn_params)]
 #![feature(with_negative_coherence)]
+#![cfg_attr(kani, feature(proc_macro_hygiene))]
+#![cfg_attr(kani, feature(kani))]
 #![rustc_preserve_ub_checks]
 // tidy-alphabetical-end
 //
diff --git a/scripts/check_kani.sh b/scripts/check_kani.sh
index c93499cb7a398..c0b79946b1c5b 100644
--- a/scripts/check_kani.sh
+++ b/scripts/check_kani.sh
@@ -44,7 +44,8 @@ cargo build-dev --release
 echo "Running tests..."
 echo
 cd "$VERIFY_RUST_STD_DIR"
-$KANI_DIR/scripts/kani verify-std -Z unstable-options $VERIFY_RUST_STD_DIR/library --target-dir "$RUNNER_TEMP" -Z function-contracts -Z mem-predicates
+$KANI_DIR/scripts/kani verify-std -Z unstable-options $VERIFY_RUST_STD_DIR/library --target-dir "$RUNNER_TEMP" -Z function-contracts -Z mem-predicates -Z loop-contracts --enable-unstable --cbmc-args --object-bits 8
+
 
 echo "Tests completed."
 echo

From eed40945837d39247fc1d7aeb5c0c1e5527b284f Mon Sep 17 00:00:00 2001
From: Michael Tautschnig <mt@debian.org>
Date: Tue, 22 Oct 2024 22:17:51 +0200
Subject: [PATCH 2/2] Update library/alloc/src/collections/binary_heap/mod.rs

---
 library/alloc/src/collections/binary_heap/mod.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/library/alloc/src/collections/binary_heap/mod.rs b/library/alloc/src/collections/binary_heap/mod.rs
index 1de0214520e05..18855d16cd766 100644
--- a/library/alloc/src/collections/binary_heap/mod.rs
+++ b/library/alloc/src/collections/binary_heap/mod.rs
@@ -155,7 +155,7 @@ use crate::collections::TryReserveError;
 use crate::slice;
 use crate::vec::{self, AsVecIntoIter, Vec};
 
-use safety::{ensures, requires};
+use safety::requires;
 
 #[cfg(kani)]
 #[unstable(feature="kani", issue="none")]