Add loop invariant support for while let loop (#4279)

This PR adds loop invariant support for `while let` loop, and also fix
the bug reported in issue #4275

I don't separate them because the bug fix is simple and necessary for
the expected test of `while let` loop.

Resolves #4275
By submitting this pull request, I confirm that my contribution is made
under the terms of the Apache 2.0 and MIT licenses.

Author: thanhnguyen-aws

kani-compiler/src/kani_middle/transform/loop_contracts.rs

@@ -701,7 +701,7 @@ impl LoopContractPass {
                 for stmt in &new_body.blocks()[bb_idx].statements {
                     if let StatementKind::Assign(place, rvalue) = &stmt.kind {
                         match rvalue {
-                            Rvalue::Ref(_,_,rplace) | Rvalue::CopyForDeref(rplace) => {
+                            Rvalue::Ref(_,_,rplace) | Rvalue::CopyForDeref(rplace) | Rvalue::Use(Operand::Copy(rplace)) => {
                                 if supported_vars.contains(&rplace.local) {
                                     supported_vars.push(place.local);
                                 } }

library/kani_macros/src/sysroot/loop_contracts/mod.rs

@@ -6,12 +6,12 @@
 
 use proc_macro::TokenStream;
 use proc_macro_error2::abort_call_site;
-use quote::{format_ident, quote};
+use quote::{format_ident, quote, quote_spanned};
 use syn::punctuated::Punctuated;
 use syn::spanned::Spanned;
 use syn::token::AndAnd;
 use syn::{
-    BinOp, Block, Expr, ExprBinary, Ident, Stmt, Token, parse_macro_input, parse_quote,
+    BinOp, Block, Expr, ExprBinary, ExprWhile, Ident, Stmt, Token, parse_macro_input, parse_quote,
     visit_mut::VisitMut,
 };
 
@@ -232,9 +232,32 @@ fn transform_break_continue(block: &mut Block) {
     block.stmts.push(return_stmt);
 }
 
+fn while_let_rewrite(loopexpr: Stmt) -> Stmt {
+    if let Stmt::Expr(ref expr, _) = loopexpr
+        && let Expr::While(ExprWhile { cond, body, .. }) = expr
+        && let Expr::Let(ref let_expr) = **cond
+    {
+        let pat = &let_expr.pat;
+        let scrutinee = &let_expr.expr;
+
+        // Transform to loop with match
+        return parse_quote! {
+            loop {
+                match #scrutinee {
+                    #pat => #body,
+                    _ => break,
+                }
+            };
+        };
+    }
+
+    loopexpr.clone()
+}
+
 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();
+    loop_stmt = while_let_rewrite(loop_stmt);
 
     // name of the loop invariant as closure of the form
     // __kani_loop_invariant_#startline_#startcol_#endline_#endcol
@@ -244,6 +267,7 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
 
     // expr of the loop invariant
     let mut inv_expr: Expr = syn::parse(attr).unwrap();
+    let original_span = inv_expr.span();
 
     // adding on_entry variables
     let mut onentry_var_prefix: String = "__kani_onentry_var".to_owned();
@@ -302,7 +326,7 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
         &mut Stmt::Expr(ref mut e, _) => match *e {
             Expr::While(ref mut ew) => {
                 let new_cond: Expr = syn::parse(
-                    quote!(
+                    quote_spanned!(original_span =>
                         #register_ident(&||->bool{#inv_expr}, 0))
                     .into(),
                 )
@@ -316,7 +340,9 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
             }
             Expr::Loop(ref mut el) => {
                 //let retexpr = get_return_statement(&el.body);
-                let invstmt: Stmt = syn::parse(quote!(if !(#register_ident(&||->bool{#inv_expr}, 0)) {assert!(false); unreachable!()};).into()).unwrap();
+                let invstmt: Stmt = syn::parse(quote_spanned!(original_span =>
+                    if !(#register_ident(&||->bool{#inv_expr}, 0)) {assert!(false); unreachable!()};)
+                    .into()).unwrap();
                 let mut new_stmts: Vec<Stmt> = Vec::new();
                 new_stmts.push(invstmt);
                 new_stmts.extend(el.body.stmts.clone());
@@ -334,7 +360,7 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
     }
 
     if has_prev {
-        quote!(
+        quote_spanned!(original_span =>
         {
         if (#loop_guard) {
         #(#onentry_decl_stms)*
@@ -363,7 +389,7 @@ pub fn loop_invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
         })
         .into()
     } else {
-        quote!(
+        quote_spanned!(original_span =>
         {
         #(#onentry_decl_stms)*
         // Dummy function used to force the compiler to capture the environment.

tests/expected/loop-contract/while_let_loop_for_slice.expected

@@ -0,0 +1,9 @@
+trim_ascii_start.loop_invariant_base.1\
+	 - Status: SUCCESS\
+	 - Description: "Check invariant before entry for loop trim_ascii_start.0"
+
+trim_ascii_start.loop_invariant_step.1\
+	 - Status: SUCCESS\
+	 - Description: "Check invariant after step for loop trim_ascii_start.0"
+
+VERIFICATION:- SUCCESSFUL

tests/expected/loop-contract/while_let_loop_for_slice.rs

@@ -0,0 +1,27 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+// kani-flags: -Z loop-contracts
+
+//! Check if while-let invariant is correctly applied.
+
+#![feature(proc_macro_hygiene)]
+#![feature(stmt_expr_attributes)]
+
+#[kani::proof]
+fn trim_ascii_start() {
+    let mut a: [u8; 10] = kani::any();
+    let s = a.as_slice();
+    let mut bytes = s;
+    #[kani::loop_invariant(
+            bytes.len() <= s.len() &&
+            (bytes.len() == 0 || (&s[s.len()-bytes.len()..]).as_ptr() == bytes.as_ptr())
+    )]
+    while let [first, rest @ ..] = bytes {
+        if first.is_ascii_whitespace() {
+            bytes = rest;
+        } else {
+            break;
+        }
+    }
+}