From 8a579daef419b00fe436844caf78476de00cc493 Mon Sep 17 00:00:00 2001
From: thanhnguyen-aws <ntson@amazon.com>
Date: Wed, 19 Mar 2025 12:04:48 -0700
Subject: [PATCH 01/18] add support or old in while loop inv

---
 .../src/sysroot/loop_contracts/mod.rs         | 198 +++++++++++++++++-
 1 file changed, 194 insertions(+), 4 deletions(-)

diff --git a/library/kani_macros/src/sysroot/loop_contracts/mod.rs b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
index 01166af3d988..738f3e8e29dc 100644
--- a/library/kani_macros/src/sysroot/loop_contracts/mod.rs
+++ b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
@@ -9,7 +9,144 @@ use proc_macro_error2::abort_call_site;
 use quote::{format_ident, quote};
 use syn::spanned::Spanned;
 use syn::token::AndAnd;
-use syn::{BinOp, Expr, ExprBinary, Stmt};
+use syn::{BinOp, Expr, ExprBinary, Stmt, Block, visit_mut::VisitMut, Ident};
+use syn::parse_macro_input;
+
+
+struct TransformationResult {
+    transformed_expr: Expr,
+    declarations_block: Block,
+    assignments_block: Block,
+}
+
+struct CallReplacer {
+    old_name: String,
+    replacements: Vec<(Expr, proc_macro2::Ident)>,
+    counter: usize,
+    var_prefix: String
+}
+
+impl CallReplacer {
+    fn new(old_name: &str, var_prefix: String) -> Self {
+        Self {
+            old_name: old_name.to_string(),
+            replacements: Vec::new(),
+            counter: 0,
+            var_prefix: var_prefix
+        }
+    }
+
+    fn generate_var_name(&mut self) -> proc_macro2::Ident {
+        let var_name = format_ident!("{}_{}", self.var_prefix, self.counter);
+        self.counter += 1;
+        var_name
+    }
+
+    fn should_replace(&self, expr_path: &syn::ExprPath) -> bool {
+        // Check both simple and qualified paths
+        if let Some(last_segment) = expr_path.path.segments.last() {
+            if last_segment.ident == self.old_name {
+                return true;
+            }
+        }
+        
+        let full_path = expr_path.path.segments
+            .iter()
+            .map(|seg| seg.ident.to_string())
+            .collect::<Vec<_>>()
+            .join("::");
+        
+        full_path.ends_with(&self.old_name)
+    }
+}
+
+impl VisitMut for CallReplacer {
+    fn visit_expr_mut(&mut self, expr: &mut Expr) {
+        // Visit nested expressions first
+        syn::visit_mut::visit_expr_mut(self, expr);
+        
+        if let Expr::Call(call) = expr {
+            if let Expr::Path(expr_path) = &*call.func {
+                if self.should_replace(expr_path) {
+                    let new_var = self.generate_var_name();
+                    self.replacements.push((expr.clone(), new_var.clone()));
+                    *expr = syn::parse_quote!(#new_var);
+                }
+            }
+        }
+    }
+}
+
+fn transform_function_calls(expr: Expr, function_name: &str, var_prefix: String) -> TransformationResult {
+    let mut replacer = CallReplacer::new(function_name, var_prefix);
+    let mut transformed_expr = expr;
+    replacer.visit_expr_mut(&mut transformed_expr);
+    
+    let mut newreplace : Vec<(Expr, Ident)> = Vec::new();
+    for (call, var) in replacer.replacements {
+        match call {
+            Expr::Call(call_expr) => {
+                let insideexpr = call_expr.args[0].clone();
+                newreplace.push((insideexpr, var.clone()));
+            }
+            _ => {}
+        }
+    }
+
+    // Generate declarations block
+    let declarations: Vec<Stmt> = newreplace
+        .iter()
+        .map(|(call, var)| {
+            syn::parse_quote!(let mut #var = #call;)
+        })
+        .collect();
+    let declarations_block: Block = syn::parse_quote!({
+        #(#declarations)*
+    });
+
+
+    // Generate assignments block
+    let assignments: Vec<Stmt> = newreplace
+        .into_iter()
+        .map(|(call, var)| {
+            syn::parse_quote!(#var = #call;)
+        })
+        .collect();
+    let assignments_block: Block = syn::parse_quote!({
+        #(#assignments)*
+    });
+
+    TransformationResult {
+        transformed_expr,
+        declarations_block,
+        assignments_block,
+    }
+}
+
+struct BreakContinueReplacer;
+
+impl VisitMut for BreakContinueReplacer {
+    fn visit_expr_mut(&mut self, expr: &mut Expr) {
+        // Visit nested expressions first
+        syn::visit_mut::visit_expr_mut(self, expr);
+        
+        // Replace the expression
+        *expr = match expr {
+            Expr::Break(_) => {
+                syn::parse_quote!(return)
+            }
+            Expr::Continue(_) => {
+                syn::parse_quote!(return)
+            }
+            _ => return
+        };
+    }
+}
+
+fn transform_break_continue(block: &mut Block) {
+    let mut replacer = BreakContinueReplacer;
+    replacer.visit_block_mut(block);
+}
 
 /// Expand loop contracts macros.
 ///
@@ -31,7 +168,7 @@ use syn::{BinOp, Expr, ExprBinary, Stmt};
 ///  }
 /// ```
 pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
-    // parse the stmt of the loop
+    // parse the stmt of the loop    
     let mut loop_stmt: Stmt = syn::parse(item.clone()).unwrap();
 
     // name of the loop invariant as closure of the form
@@ -41,7 +178,37 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
     inv_name.push_str(&loop_id);
 
     // expr of the loop invariant
-    let inv_expr: Expr = syn::parse(attr).unwrap();
+    let mut inv_expr: Expr = syn::parse(attr).unwrap();
+    // adding remember variables 
+    let mut var_prefix: String = "__kani_remember_var".to_owned();
+    var_prefix.push_str(&loop_id);
+    let transform_inv: TransformationResult = transform_function_calls(inv_expr.clone(), "old", var_prefix);
+    let has_old = !transform_inv.declarations_block.stmts.is_empty();
+    let decl_stms = transform_inv.declarations_block.stmts.clone();
+    let mut assign_stms = transform_inv.assignments_block.stmts.clone();
+    let (mut loop_body, loop_guard) = match loop_stmt {
+        Stmt::Expr(ref mut e, _) => match e {
+            Expr::While(ew) => (ew.body.clone(), ew.cond.clone()),
+            _ => panic!()
+        },
+        _ => panic!()
+    };
+    let loop_body_stms = loop_body.stmts.clone();
+    assign_stms.extend(loop_body_stms);
+    transform_break_continue(&mut loop_body);
+    let mut loop_body_closure_name: String = "__kani_loop_body_closure".to_owned();
+    loop_body_closure_name.push_str(&loop_id);
+    let loop_body_closure = format_ident!("{}", loop_body_closure_name);
+    if has_old {
+        inv_expr = transform_inv.transformed_expr.clone();
+        match loop_stmt {
+            Stmt::Expr(ref mut e, _) => match e {
+                Expr::While(ew) => ew.body.stmts = assign_stms.clone(),
+                _ => panic!()
+            },
+            _ => panic!()
+        };
+    }
 
     // ident of the register function
     let mut register_name: String = "kani_register_loop_contract".to_owned();
@@ -74,8 +241,15 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
             note = "for now, loop contracts is only supported for while-loops.";
         ),
     }
+    
+    if has_old {
     quote!(
         {
+        assert!(#loop_guard);
+        #(#decl_stms)*
+        let mut #loop_body_closure = ||
+        #loop_body;
+        #loop_body_closure ();
         // Dummy function used to force the compiler to capture the environment.
         // We cannot call closures inside constant functions.
         // This function gets replaced by `kani::internal::call_closure`.
@@ -84,8 +258,24 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
         const fn #register_ident<F: Fn() -> bool>(_f: &F, _transformed: usize) -> bool {
             true
         }
-        #loop_stmt})
+        #loop_stmt
+        })
     .into()
+    } else {
+    quote!(
+        {
+        // Dummy function used to force the compiler to capture the environment.
+        // We cannot call closures inside constant functions.
+        // This function gets replaced by `kani::internal::call_closure`.
+        #[inline(never)]
+        #[kanitool::fn_marker = "kani_register_loop_contract"]
+        const fn #register_ident<F: Fn() -> bool>(_f: &F, _transformed: usize) -> bool {
+            true
+        }
+        #loop_stmt
+        })
+    .into()
+    }
 }
 
 fn generate_unique_id_from_span(stmt: &Stmt) -> String {

From 7520bc3b68ecb178341283cee6a6b34d0d27a028 Mon Sep 17 00:00:00 2001
From: thanhnguyen-aws <ntson@amazon.com>
Date: Wed, 19 Mar 2025 13:55:43 -0700
Subject: [PATCH 02/18] add expected

---
 .../src/sysroot/loop_contracts/mod.rs         | 81 +++++++++----------
 1 file changed, 37 insertions(+), 44 deletions(-)

diff --git a/library/kani_macros/src/sysroot/loop_contracts/mod.rs b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
index 738f3e8e29dc..7da71739f63b 100644
--- a/library/kani_macros/src/sysroot/loop_contracts/mod.rs
+++ b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
@@ -7,11 +7,10 @@
 use proc_macro::TokenStream;
 use proc_macro_error2::abort_call_site;
 use quote::{format_ident, quote};
+use syn::parse_macro_input;
 use syn::spanned::Spanned;
 use syn::token::AndAnd;
-use syn::{BinOp, Expr, ExprBinary, Stmt, Block, visit_mut::VisitMut, Ident};
-use syn::parse_macro_input;
-
+use syn::{BinOp, Block, Expr, ExprBinary, Ident, Stmt, visit_mut::VisitMut};
 
 struct TransformationResult {
     transformed_expr: Expr,
@@ -23,7 +22,7 @@ struct CallReplacer {
     old_name: String,
     replacements: Vec<(Expr, proc_macro2::Ident)>,
     counter: usize,
-    var_prefix: String
+    var_prefix: String,
 }
 
 impl CallReplacer {
@@ -32,7 +31,7 @@ impl CallReplacer {
             old_name: old_name.to_string(),
             replacements: Vec::new(),
             counter: 0,
-            var_prefix: var_prefix
+            var_prefix: var_prefix,
         }
     }
 
@@ -49,13 +48,15 @@ impl CallReplacer {
                 return true;
             }
         }
-        
-        let full_path = expr_path.path.segments
+
+        let full_path = expr_path
+            .path
+            .segments
             .iter()
             .map(|seg| seg.ident.to_string())
             .collect::<Vec<_>>()
             .join("::");
-        
+
         full_path.ends_with(&self.old_name)
     }
 }
