Ensure that contract closures are FnOnce (#4151)

Rust believes that Kani's contract closures are `FnMut`. This prevents
us from writing contracts for functions that return mutable references
to their input arguments (#3764).

To ensure Rust correctly infers these closures as `FnOnce`, they need to
be wrapped in a dummy function that explicitly requires an `FnOnce`.
This wrapping must be done at the point of closure definition, as doing
it later, when calling the function, doesn't seem to have any effect.

Resolves #3764

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

---------

Co-authored-by: Fedor Ryabinin <rfedor@amazon.com>
Co-authored-by: Carolyn Zech <cmzech@amazon.com>

Author: 3 people

library/kani_macros/src/sysroot/contracts/assert.rs

@@ -29,7 +29,7 @@ impl<'a> ContractConditionsHandler<'a> {
         quote!(
             #[kanitool::is_contract_generated(assert)]
             #[allow(dead_code, unused_variables, unused_mut)]
-            let mut #assert_ident = || #output #body;
+            let mut #assert_ident = kani_force_fn_once(|| #output #body);
         )
     }
 
@@ -49,9 +49,9 @@ impl<'a> ContractConditionsHandler<'a> {
         let result = Ident::new(INTERNAL_RESULT_IDENT, Span::call_site());
 
         parse_quote!(
-            let mut body_wrapper = || #output {
+            let mut body_wrapper = kani_force_fn_once(|| #output {
                 #(#stmts)*
-            };
+            });
             let #result : #return_type = #body_wrapper_ident();
             #result
         )

library/kani_macros/src/sysroot/contracts/bootstrap.rs

@@ -47,6 +47,20 @@ impl<'a> ContractConditionsHandler<'a> {
             #[kanitool::asserted_with = #assert_name]
             #[kanitool::modifies_wrapper = #modifies_name]
             #vis #sig {
+                // Dummy functions used to force the compiler to annotate Kani's
+                // closures as `FnOnce`. Without this, Rust infers the generated closures as `FnMut`,
+                // preventing us from writing contracts for functions returning mutable references.
+                // See https://github.com/model-checking/kani/issues/3764
+                #[inline(never)]
+                #[kanitool::fn_marker = "kani_force_fn_once"]
+                const fn kani_force_fn_once<T, F: FnOnce() -> T>(f: F) -> F {
+                    f
+                }
+                #[inline(never)]
+                #[kanitool::fn_marker = "kani_force_fn_once_with_args"]
+                const fn kani_force_fn_once_with_args<A, T, F: FnOnce(A) -> T>(f: F) -> F {
+                    f
+                }
                 // 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`.
@@ -125,7 +139,7 @@ impl<'a> ContractConditionsHandler<'a> {
         quote!(
             #[kanitool::is_contract_generated(recursion_check)]
             #[allow(dead_code, unused_variables, unused_mut)]
-            let mut #recursion_ident = || #output
+            let mut #recursion_ident = kani_force_fn_once(|| #output
             {
                 #[kanitool::recursion_tracker]
                 static mut REENTRY: bool = false;
@@ -139,7 +153,7 @@ impl<'a> ContractConditionsHandler<'a> {
                     unsafe { REENTRY = false };
                     #result
                 }
-            };
+            });
         )
     }
 

library/kani_macros/src/sysroot/contracts/check.rs

@@ -114,7 +114,7 @@ impl<'a> ContractConditionsHandler<'a> {
         quote!(
             #[kanitool::is_contract_generated(check)]
             #[allow(dead_code, unused_variables, unused_mut)]
-            let mut #check_ident = || #output #body;
+            let mut #check_ident = kani_force_fn_once(|| #output #body);
         )
     }
 
@@ -144,10 +144,10 @@ impl<'a> ContractConditionsHandler<'a> {
         quote!(
             #[kanitool::is_contract_generated(wrapper)]
             #[allow(dead_code, unused_variables, unused_mut)]
-            let mut #modifies_ident = |#wrapper_ident: _| #output {
+            let mut #modifies_ident = kani_force_fn_once_with_args(|#wrapper_ident: _| #output {
                 #redefs
                 #(#stmts)*
-            };
+            });
         )
     }
 
@@ -157,7 +157,13 @@ impl<'a> ContractConditionsHandler<'a> {
             let Stmt::Local(Local { init: Some(LocalInit { expr, .. }), .. }) = closure_stmt else {
                 unreachable!()
             };
-            let Expr::Closure(closure) = expr.as_ref() else { unreachable!() };
+            // The closure is wrapped into `kani_force_fn_once_with_args`
+            let Expr::Call(call) = expr.as_ref() else { unreachable!() };
+            if call.args.len() != 1 {
+                unreachable!()
+            }
+            let expr = call.args.first().unwrap();
+            let Expr::Closure(closure) = expr else { unreachable!() };
             let Expr::Block(body) = closure.body.as_ref() else { unreachable!() };
             let stream = self.modifies_closure(&closure.output, &body.block, TokenStream2::new());
             *closure_stmt = syn::parse2(stream).unwrap();

library/kani_macros/src/sysroot/contracts/helpers.rs

@@ -123,9 +123,25 @@ pub fn closure_body(closure: &mut Stmt) -> &mut ExprBlock {
     let Stmt::Local(Local { init: Some(LocalInit { expr, .. }), .. }) = closure else {
         unreachable!()
     };
-    let Expr::Closure(closure) = expr.as_mut() else { unreachable!() };
-    let Expr::Block(body) = closure.body.as_mut() else { unreachable!() };
-    body
+    match expr.as_mut() {
+        // The case of closures wrapped in `kani_force_fn_once`
+        Expr::Call(call) if call.args.len() == 1 => {
+            let arg = call.args.first_mut().unwrap();
+            match arg {
+                Expr::Closure(closure) => {
+                    let Expr::Block(body) = closure.body.as_mut() else { unreachable!() };
+                    body
+                }
+                _ => unreachable!(),
+            }
+        }
+
+        Expr::Closure(closure) => {
+            let Expr::Block(body) = closure.body.as_mut() else { unreachable!() };
+            body
+        }
+        _ => unreachable!(),
+    }
 }
 
 /// Does the provided path have the same chain of identifiers as `mtch` (match)