From 57cecd4d185b0686c85f08dfe3a39cf5c4b1a978 Mon Sep 17 00:00:00 2001
From: Qinheping Hu <qinhh@amazon.com>
Date: Fri, 25 Oct 2024 01:10:29 -0500
Subject: [PATCH] Add loop_invariants predicate

---
 library/contracts/safety/src/kani.rs    | 16 +++++++++++++++-
 library/contracts/safety/src/lib.rs     |  6 ++++++
 library/contracts/safety/src/runtime.rs |  9 ++++++++-
 tool_config/kani-version.toml           |  2 +-
 4 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/library/contracts/safety/src/kani.rs b/library/contracts/safety/src/kani.rs
index a8b4ab360dc8e..82f7d9b86b346 100644
--- a/library/contracts/safety/src/kani.rs
+++ b/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 syn::{ItemFn, parse_macro_input, Stmt};
 
 pub(crate) fn requires(attr: TokenStream, item: TokenStream) -> TokenStream {
     rewrite_attr(attr, item, "requires")
@@ -10,6 +10,20 @@ pub(crate) fn ensures(attr: TokenStream, item: TokenStream) -> TokenStream {
     rewrite_attr(attr, item, "ensures")
 }
 
+pub(crate) fn loop_invariant(attr: TokenStream, stmt: TokenStream) -> TokenStream {
+    rewrite_stmt_attr(attr, stmt, "loop_invariant")
+}
+
+fn rewrite_stmt_attr(attr: TokenStream, stmt_stream: TokenStream, name: &str) -> TokenStream {
+    let args = proc_macro2::TokenStream::from(attr);
+    let stmt = parse_macro_input!(stmt_stream as Stmt);
+    let attribute = format_ident!("{}", name);
+    quote!(
+        #[cfg_attr(kani, kani_core::#attribute(#args))]
+        #stmt
+    ).into()
+}
+
 fn rewrite_attr(attr: TokenStream, item: TokenStream, name: &str) -> TokenStream {
     let args = proc_macro2::TokenStream::from(attr);
     let fn_item = parse_macro_input!(item as ItemFn);
diff --git a/library/contracts/safety/src/lib.rs b/library/contracts/safety/src/lib.rs
index 9fe852a622de3..feb572865f891 100644
--- a/library/contracts/safety/src/lib.rs
+++ b/library/contracts/safety/src/lib.rs
@@ -23,3 +23,9 @@ pub fn requires(attr: TokenStream, item: TokenStream) -> TokenStream {
 pub fn ensures(attr: TokenStream, item: TokenStream) -> TokenStream {
     tool::ensures(attr, item)
 }
+
+#[proc_macro_error]
+#[proc_macro_attribute]
+pub fn loop_invariant(attr: TokenStream, stmt_stream: TokenStream) -> TokenStream {
+    tool::loop_invariant(attr, stmt_stream)
+}
diff --git a/library/contracts/safety/src/runtime.rs b/library/contracts/safety/src/runtime.rs
index 78e8b1dc354d2..3d1cb1d96c90e 100644
--- a/library/contracts/safety/src/runtime.rs
+++ b/library/contracts/safety/src/runtime.rs
@@ -7,9 +7,16 @@ pub(crate) fn requires(_attr: TokenStream, item: TokenStream) -> TokenStream {
     item
 }
 
-/// For now, runtime requires is a no-op.
+/// For now, runtime ensures is a no-op.
 ///
 /// TODO: At runtime the `ensures` should become an assert as well.
 pub(crate) fn ensures(_attr: TokenStream, item: TokenStream) -> TokenStream {
     item
 }
+
+/// For now, runtime loop_invariant is a no-op.
+///
+/// TODO: At runtime the `loop_invariant` should become an assert as well.
+pub(crate) fn loop_invariant(_attr: TokenStream, stmt_stream: TokenStream) -> TokenStream {
+    stmt_stream
+}
diff --git a/tool_config/kani-version.toml b/tool_config/kani-version.toml
index 59b19da6a5958..5c0b4c857b6b0 100644
--- a/tool_config/kani-version.toml
+++ b/tool_config/kani-version.toml
@@ -1,5 +1,5 @@
 # This version should be updated whenever there is a change that makes this version of kani
-# incomaptible with the verify-std repo.
+# incompatible with the verify-std repo.
 
 [kani]
 commit = "5f8f513d297827cfdce4c48065e51973ba563068"