add while let loop transformation

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 {
+        if let Expr::While(ExprWhile { cond, body, .. }) = expr {
+            if 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,
+                        }
+                    };
+                };
+            }
+        }
+    }
+    return 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.