@@ -64,7 +65,7 @@ impl VisitMut for CallReplacer {
     fn visit_expr_mut(&mut self, expr: &mut Expr) {
         // Visit nested expressions first
         syn::visit_mut::visit_expr_mut(self, expr);
-        
+
         if let Expr::Call(call) = expr {
             if let Expr::Path(expr_path) = &*call.func {
                 if self.should_replace(expr_path) {
@@ -77,12 +78,16 @@ impl VisitMut for CallReplacer {
     }
 }
 
-fn transform_function_calls(expr: Expr, function_name: &str, var_prefix: String) -> TransformationResult {
+fn transform_function_calls(
+    expr: Expr,
+    function_name: &str,
+    var_prefix: String,
+) -> TransformationResult {
     let mut replacer = CallReplacer::new(function_name, var_prefix);
     let mut transformed_expr = expr;
     replacer.visit_expr_mut(&mut transformed_expr);
-    
-    let mut newreplace : Vec<(Expr, Ident)> = Vec::new();
+
+    let mut newreplace: Vec<(Expr, Ident)> = Vec::new();
     for (call, var) in replacer.replacements {
         match call {
             Expr::Call(call_expr) => {
@@ -94,33 +99,20 @@ fn transform_function_calls(expr: Expr, function_name: &str, var_prefix: String)
     }
 
     // Generate declarations block
-    let declarations: Vec<Stmt> = newreplace
-        .iter()
-        .map(|(call, var)| {
-            syn::parse_quote!(let mut #var = #call;)
-        })
-        .collect();
+    let declarations: Vec<Stmt> =
+        newreplace.iter().map(|(call, var)| syn::parse_quote!(let mut #var = #call;)).collect();
     let declarations_block: Block = syn::parse_quote!({
         #(#declarations)*
     });
 
-
     // Generate assignments block
-    let assignments: Vec<Stmt> = newreplace
-        .into_iter()
-        .map(|(call, var)| {
-            syn::parse_quote!(#var = #call;)
-        })
-        .collect();
+    let assignments: Vec<Stmt> =
+        newreplace.into_iter().map(|(call, var)| syn::parse_quote!(#var = #call;)).collect();
     let assignments_block: Block = syn::parse_quote!({
         #(#assignments)*
     });
 
-    TransformationResult {
-        transformed_expr,
-        declarations_block,
-        assignments_block,
-    }
+    TransformationResult { transformed_expr, declarations_block, assignments_block }
 }
 
 struct BreakContinueReplacer;
@@ -129,7 +121,7 @@ impl VisitMut for BreakContinueReplacer {
     fn visit_expr_mut(&mut self, expr: &mut Expr) {
         // Visit nested expressions first
         syn::visit_mut::visit_expr_mut(self, expr);
-        
+
         // Replace the expression
         *expr = match expr {
             Expr::Break(_) => {
@@ -138,7 +130,7 @@ impl VisitMut for BreakContinueReplacer {
             Expr::Continue(_) => {
                 syn::parse_quote!(return)
             }
-            _ => return
+            _ => return,
         };
     }
 }
@@ -168,7 +160,7 @@ fn transform_break_continue(block: &mut Block) {
 ///  }
 /// ```
 pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
-    // parse the stmt of the loop    
+    // parse the stmt of the loop
     let mut loop_stmt: Stmt = syn::parse(item.clone()).unwrap();
 
     // name of the loop invariant as closure of the form
@@ -179,19 +171,20 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
 
     // expr of the loop invariant
     let mut inv_expr: Expr = syn::parse(attr).unwrap();
-    // adding remember variables 
+    // adding remember variables
     let mut var_prefix: String = "__kani_remember_var".to_owned();
     var_prefix.push_str(&loop_id);
-    let transform_inv: TransformationResult = transform_function_calls(inv_expr.clone(), "old", var_prefix);
+    let transform_inv: TransformationResult =
+        transform_function_calls(inv_expr.clone(), "old", var_prefix);
     let has_old = !transform_inv.declarations_block.stmts.is_empty();
     let decl_stms = transform_inv.declarations_block.stmts.clone();
     let mut assign_stms = transform_inv.assignments_block.stmts.clone();
     let (mut loop_body, loop_guard) = match loop_stmt {
         Stmt::Expr(ref mut e, _) => match e {
             Expr::While(ew) => (ew.body.clone(), ew.cond.clone()),
-            _ => panic!()
+            _ => panic!(),
         },
-        _ => panic!()
+        _ => panic!(),
     };
     let loop_body_stms = loop_body.stmts.clone();
     assign_stms.extend(loop_body_stms);
@@ -204,9 +197,9 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
         match loop_stmt {
             Stmt::Expr(ref mut e, _) => match e {
                 Expr::While(ew) => ew.body.stmts = assign_stms.clone(),
-                _ => panic!()
+                _ => panic!(),
             },
-            _ => panic!()
+            _ => panic!(),
         };
     }
 
@@ -241,9 +234,9 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
             note = "for now, loop contracts is only supported for while-loops.";
         ),
     }
-    
+
     if has_old {
-    quote!(
+        quote!(
         {
         assert!(#loop_guard);
         #(#decl_stms)*
@@ -260,9 +253,9 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
         }
         #loop_stmt
         })
-    .into()
+        .into()
     } else {
-    quote!(
+        quote!(
         {
         // Dummy function used to force the compiler to capture the environment.
         // We cannot call closures inside constant functions.
@@ -274,7 +267,7 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
         }
         #loop_stmt
         })
-    .into()
+        .into()
     }
 }
 

From c3934394d055f747faa22a1d007afd129e22f85a Mon Sep 17 00:00:00 2001
From: thanhnguyen-aws <ntson@amazon.com>
Date: Wed, 19 Mar 2025 14:00:25 -0700
Subject: [PATCH 03/18] add expected

---
 library/kani_macros/src/sysroot/loop_contracts/mod.rs | 1 -
 1 file changed, 1 deletion(-)

diff --git a/library/kani_macros/src/sysroot/loop_contracts/mod.rs b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
index 7da71739f63b..9235493e4eaf 100644
--- a/library/kani_macros/src/sysroot/loop_contracts/mod.rs
+++ b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
@@ -7,7 +7,6 @@
 use proc_macro::TokenStream;
 use proc_macro_error2::abort_call_site;
 use quote::{format_ident, quote};
-use syn::parse_macro_input;
 use syn::spanned::Spanned;
 use syn::token::AndAnd;
 use syn::{BinOp, Block, Expr, ExprBinary, Ident, Stmt, visit_mut::VisitMut};

From 87149319440553711f87282efa58b11c626b88c7 Mon Sep 17 00:00:00 2001
From: thanhnguyen-aws <ntson@amazon.com>
Date: Wed, 19 Mar 2025 14:00:37 -0700
Subject: [PATCH 04/18] add expected

---
 .../loop-contract/loop_with_old.expected      |  1 +
 tests/expected/loop-contract/loop_with_old.rs | 24 +++++++++++++++++++
 2 files changed, 25 insertions(+)
 create mode 100644 tests/expected/loop-contract/loop_with_old.expected
 create mode 100644 tests/expected/loop-contract/loop_with_old.rs

diff --git a/tests/expected/loop-contract/loop_with_old.expected b/tests/expected/loop-contract/loop_with_old.expected
new file mode 100644
index 000000000000..34c886c358cb
--- /dev/null
+++ b/tests/expected/loop-contract/loop_with_old.expected
@@ -0,0 +1 @@
+VERIFICATION:- SUCCESSFUL
diff --git a/tests/expected/loop-contract/loop_with_old.rs b/tests/expected/loop-contract/loop_with_old.rs
new file mode 100644
index 000000000000..2db067ee2daf
--- /dev/null
+++ b/tests/expected/loop-contract/loop_with_old.rs
@@ -0,0 +1,24 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+// kani-flags: -Z loop-contracts
+
+//! Check if loop contracts is correctly applied.
+
+#![feature(stmt_expr_attributes)]
+#![feature(proc_macro_hygiene)]
+
+#[kani::proof]
+pub fn loop_with_old() {
+    let mut i = 100;
+    let mut j = 100;
+    #[kani::loop_invariant((i >= 2) && (i <= 100) && (i % 2 == 0) && (j == 2*i-100) && (old(i) == i + 2) && (old(j) == j + 4) && (old(i-j) == i-j-2) )]
+    while i > 2 {
+        if i == 1 {
+            break;
+        }
+        i = i - 2;
+        j = j - 4
+    }
+    assert!(i == 2);
+}

From bb7fdff02269be048e8fa86bb9e5fc263cd7af0e Mon Sep 17 00:00:00 2001
From: thanhnguyen-aws <ntson@amazon.com>
Date: Wed, 19 Mar 2025 17:08:57 -0700
Subject: [PATCH 05/18] add more expected tests

---
 .../src/sysroot/loop_contracts/mod.rs         | 26 ++++++++++++++-----
 tests/expected/loop-contract/loop_with_old.rs | 18 ++++++-------
 .../loop_with_old_and_prev.expected           |  1 +
 .../loop-contract/loop_with_old_and_prev.rs   | 22 ++++++++++++++++
 .../expected/loop-contract/loop_with_prev.rs  | 24 +++++++++++++++++
 .../loop-contract/loop_with_previous.expected |  1 +
 6 files changed, 76 insertions(+), 16 deletions(-)
 create mode 100644 tests/expected/loop-contract/loop_with_old_and_prev.expected
 create mode 100644 tests/expected/loop-contract/loop_with_old_and_prev.rs
 create mode 100644 tests/expected/loop-contract/loop_with_prev.rs
 create mode 100644 tests/expected/loop-contract/loop_with_previous.expected

diff --git a/library/kani_macros/src/sysroot/loop_contracts/mod.rs b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
index 9235493e4eaf..9a2256706242 100644
--- a/library/kani_macros/src/sysroot/loop_contracts/mod.rs
+++ b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
@@ -170,12 +170,21 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
 
     // expr of the loop invariant
     let mut inv_expr: Expr = syn::parse(attr).unwrap();
-    // adding remember variables
-    let mut var_prefix: String = "__kani_remember_var".to_owned();
-    var_prefix.push_str(&loop_id);
+
+    // adding old variables
+    let mut old_var_prefix: String = "__kani_old_var".to_owned();
+    old_var_prefix.push_str(&loop_id);
+    let replace_old: TransformationResult =
+        transform_function_calls(inv_expr.clone(), "old", old_var_prefix);
+    inv_expr = replace_old.transformed_expr.clone();
+    let old_decl_stms = replace_old.declarations_block.stmts.clone();
+
+    // adding prev variables
+    let mut prev_var_prefix: String = "__kani_prev_var".to_owned();
+    prev_var_prefix.push_str(&loop_id);
     let transform_inv: TransformationResult =
-        transform_function_calls(inv_expr.clone(), "old", var_prefix);
-    let has_old = !transform_inv.declarations_block.stmts.is_empty();
+        transform_function_calls(inv_expr.clone(), "prev", prev_var_prefix);
+    let has_prev = !transform_inv.declarations_block.stmts.is_empty();
     let decl_stms = transform_inv.declarations_block.stmts.clone();
     let mut assign_stms = transform_inv.assignments_block.stmts.clone();
     let (mut loop_body, loop_guard) = match loop_stmt {
@@ -191,7 +200,7 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
     let mut loop_body_closure_name: String = "__kani_loop_body_closure".to_owned();
     loop_body_closure_name.push_str(&loop_id);
     let loop_body_closure = format_ident!("{}", loop_body_closure_name);
-    if has_old {
+    if has_prev {
         inv_expr = transform_inv.transformed_expr.clone();
         match loop_stmt {
             Stmt::Expr(ref mut e, _) => match e {
@@ -234,14 +243,16 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
         ),
     }
 
-    if has_old {
+    if has_prev {
         quote!(
         {
+        #(#old_decl_stms)*
         assert!(#loop_guard);
         #(#decl_stms)*
         let mut #loop_body_closure = ||
         #loop_body;
         #loop_body_closure ();
+        assert!(#inv_expr);
         // Dummy function used to force the compiler to capture the environment.
         // We cannot call closures inside constant functions.
         // This function gets replaced by `kani::internal::call_closure`.
@@ -256,6 +267,7 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
     } else {
         quote!(
         {
+        #(#old_decl_stms)*
         // Dummy function used to force the compiler to capture the environment.
         // We cannot call closures inside constant functions.
         // This function gets replaced by `kani::internal::call_closure`.
diff --git a/tests/expected/loop-contract/loop_with_old.rs b/tests/expected/loop-contract/loop_with_old.rs
index 2db067ee2daf..0d83b050cfcd 100644
--- a/tests/expected/loop-contract/loop_with_old.rs
+++ b/tests/expected/loop-contract/loop_with_old.rs
@@ -10,15 +10,15 @@
 
 #[kani::proof]
 pub fn loop_with_old() {
-    let mut i = 100;
-    let mut j = 100;
-    #[kani::loop_invariant((i >= 2) && (i <= 100) && (i % 2 == 0) && (j == 2*i-100) && (old(i) == i + 2) && (old(j) == j + 4) && (old(i-j) == i-j-2) )]
-    while i > 2 {
-        if i == 1 {
-            break;
+    let mut x: u8 = kani::any_where(|v| *v < 10);
+    let mut y: u8 = kani::any();
+    let mut i = 0;
+    #[kani::loop_invariant( (i<=5) && (x <= old(x) + i) && (old(x) + i == old(i) + x))]
+    while i < 5 {
+        if i == 0 {
+            y = x
         }
-        i = i - 2;
-        j = j - 4
+        x += 1;
+        i += 1;
     }
-    assert!(i == 2);
 }
diff --git a/tests/expected/loop-contract/loop_with_old_and_prev.expected b/tests/expected/loop-contract/loop_with_old_and_prev.expected
new file mode 100644
index 000000000000..34c886c358cb
--- /dev/null
+++ b/tests/expected/loop-contract/loop_with_old_and_prev.expected
@@ -0,0 +1 @@
+VERIFICATION:- SUCCESSFUL
diff --git a/tests/expected/loop-contract/loop_with_old_and_prev.rs b/tests/expected/loop-contract/loop_with_old_and_prev.rs
new file mode 100644
index 000000000000..851c27185e88
--- /dev/null
+++ b/tests/expected/loop-contract/loop_with_old_and_prev.rs
@@ -0,0 +1,22 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+// kani-flags: -Z loop-contracts
+
+//! Check if loop contracts is correctly applied.
+
+#![feature(stmt_expr_attributes)]
+#![feature(proc_macro_hygiene)]
+
+#[kani::proof]
+pub fn loop_with_old_and_prev() {
+    let mut i = 100;
+    #[kani::loop_invariant((i >= 2) && (i <= 100) && (i % 2 == 0) && (old(i) == 100) && (prev(i) == i + 2))]
+    while i > 2 {
+        if i == 1 {
+            break;
+        }
+        i = i - 2;
+    }
+    assert!(i == 2);
+}
diff --git a/tests/expected/loop-contract/loop_with_prev.rs b/tests/expected/loop-contract/loop_with_prev.rs
new file mode 100644
index 000000000000..8a5dfb36e230
--- /dev/null
+++ b/tests/expected/loop-contract/loop_with_prev.rs
@@ -0,0 +1,24 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+// kani-flags: -Z loop-contracts
+
+//! Check if loop contracts is correctly applied.
+
+#![feature(stmt_expr_attributes)]
+#![feature(proc_macro_hygiene)]
+
+#[kani::proof]
+pub fn loop_with_prev() {
+    let mut i = 100;
+    let mut j = 100;
+    #[kani::loop_invariant((i >= 2) && (i <= 100) && (i % 2 == 0) && (j == 2*i-100) && (prev(i) == i + 2) && (prev(j) == j + 4) && (prev(i-j) == i-j-2) )]
+    while i > 2 {
+        if i == 1 {
+            break;
+        }
+        i = i - 2;
+        j = j - 4
+    }
+    assert!(i == 2);
+}
diff --git a/tests/expected/loop-contract/loop_with_previous.expected b/tests/expected/loop-contract/loop_with_previous.expected
new file mode 100644
index 000000000000..34c886c358cb
--- /dev/null
+++ b/tests/expected/loop-contract/loop_with_previous.expected
@@ -0,0 +1 @@
+VERIFICATION:- SUCCESSFUL

From 1c168bd53eabad926f922bf9a067b6f7a4404877 Mon Sep 17 00:00:00 2001
From: thanhnguyen-aws <ntson@amazon.com>
Date: Thu, 20 Mar 2025 08:44:54 -0700
Subject: [PATCH 06/18] add clone to assignment, fix expected file name

---
 library/kani_macros/src/sysroot/loop_contracts/mod.rs         | 4 ++--
 .../{loop_with_previous.expected => loop_with_prev.expected}  | 0
 2 files changed, 2 insertions(+), 2 deletions(-)
 rename tests/expected/loop-contract/{loop_with_previous.expected => loop_with_prev.expected} (100%)

diff --git a/library/kani_macros/src/sysroot/loop_contracts/mod.rs b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
index 9a2256706242..46d3cabfd6b1 100644
--- a/library/kani_macros/src/sysroot/loop_contracts/mod.rs
+++ b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
@@ -99,14 +99,14 @@ fn transform_function_calls(
 
     // Generate declarations block
     let declarations: Vec<Stmt> =
-        newreplace.iter().map(|(call, var)| syn::parse_quote!(let mut #var = #call;)).collect();
+        newreplace.iter().map(|(call, var)| syn::parse_quote!(let mut #var = #call.clone();)).collect();
     let declarations_block: Block = syn::parse_quote!({
         #(#declarations)*
     });
 
     // Generate assignments block
     let assignments: Vec<Stmt> =
-        newreplace.into_iter().map(|(call, var)| syn::parse_quote!(#var = #call;)).collect();
+        newreplace.into_iter().map(|(call, var)| syn::parse_quote!(#var = #call.clone();)).collect();
     let assignments_block: Block = syn::parse_quote!({
         #(#assignments)*
     });
diff --git a/tests/expected/loop-contract/loop_with_previous.expected b/tests/expected/loop-contract/loop_with_prev.expected
similarity index 100%
rename from tests/expected/loop-contract/loop_with_previous.expected
rename to tests/expected/loop-contract/loop_with_prev.expected

From b981fa219af17ec0394478101670419ccda17302 Mon Sep 17 00:00:00 2001
From: thanhnguyen-aws <ntson@amazon.com>
Date: Fri, 21 Mar 2025 08:28:05 -0700
Subject: [PATCH 07/18] fix expected files and add commnents in tests

---
 tests/expected/loop-contract/loop_with_old.expected         | 5 +++++
 tests/expected/loop-contract/loop_with_old.rs               | 2 +-
 .../expected/loop-contract/loop_with_old_and_prev.expected  | 6 ++++++
 tests/expected/loop-contract/loop_with_old_and_prev.rs      | 2 +-
 tests/expected/loop-contract/loop_with_prev.expected        | 5 +++++
 tests/expected/loop-contract/loop_with_prev.rs              | 2 +-
 6 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/tests/expected/loop-contract/loop_with_old.expected b/tests/expected/loop-contract/loop_with_old.expected
index 34c886c358cb..ec9403766c7e 100644
--- a/tests/expected/loop-contract/loop_with_old.expected
+++ b/tests/expected/loop-contract/loop_with_old.expected
@@ -1 +1,6 @@
+Check 10: loop_with_old.loop_invariant_base.1
+	 - Status: SUCCESS
+	 - Description: "Check invariant before entry for loop loop_with_old.0"
+	 - Location: src/loop_with_old.rs:16:5 in function loop_with_old
+
 VERIFICATION:- SUCCESSFUL
diff --git a/tests/expected/loop-contract/loop_with_old.rs b/tests/expected/loop-contract/loop_with_old.rs
index 0d83b050cfcd..e5340a38e53e 100644
--- a/tests/expected/loop-contract/loop_with_old.rs
+++ b/tests/expected/loop-contract/loop_with_old.rs
@@ -3,7 +3,7 @@
 
 // kani-flags: -Z loop-contracts
 
-//! Check if loop contracts is correctly applied.
+//! Check the use of "old" in loop invariant
 
 #![feature(stmt_expr_attributes)]
 #![feature(proc_macro_hygiene)]
diff --git a/tests/expected/loop-contract/loop_with_old_and_prev.expected b/tests/expected/loop-contract/loop_with_old_and_prev.expected
index 34c886c358cb..4f3b992aa3a6 100644
--- a/tests/expected/loop-contract/loop_with_old_and_prev.expected
+++ b/tests/expected/loop-contract/loop_with_old_and_prev.expected
@@ -1 +1,7 @@
+Check 24: loop_with_old_and_prev.loop_invariant_base.1
+	 - Status: SUCCESS
+	 - Description: "Check invariant before entry for loop loop_with_old_and_prev.0"
+	 - Location: src/loop_with_old_and_prev.rs:14:5 in function loop_with_old_and_prev
+
+
 VERIFICATION:- SUCCESSFUL
diff --git a/tests/expected/loop-contract/loop_with_old_and_prev.rs b/tests/expected/loop-contract/loop_with_old_and_prev.rs
index 851c27185e88..ca3a8896005e 100644
--- a/tests/expected/loop-contract/loop_with_old_and_prev.rs
+++ b/tests/expected/loop-contract/loop_with_old_and_prev.rs
@@ -3,7 +3,7 @@
 
 // kani-flags: -Z loop-contracts
 
-//! Check if loop contracts is correctly applied.
+//! Check the use of both "old" and "prev" in loop invariant
 
 #![feature(stmt_expr_attributes)]
 #![feature(proc_macro_hygiene)]
diff --git a/tests/expected/loop-contract/loop_with_prev.expected b/tests/expected/loop-contract/loop_with_prev.expected
index 34c886c358cb..20bbff88d6b9 100644
--- a/tests/expected/loop-contract/loop_with_prev.expected
+++ b/tests/expected/loop-contract/loop_with_prev.expected
@@ -1 +1,6 @@
+Check 33: loop_with_prev.loop_invariant_step.1
+	 - Status: SUCCESS
+	 - Description: "Check invariant after step for loop loop_with_prev.0"
+	 - Location: src/loop_with_prev.rs:15:5 in function loop_with_prev
+
 VERIFICATION:- SUCCESSFUL
diff --git a/tests/expected/loop-contract/loop_with_prev.rs b/tests/expected/loop-contract/loop_with_prev.rs
index 8a5dfb36e230..d64e766a4abf 100644
--- a/tests/expected/loop-contract/loop_with_prev.rs
+++ b/tests/expected/loop-contract/loop_with_prev.rs
@@ -3,7 +3,7 @@
 
 // kani-flags: -Z loop-contracts
 
-//! Check if loop contracts is correctly applied.
+//! Check the use of "prev" in loop invariant
 
 #![feature(stmt_expr_attributes)]
 #![feature(proc_macro_hygiene)]

From 5fc1da162ccbd3025a4258225a5e371103aa9cbe Mon Sep 17 00:00:00 2001
From: thanhnguyen-aws <ntson@amazon.com>
Date: Mon, 24 Mar 2025 10:18:25 -0700
Subject: [PATCH 08/18] fix expected files

---
 .../kani_macros/src/sysroot/loop_contracts/mod.rs  | 14 +++++++++-----
 .../expected/loop-contract/loop_with_old.expected  |  5 ++---
 tests/expected/loop-contract/loop_with_old.rs      |  2 +-
 .../loop-contract/loop_with_old_and_prev.expected  |  7 +++----
 .../loop-contract/loop_with_old_and_prev.rs        |  2 +-
 .../expected/loop-contract/loop_with_prev.expected |  5 ++---
 6 files changed, 18 insertions(+), 17 deletions(-)

diff --git a/library/kani_macros/src/sysroot/loop_contracts/mod.rs b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
index 46d3cabfd6b1..7203fb948fb3 100644
--- a/library/kani_macros/src/sysroot/loop_contracts/mod.rs
+++ b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
@@ -98,15 +98,19 @@ fn transform_function_calls(
     }
 
     // Generate declarations block
-    let declarations: Vec<Stmt> =
-        newreplace.iter().map(|(call, var)| syn::parse_quote!(let mut #var = #call.clone();)).collect();
+    let declarations: Vec<Stmt> = newreplace
+        .iter()
+        .map(|(call, var)| syn::parse_quote!(let mut #var = #call.clone();))
+        .collect();
     let declarations_block: Block = syn::parse_quote!({
         #(#declarations)*
     });
 
     // Generate assignments block
-    let assignments: Vec<Stmt> =
-        newreplace.into_iter().map(|(call, var)| syn::parse_quote!(#var = #call.clone();)).collect();
+    let assignments: Vec<Stmt> = newreplace
+        .into_iter()
+        .map(|(call, var)| syn::parse_quote!(#var = #call.clone();))
+        .collect();
     let assignments_block: Block = syn::parse_quote!({
         #(#assignments)*
     });
@@ -175,7 +179,7 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
     let mut old_var_prefix: String = "__kani_old_var".to_owned();
     old_var_prefix.push_str(&loop_id);
     let replace_old: TransformationResult =
-        transform_function_calls(inv_expr.clone(), "old", old_var_prefix);
+        transform_function_calls(inv_expr.clone(), "on_entry", old_var_prefix);
     inv_expr = replace_old.transformed_expr.clone();
     let old_decl_stms = replace_old.declarations_block.stmts.clone();
 
diff --git a/tests/expected/loop-contract/loop_with_old.expected b/tests/expected/loop-contract/loop_with_old.expected
index ec9403766c7e..33035940db68 100644
--- a/tests/expected/loop-contract/loop_with_old.expected
+++ b/tests/expected/loop-contract/loop_with_old.expected
@@ -1,6 +1,5 @@
-Check 10: loop_with_old.loop_invariant_base.1
-	 - Status: SUCCESS
+loop_with_old.loop_invariant_base.1\
+	 - Status: SUCCESS\
 	 - Description: "Check invariant before entry for loop loop_with_old.0"
-	 - Location: src/loop_with_old.rs:16:5 in function loop_with_old
 
 VERIFICATION:- SUCCESSFUL
diff --git a/tests/expected/loop-contract/loop_with_old.rs b/tests/expected/loop-contract/loop_with_old.rs
index e5340a38e53e..32371b926831 100644
--- a/tests/expected/loop-contract/loop_with_old.rs
+++ b/tests/expected/loop-contract/loop_with_old.rs
@@ -13,7 +13,7 @@ pub fn loop_with_old() {
     let mut x: u8 = kani::any_where(|v| *v < 10);
     let mut y: u8 = kani::any();
     let mut i = 0;
-    #[kani::loop_invariant( (i<=5) && (x <= old(x) + i) && (old(x) + i == old(i) + x))]
+    #[kani::loop_invariant( (i<=5) && (x <= on_entry(x) + i) && (on_entry(x) + i == on_entry(i) + x))]
     while i < 5 {
         if i == 0 {
             y = x
diff --git a/tests/expected/loop-contract/loop_with_old_and_prev.expected b/tests/expected/loop-contract/loop_with_old_and_prev.expected
index 4f3b992aa3a6..3d74c06488b8 100644
--- a/tests/expected/loop-contract/loop_with_old_and_prev.expected
+++ b/tests/expected/loop-contract/loop_with_old_and_prev.expected
@@ -1,7 +1,6 @@
-Check 24: loop_with_old_and_prev.loop_invariant_base.1
-	 - Status: SUCCESS
-	 - Description: "Check invariant before entry for loop loop_with_old_and_prev.0"
-	 - Location: src/loop_with_old_and_prev.rs:14:5 in function loop_with_old_and_prev
+loop_with_old_and_prev.loop_invariant_base.1\
+	 - Status: SUCCESS\
+	 - Description: "Check invariant before entry for loop loop_with_old_and_prev.0"\
 
 
 VERIFICATION:- SUCCESSFUL
diff --git a/tests/expected/loop-contract/loop_with_old_and_prev.rs b/tests/expected/loop-contract/loop_with_old_and_prev.rs
index ca3a8896005e..87273c8d79df 100644
--- a/tests/expected/loop-contract/loop_with_old_and_prev.rs
+++ b/tests/expected/loop-contract/loop_with_old_and_prev.rs
@@ -11,7 +11,7 @@
 #[kani::proof]
 pub fn loop_with_old_and_prev() {
     let mut i = 100;
-    #[kani::loop_invariant((i >= 2) && (i <= 100) && (i % 2 == 0) && (old(i) == 100) && (prev(i) == i + 2))]
+    #[kani::loop_invariant((i >= 2) && (i <= 100) && (i % 2 == 0) && (on_entry(i) == 100) && (prev(i) == i + 2))]
     while i > 2 {
         if i == 1 {
             break;
diff --git a/tests/expected/loop-contract/loop_with_prev.expected b/tests/expected/loop-contract/loop_with_prev.expected
index 20bbff88d6b9..1ff35b1c3e37 100644
--- a/tests/expected/loop-contract/loop_with_prev.expected
+++ b/tests/expected/loop-contract/loop_with_prev.expected
@@ -1,6 +1,5 @@
-Check 33: loop_with_prev.loop_invariant_step.1
-	 - Status: SUCCESS
+loop_with_prev.loop_invariant_step.1\
+	 - Status: SUCCESS\
 	 - Description: "Check invariant after step for loop loop_with_prev.0"
-	 - Location: src/loop_with_prev.rs:15:5 in function loop_with_prev
 
 VERIFICATION:- SUCCESSFUL

From 6a30e72acbb2d28cce3521eec97bec0234df98fb Mon Sep 17 00:00:00 2001
From: thanhnguyen-aws <ntson@amazon.com>
Date: Mon, 24 Mar 2025 14:16:00 -0700
Subject: [PATCH 09/18] added comments to expain the code

---
 .../src/sysroot/loop_contracts/mod.rs         | 110 +++++++++++++-----
 1 file changed, 81 insertions(+), 29 deletions(-)

diff --git a/library/kani_macros/src/sysroot/loop_contracts/mod.rs b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
index 7203fb948fb3..e3ae51e853e3 100644
--- a/library/kani_macros/src/sysroot/loop_contracts/mod.rs
+++ b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
@@ -11,6 +11,80 @@ use syn::spanned::Spanned;
 use syn::token::AndAnd;
 use syn::{BinOp, Block, Expr, ExprBinary, Ident, Stmt, visit_mut::VisitMut};
 
+/*
+Transform the loop to support on_entry(expr) : the value of expr before entering the loop
+1. For each on_entry(expr) in the loop variant, replace it with a newly generated "memory" variable old_k
+2. Add the declaration of i before the loop: let old_k = expr
+For example:
+#[kani::loop_invariant(on_entry(x+y) = x + y -1)]
+while(....)
+
+is transformed into
+let old_1 = x + y
+#[kani::loop_invariant(old_1 = x + y -1)]
+while(....)
+
+Then the loop_invartiant is transformed
+
+*/
+
+/*
+Transform the loop to support prev(expr) : the value of expr at the end of the previous iteration
+If there is a prev(expr) in the loop_invariant:
+1. Firstly, Add an assertion of the loop_quard: the loop must run at least once, otherwise prev(expr) is undefined.
+1. For each prev(expr) in the loop variant, replace it with a newly generated "memory" variable prev_k
+2. Add the declaration of i before the loop: let mut prev_k = expr
+3. Define a mut closure whose body is exactly the loop body, but replace all break and continue statements with return statements
+4. Call the closure (the same as run the loop once)
+5. Add the assertion for the loop_invariant (not includes the loop_quard): check if the loop_invariant holds after the first iteration.
+6. Add the loop
+6. Add the assignment statements (exactly the same as the declarations without the "let mut") on the top of the loop body to update the "memory" variables
+For example:
+#[kani::loop_invariant(prev(x+y) = x + y -1 && ...)]
+while(loop_guard)
+{
+    loop_body
+}
+
+is transformed into
+
+assert!(loop_guard);
+let mut prev_1 = x + y;
+let mut loop_body_closure = || {
+    loop_body_replaced //replace breaks/continues in loop_body with returns
+};
+loop_body_closure();
+assert!(prev_1 = x + y -1 && ...);
+#[kani::loop_invariant(prev_1  = x + y -1)]
+while(loop_guard)
+{
+    prev_1 = x + y;
+    loop_body
+}
+
+*/
+
+/// After that:
+/// Expand loop contracts macros.
+///
+/// A while loop of the form
+/// ``` rust
+///  while guard {
+///      body
+///  }
+/// ```
+/// will be annotated as
+/// ``` rust
+/// #[inline(never)]
+/// #[kanitool::fn_marker = "kani_register_loop_contract"]
+/// const fn kani_register_loop_contract_id<T, F: FnOnce() -> T>(f: F) -> T {
+///     unreachable!()
+/// }
+///  while kani_register_loop_contract_id(|| -> bool {inv};) && guard {
+///      body
+///  }
+/// ```
+
 struct TransformationResult {
     transformed_expr: Expr,
     declarations_block: Block,
@@ -24,6 +98,7 @@ struct CallReplacer {
     var_prefix: String,
 }
 
+// This impl replaces any function call of a function name : old_name with a newly generated variable.
 impl CallReplacer {
     fn new(old_name: &str, var_prefix: String) -> Self {
         Self {
@@ -40,14 +115,8 @@ impl CallReplacer {
         var_name
     }
 
+    //Check if the function name is old_name
     fn should_replace(&self, expr_path: &syn::ExprPath) -> bool {
-        // Check both simple and qualified paths
-        if let Some(last_segment) = expr_path.path.segments.last() {
-            if last_segment.ident == self.old_name {
-                return true;
-            }
-        }
-
         let full_path = expr_path
             .path
             .segments
@@ -56,7 +125,7 @@ impl CallReplacer {
             .collect::<Vec<_>>()
             .join("::");
 
-        full_path.ends_with(&self.old_name)
+        full_path == self.old_name
     }
 }
 
@@ -77,6 +146,7 @@ impl VisitMut for CallReplacer {
     }
 }
 
+// The main function to replace the function call with the variables
 fn transform_function_calls(
     expr: Expr,
     function_name: &str,
@@ -97,7 +167,7 @@ fn transform_function_calls(
         }
     }
 
-    // Generate declarations block
+    // Generate declarations block of the newly generated variables (will added before the loop)
     let declarations: Vec<Stmt> = newreplace
         .iter()
         .map(|(call, var)| syn::parse_quote!(let mut #var = #call.clone();))
@@ -106,7 +176,7 @@ fn transform_function_calls(
         #(#declarations)*
     });
 
-    // Generate assignments block
+    // Generate declarations block of the newly generated variables (will be added on the loop of the loop body)
     let assignments: Vec<Stmt> = newreplace
         .into_iter()
         .map(|(call, var)| syn::parse_quote!(#var = #call.clone();))
@@ -138,30 +208,12 @@ impl VisitMut for BreakContinueReplacer {
     }
 }
 
+// This function replace the break/continue statements inside a loop body with return statements
 fn transform_break_continue(block: &mut Block) {
     let mut replacer = BreakContinueReplacer;
     replacer.visit_block_mut(block);
 }
 
-/// Expand loop contracts macros.
-///
-/// A while loop of the form
-/// ``` rust
-///  while guard {
-///      body
-///  }
-/// ```
-/// will be annotated as
-/// ``` rust
-/// #[inline(never)]
-/// #[kanitool::fn_marker = "kani_register_loop_contract"]
-/// const fn kani_register_loop_contract_id<T, F: FnOnce() -> T>(f: F) -> T {
-///     unreachable!()
-/// }
-///  while kani_register_loop_contract_id(|| -> bool {inv};) && guard {
-///      body
-///  }
-/// ```
 pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
     // parse the stmt of the loop
     let mut loop_stmt: Stmt = syn::parse(item.clone()).unwrap();

From 6a9e97c4e7aee650262c08f528cf0a3574082163 Mon Sep 17 00:00:00 2001
From: thanhnguyen-aws <ntson@amazon.com>
Date: Tue, 25 Mar 2025 08:38:04 -0700
Subject: [PATCH 10/18] separate continue/break for the case where the loop
 break on the first iter

---
 .../src/sysroot/loop_contracts/mod.rs         | 63 ++++++++++++-------
 .../loop_with_prev_break_first_iter.expected  |  4 ++
 .../loop_with_prev_break_first_iter.rs        | 22 +++++++
 3 files changed, 68 insertions(+), 21 deletions(-)
 create mode 100644 tests/expected/loop-contract/loop_with_prev_break_first_iter.expected
 create mode 100644 tests/expected/loop-contract/loop_with_prev_break_first_iter.rs

diff --git a/library/kani_macros/src/sysroot/loop_contracts/mod.rs b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
index e3ae51e853e3..63f821be86b2 100644
--- a/library/kani_macros/src/sysroot/loop_contracts/mod.rs
+++ b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
@@ -34,11 +34,13 @@ If there is a prev(expr) in the loop_invariant:
 1. Firstly, Add an assertion of the loop_quard: the loop must run at least once, otherwise prev(expr) is undefined.
 1. For each prev(expr) in the loop variant, replace it with a newly generated "memory" variable prev_k
 2. Add the declaration of i before the loop: let mut prev_k = expr
-3. Define a mut closure whose body is exactly the loop body, but replace all break and continue statements with return statements
-4. Call the closure (the same as run the loop once)
-5. Add the assertion for the loop_invariant (not includes the loop_quard): check if the loop_invariant holds after the first iteration.
-6. Add the loop
-6. Add the assignment statements (exactly the same as the declarations without the "let mut") on the top of the loop body to update the "memory" variables
+3. Define a mut closure whose body is exactly the loop body, but replace all continue/break statements with return true/false statements,
+        then add a final return true statement at the end of it
+4. Add an if statement with condition tobe the that closure's call (the same as run the loop once):
+    True block: add the loop with expanded macros (see next section) and inside the loop body:
+        add the assignment statements (exactly the same as the declarations without the "let mut") on the top to update the "memory" variables
+    Else block: Add the assertion for the loop_invariant (not includes the loop_quard): check if the loop_invariant holds after the first iteration.
+
 For example:
 #[kani::loop_invariant(prev(x+y) = x + y -1 && ...)]
 while(loop_guard)
@@ -53,14 +55,18 @@ let mut prev_1 = x + y;
 let mut loop_body_closure = || {
     loop_body_replaced //replace breaks/continues in loop_body with returns
 };
-loop_body_closure();
-assert!(prev_1 = x + y -1 && ...);
-#[kani::loop_invariant(prev_1  = x + y -1)]
-while(loop_guard)
-{
-    prev_1 = x + y;
-    loop_body
+if loop_body_closure(){
+    #[kani::loop_invariant(prev_1  = x + y -1)]
+    while(loop_guard)
+    {
+        prev_1 = x + y;
+        loop_body
+    }
 }
+else{
+    assert!(prev_1 = x + y -1 && ...);
+}
+
 
 */
 
@@ -198,10 +204,10 @@ impl VisitMut for BreakContinueReplacer {
         // Replace the expression
         *expr = match expr {
             Expr::Break(_) => {
-                syn::parse_quote!(return)
+                syn::parse_quote!(return false)
             }
             Expr::Continue(_) => {
-                syn::parse_quote!(return)
+                syn::parse_quote!(return true)
             }
             _ => return,
         };
@@ -212,6 +218,18 @@ impl VisitMut for BreakContinueReplacer {
 fn transform_break_continue(block: &mut Block) {
     let mut replacer = BreakContinueReplacer;
     replacer.visit_block_mut(block);
+    let return_stmt: Stmt = syn::parse_quote! {
+        return true;
+    };
+    // Add semicolon to the last statement if it's an expression without semicolon
+    if let Some(last_stmt) = block.stmts.last_mut() {
+        if let Stmt::Expr(expr, ref mut semi) = last_stmt {
+            if semi.is_none() {
+                *semi = Some(Default::default());
+            }
+        }
+    }
+    block.stmts.push(return_stmt);
 }
 
 pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
@@ -307,17 +325,20 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
         #(#decl_stms)*
         let mut #loop_body_closure = ||
         #loop_body;
-        #loop_body_closure ();
-        assert!(#inv_expr);
+        if #loop_body_closure () {
         // Dummy function used to force the compiler to capture the environment.
         // We cannot call closures inside constant functions.
         // This function gets replaced by `kani::internal::call_closure`.
-        #[inline(never)]
-        #[kanitool::fn_marker = "kani_register_loop_contract"]
-        const fn #register_ident<F: Fn() -> bool>(_f: &F, _transformed: usize) -> bool {
-            true
+            #[inline(never)]
+            #[kanitool::fn_marker = "kani_register_loop_contract"]
+            const fn #register_ident<F: Fn() -> bool>(_f: &F, _transformed: usize) -> bool {
+                true
+            }
+            #loop_stmt
         }
-        #loop_stmt
+        else {
+            assert!(#inv_expr);
+        };
         })
         .into()
     } else {
diff --git a/tests/expected/loop-contract/loop_with_prev_break_first_iter.expected b/tests/expected/loop-contract/loop_with_prev_break_first_iter.expected
new file mode 100644
index 000000000000..b55ae59c4b5f
--- /dev/null
+++ b/tests/expected/loop-contract/loop_with_prev_break_first_iter.expected
@@ -0,0 +1,4 @@
+Failed Checks: assertion failed: (i >= 2) && (i <= 100) && (__kani_prev_var_15_5_18_6_0 == i + 2)
+ File: "src/exptest.rs", line 14, in loop_with_prev
+
+VERIFICATION:- FAILED
diff --git a/tests/expected/loop-contract/loop_with_prev_break_first_iter.rs b/tests/expected/loop-contract/loop_with_prev_break_first_iter.rs
new file mode 100644
index 000000000000..87cab0426373
--- /dev/null
+++ b/tests/expected/loop-contract/loop_with_prev_break_first_iter.rs
@@ -0,0 +1,22 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+// kani-flags: -Z loop-contracts
+
+//! Check if loop contracts is correctly applied.
+
+#![feature(stmt_expr_attributes)]
+#![feature(proc_macro_hygiene)]
+
+#[kani::proof]
+pub fn loop_with_prev() {
+    let mut i = 100;
+    #[kani::loop_invariant((i >= 2) && (i <= 100) && (prev(i) == i + 2))]
+    while i > 2 {
+        if i == 100 {
+            break;
+        }
+        i = i - 2;
+    }
+    assert!(i == 2);
+}

From 4338359d72a783943e191361e191eb47896534dc Mon Sep 17 00:00:00 2001
From: thanhnguyen-aws <ntson@amazon.com>
Date: Tue, 25 Mar 2025 09:58:26 -0700
Subject: [PATCH 11/18] fixed expected

---
 .../loop-contract/loop_with_prev_break_first_iter.expected       | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tests/expected/loop-contract/loop_with_prev_break_first_iter.expected b/tests/expected/loop-contract/loop_with_prev_break_first_iter.expected
index b55ae59c4b5f..c37aa96eca8d 100644
--- a/tests/expected/loop-contract/loop_with_prev_break_first_iter.expected
+++ b/tests/expected/loop-contract/loop_with_prev_break_first_iter.expected
@@ -1,4 +1,3 @@
 Failed Checks: assertion failed: (i >= 2) && (i <= 100) && (__kani_prev_var_15_5_18_6_0 == i + 2)
- File: "src/exptest.rs", line 14, in loop_with_prev
 
 VERIFICATION:- FAILED

From 000c3847ca19e17d0fb03ac6b576a4466119595e Mon Sep 17 00:00:00 2001
From: thanhnguyen-aws <ntson@amazon.com>
Date: Tue, 25 Mar 2025 13:22:58 -0700
Subject: [PATCH 12/18] fixed expected

---
 .../loop-contract/loop_with_prev_break_first_iter.expected      | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/expected/loop-contract/loop_with_prev_break_first_iter.expected b/tests/expected/loop-contract/loop_with_prev_break_first_iter.expected
index c37aa96eca8d..68dfcd51d085 100644
--- a/tests/expected/loop-contract/loop_with_prev_break_first_iter.expected
+++ b/tests/expected/loop-contract/loop_with_prev_break_first_iter.expected
@@ -1,3 +1,3 @@
-Failed Checks: assertion failed: (i >= 2) && (i <= 100) && (__kani_prev_var_15_5_18_6_0 == i + 2)
+Failed Checks: assertion failed: (i >= 2) && (i <= 100) && (__kani_prev_var_
 
 VERIFICATION:- FAILED

From 8f7f0674084adf9ad021a8abbe9b33e8ec5a6f29 Mon Sep 17 00:00:00 2001
From: thanhnguyen-aws <ntson@amazon.com>
Date: Tue, 1 Apr 2025 13:13:46 -0700
Subject: [PATCH 13/18] fix some typos

---
 library/kani_macros/src/sysroot/loop_contracts/mod.rs | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/library/kani_macros/src/sysroot/loop_contracts/mod.rs b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
index 63f821be86b2..0c0a3dff1824 100644
--- a/library/kani_macros/src/sysroot/loop_contracts/mod.rs
+++ b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
@@ -31,12 +31,12 @@ Then the loop_invartiant is transformed
 /*
 Transform the loop to support prev(expr) : the value of expr at the end of the previous iteration
 If there is a prev(expr) in the loop_invariant:
-1. Firstly, Add an assertion of the loop_quard: the loop must run at least once, otherwise prev(expr) is undefined.
-1. For each prev(expr) in the loop variant, replace it with a newly generated "memory" variable prev_k
-2. Add the declaration of i before the loop: let mut prev_k = expr
-3. Define a mut closure whose body is exactly the loop body, but replace all continue/break statements with return true/false statements,
+1. Firstly, add an assertion of the loop_quard: the loop must run at least once, otherwise prev(expr) is undefined.
+2. For each prev(expr) in the loop variant, replace it with a newly generated "memory" variable prev_k
+3. Add the declaration of prev_k before the loop: let mut prev_k = expr
+4. Define a mut closure whose body is exactly the loop body, but replace all continue/break statements with return true/false statements,
         then add a final return true statement at the end of it
-4. Add an if statement with condition tobe the that closure's call (the same as run the loop once):
+5. Add an if statement with condition to be the that closure's call (the same as run the loop once):
     True block: add the loop with expanded macros (see next section) and inside the loop body:
         add the assignment statements (exactly the same as the declarations without the "let mut") on the top to update the "memory" variables
     Else block: Add the assertion for the loop_invariant (not includes the loop_quard): check if the loop_invariant holds after the first iteration.

From 8bb5b10c6cfdbbf0487ebc575f03ea0f91ca92fb Mon Sep 17 00:00:00 2001
From: thanhnguyen-aws <ntson@amazon.com>
Date: Thu, 3 Apr 2025 10:07:35 -0700
Subject: [PATCH 14/18] update prev semantic

---
 .../kani_macros/src/sysroot/loop_contracts/mod.rs   | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/library/kani_macros/src/sysroot/loop_contracts/mod.rs b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
index 0c0a3dff1824..001fbd8c2d8e 100644
--- a/library/kani_macros/src/sysroot/loop_contracts/mod.rs
+++ b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
@@ -30,8 +30,12 @@ Then the loop_invartiant is transformed
 
 /*
 Transform the loop to support prev(expr) : the value of expr at the end of the previous iteration
+Semantic: If the loop has at least 1 iteration: prev(expr) is the value of expr at the end of the previous iteration. Otherwise, just remove the loop (without check for the invariant too).
+
+Transformation: basically, if the loop has at least 1 iteration (loop_quard is satisfied at the beginning), we unfold the loop once, declare the variables for prev values and update them at the beginning of the loop body.
+Otherwise, we remove the loop.
 If there is a prev(expr) in the loop_invariant:
-1. Firstly, add an assertion of the loop_quard: the loop must run at least once, otherwise prev(expr) is undefined.
+1. Firstly, add an if block whose condition is the loop_quard, inside its body add/do the followings:
 2. For each prev(expr) in the loop variant, replace it with a newly generated "memory" variable prev_k
 3. Add the declaration of prev_k before the loop: let mut prev_k = expr
 4. Define a mut closure whose body is exactly the loop body, but replace all continue/break statements with return true/false statements,
@@ -259,7 +263,7 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
     let transform_inv: TransformationResult =
         transform_function_calls(inv_expr.clone(), "prev", prev_var_prefix);
     let has_prev = !transform_inv.declarations_block.stmts.is_empty();
-    let decl_stms = transform_inv.declarations_block.stmts.clone();
+    let prev_decl_stms = transform_inv.declarations_block.stmts.clone();
     let mut assign_stms = transform_inv.assignments_block.stmts.clone();
     let (mut loop_body, loop_guard) = match loop_stmt {
         Stmt::Expr(ref mut e, _) => match e {
@@ -320,9 +324,9 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
     if has_prev {
         quote!(
         {
+        if (#loop_guard) {
         #(#old_decl_stms)*
-        assert!(#loop_guard);
-        #(#decl_stms)*
+        #(#prev_decl_stms)*
         let mut #loop_body_closure = ||
         #loop_body;
         if #loop_body_closure () {
@@ -339,6 +343,7 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
         else {
             assert!(#inv_expr);
         };
+        }
         })
         .into()
     } else {

From fcd878dbbb9ea8ec8509cabaafca2b3f4b5ff6eb Mon Sep 17 00:00:00 2001
From: thanhnguyen-aws <ntson@amazon.com>
Date: Thu, 10 Apr 2025 13:45:17 -0700
Subject: [PATCH 15/18] fix comment

---
 library/kani_macros/src/sysroot/loop_contracts/mod.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/library/kani_macros/src/sysroot/loop_contracts/mod.rs b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
index 001fbd8c2d8e..51f95b1f08d6 100644
--- a/library/kani_macros/src/sysroot/loop_contracts/mod.rs
+++ b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
@@ -249,7 +249,7 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
     // expr of the loop invariant
     let mut inv_expr: Expr = syn::parse(attr).unwrap();
 
-    // adding old variables
+    // adding on_entry variables
     let mut old_var_prefix: String = "__kani_old_var".to_owned();
     old_var_prefix.push_str(&loop_id);
     let replace_old: TransformationResult =

From 132723eea2fc57b9c1e6ee93aa082858702f8efd Mon Sep 17 00:00:00 2001
From: thanhnguyen-aws <ntson@amazon.com>
Date: Thu, 10 Apr 2025 15:27:25 -0700
Subject: [PATCH 16/18] separate return statement in loop body

---
 .../src/sysroot/loop_contracts/mod.rs         | 45 +++++++++++++------
 1 file changed, 32 insertions(+), 13 deletions(-)

diff --git a/library/kani_macros/src/sysroot/loop_contracts/mod.rs b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
index 51f95b1f08d6..22c265205982 100644
--- a/library/kani_macros/src/sysroot/loop_contracts/mod.rs
+++ b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
@@ -208,22 +208,29 @@ impl VisitMut for BreakContinueReplacer {
         // Replace the expression
         *expr = match expr {
             Expr::Break(_) => {
-                syn::parse_quote!(return false)
+                syn::parse_quote!(return (false, None))
             }
             Expr::Continue(_) => {
-                syn::parse_quote!(return true)
+                syn::parse_quote!(return (true, None))
+            }
+            Expr::Return(rexpr) => {
+                match rexpr.expr.clone() {
+                    Some(ret) => syn::parse_quote!(return (false, Some(#ret))),
+                    _ => syn::parse_quote!(return (false, Some(())))
+                }
             }
             _ => return,
         };
     }
 }
 
+
 // This function replace the break/continue statements inside a loop body with return statements
 fn transform_break_continue(block: &mut Block) {
     let mut replacer = BreakContinueReplacer;
     replacer.visit_block_mut(block);
     let return_stmt: Stmt = syn::parse_quote! {
-        return true;
+        return (true, None);
     };
     // Add semicolon to the last statement if it's an expression without semicolon
     if let Some(last_stmt) = block.stmts.last_mut() {
@@ -236,6 +243,8 @@ fn transform_break_continue(block: &mut Block) {
     block.stmts.push(return_stmt);
 }
 
+
+
 pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
     // parse the stmt of the loop
     let mut loop_stmt: Stmt = syn::parse(item.clone()).unwrap();
@@ -250,12 +259,12 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
     let mut inv_expr: Expr = syn::parse(attr).unwrap();
 
     // adding on_entry variables
-    let mut old_var_prefix: String = "__kani_old_var".to_owned();
-    old_var_prefix.push_str(&loop_id);
-    let replace_old: TransformationResult =
-        transform_function_calls(inv_expr.clone(), "on_entry", old_var_prefix);
-    inv_expr = replace_old.transformed_expr.clone();
-    let old_decl_stms = replace_old.declarations_block.stmts.clone();
+    let mut onentry_var_prefix: String = "__kani_onentry_var".to_owned();
+    onentry_var_prefix.push_str(&loop_id);
+    let replace_onentry: TransformationResult =
+        transform_function_calls(inv_expr.clone(), "on_entry", onentry_var_prefix);
+    inv_expr = replace_onentry.transformed_expr.clone();
+    let onentry_decl_stms = replace_onentry.declarations_block.stmts.clone();
 
     // adding prev variables
     let mut prev_var_prefix: String = "__kani_prev_var".to_owned();
@@ -278,6 +287,12 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
     let mut loop_body_closure_name: String = "__kani_loop_body_closure".to_owned();
     loop_body_closure_name.push_str(&loop_id);
     let loop_body_closure = format_ident!("{}", loop_body_closure_name);
+    let mut loop_body_closure_ret_1_name: String = "__kani_loop_body_closure_ret_1".to_owned();
+    loop_body_closure_ret_1_name.push_str(&loop_id);
+    let loop_body_closure_ret_1 = format_ident!("{}", loop_body_closure_ret_1_name);
+    let mut loop_body_closure_ret_2_name: String = "__kani_loop_body_closure_ret_2".to_owned();
+    loop_body_closure_ret_2_name.push_str(&loop_id);
+    let loop_body_closure_ret_2 = format_ident!("{}", loop_body_closure_ret_2_name);
     if has_prev {
         inv_expr = transform_inv.transformed_expr.clone();
         match loop_stmt {
@@ -325,11 +340,15 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
         quote!(
         {
         if (#loop_guard) {
-        #(#old_decl_stms)*
+        #(#onentry_decl_stms)*
         #(#prev_decl_stms)*
         let mut #loop_body_closure = ||
         #loop_body;
-        if #loop_body_closure () {
+        let (#loop_body_closure_ret_1, #loop_body_closure_ret_2) = #loop_body_closure ();
+        if #loop_body_closure_ret_2.is_some() {
+            return #loop_body_closure_ret_2.unwrap();
+        }
+        if #loop_body_closure_ret_1 {
         // Dummy function used to force the compiler to capture the environment.
         // We cannot call closures inside constant functions.
         // This function gets replaced by `kani::internal::call_closure`.
@@ -345,11 +364,11 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
         };
         }
         })
-        .into()
+        .into();
     } else {
         quote!(
         {
-        #(#old_decl_stms)*
+        #(#onentry_decl_stms)*
         // Dummy function used to force the compiler to capture the environment.
         // We cannot call closures inside constant functions.
         // This function gets replaced by `kani::internal::call_closure`.

From 609c38841c6533302fa23e32b1685f53e1afa3bc Mon Sep 17 00:00:00 2001
From: thanhnguyen-aws <ntson@amazon.com>
Date: Fri, 11 Apr 2025 08:34:29 -0700
Subject: [PATCH 17/18] fix format

---
 .../kani_macros/src/sysroot/loop_contracts/mod.rs   | 13 ++++---------
 1 file changed, 4 insertions(+), 9 deletions(-)

diff --git a/library/kani_macros/src/sysroot/loop_contracts/mod.rs b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
index 22c265205982..dbcc7264004a 100644
--- a/library/kani_macros/src/sysroot/loop_contracts/mod.rs
+++ b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
@@ -213,18 +213,15 @@ impl VisitMut for BreakContinueReplacer {
             Expr::Continue(_) => {
                 syn::parse_quote!(return (true, None))
             }
-            Expr::Return(rexpr) => {
-                match rexpr.expr.clone() {
-                    Some(ret) => syn::parse_quote!(return (false, Some(#ret))),
-                    _ => syn::parse_quote!(return (false, Some(())))
-                }
-            }
+            Expr::Return(rexpr) => match rexpr.expr.clone() {
+                Some(ret) => syn::parse_quote!(return (false, Some(#ret))),
+                _ => syn::parse_quote!(return (false, Some(()))),
+            },
             _ => return,
         };
     }
 }
 
-
 // This function replace the break/continue statements inside a loop body with return statements
 fn transform_break_continue(block: &mut Block) {
     let mut replacer = BreakContinueReplacer;
@@ -243,8 +240,6 @@ fn transform_break_continue(block: &mut Block) {
     block.stmts.push(return_stmt);
 }
 
-
-
 pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
     // parse the stmt of the loop
     let mut loop_stmt: Stmt = syn::parse(item.clone()).unwrap();

From 747d2a38da1b08f826572aa4ddf6cbbecf09eb3f Mon Sep 17 00:00:00 2001
From: thanhnguyen-aws <ntson@amazon.com>
Date: Fri, 11 Apr 2025 08:45:00 -0700
Subject: [PATCH 18/18] fix typo

---
 library/kani_macros/src/sysroot/loop_contracts/mod.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/library/kani_macros/src/sysroot/loop_contracts/mod.rs b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
index dbcc7264004a..844ca7be8165 100644
--- a/library/kani_macros/src/sysroot/loop_contracts/mod.rs
+++ b/library/kani_macros/src/sysroot/loop_contracts/mod.rs
@@ -359,7 +359,7 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
         };
         }
         })
-        .into();
+        .into()
     } else {
         quote!(
         {