Ensure that contract closures are FnOnce

kani-compiler/src/kani_middle/kani_functions.rs

@@ -91,6 +91,10 @@ pub enum KaniModel {
     PtrOffsetFromUnsigned,
     #[strum(serialize = "RunContractModel")]
     RunContract,
+    #[strum(serialize = "ForceContractTypeModel")]
+    ForceContractType,
+    #[strum(serialize = "ForceContractTypeWithArgumentsModel")]
+    ForceContractTypeWithArguments,
     #[strum(serialize = "RunLoopContractModel")]
     RunLoopContract,
     #[strum(serialize = "SetPtrInitializedModel")]

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

@@ -282,6 +282,10 @@ pub struct FunctionWithContractPass {
     unused_closures: HashSet<ClosureDef>,
     /// Cache KaniRunContract function used to implement contracts.
     run_contract_fn: Option<FnDef>,
+    /// Kani ForceFnOnce function used to enforce correct closure annotation.
+    force_fn_once: Option<FnDef>,
+    /// Kani ForceFnOnceWithArgs function used to enforce correct closure annotation.
+    force_fn_once_with_args: Option<FnDef>,
 }
 
 impl TransformPass for FunctionWithContractPass {
@@ -313,6 +317,17 @@ impl TransformPass for FunctionWithContractPass {
                 {
                     let run = Instance::resolve(self.run_contract_fn.unwrap(), args).unwrap();
                     (true, run.body().unwrap())
+                } else if KaniAttributes::for_instance(tcx, instance).fn_marker()
+                    == Some(Symbol::intern("kani_force_fn_once"))
+                {
+                    let run = Instance::resolve(self.force_fn_once.unwrap(), args).unwrap();
+                    (true, run.body().unwrap())
+                } else if KaniAttributes::for_instance(tcx, instance).fn_marker()
+                    == Some(Symbol::intern("kani_force_fn_once_with_args"))
+                {
+                    let run =
+                        Instance::resolve(self.force_fn_once_with_args.unwrap(), args).unwrap();
+                    (true, run.body().unwrap())
                 } else {
                     // Not a contract annotated function
                     (false, body)
@@ -372,12 +387,28 @@ impl FunctionWithContractPass {
             let run_contract_fn =
                 queries.kani_functions().get(&KaniModel::RunContract.into()).copied();
             assert!(run_contract_fn.is_some(), "Failed to find Kani run contract function");
+
+            let force_fn_once =
+                queries.kani_functions().get(&KaniModel::ForceContractType.into()).copied();
+            assert!(force_fn_once.is_some(), "Failed to find Kani force_fn_once function");
+
+            let force_fn_once_with_args = queries
+                .kani_functions()
+                .get(&KaniModel::ForceContractTypeWithArguments.into())
+                .copied();
+            assert!(
+                force_fn_once_with_args.is_some(),
+                "Failed to find Kani force_fn_once_with_args function"
+            );
+
             FunctionWithContractPass {
                 check_fn,
                 replace_fns,
                 assert_contracts: !queries.args().no_assert_contracts,
                 unused_closures: Default::default(),
                 run_contract_fn,
+                force_fn_once,
+                force_fn_once_with_args,
             }
         } else {
             // If reachability mode is PubFns or Tests, we just remove any contract logic.

library/kani_core/src/lib.rs

@@ -567,6 +567,30 @@ macro_rules! kani_intrinsics {
                 func()
             }
 
+            /// Function that wraps a closure without arguments used to implement contracts.
+            ///
+            /// To ensure Rust correctly infers type of a contract closure as FnOnce, we wrap it in a dummy
+            /// function that explicitly requires an FnOnce. This wrapping is done at the point of closure
+            /// definition.
+            #[doc(hidden)]
+            #[allow(dead_code)]
+            #[kanitool::fn_marker = "ForceContractTypeModel"]
+            fn force_fn_once<T, F: FnOnce() -> T>(func: F) -> F {
+                func
+            }
+
+            /// Function that wraps a closure with one argument used to implement contracts.
+            ///
+            /// To ensure Rust correctly infers type of a contract closure as FnOnce, we wrap it in a dummy
+            /// function that explicitly requires an FnOnce. This wrapping is done at the point of closure
+            /// definition.
+            #[doc(hidden)]
+            #[allow(dead_code)]
+            #[kanitool::fn_marker = "ForceContractTypeWithArgumentsModel"]
+            fn force_fn_once_with_args<A, T, F: FnOnce(A) -> T>(func: F) -> F {
+                func
+            }
+
             /// Function that calls a closure used to implement loop contracts.
             ///
             /// In contracts, we cannot invoke the generated closures directly, instead, we call register

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,18 @@ 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.
+                #[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 +137,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 +151,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)

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

@@ -133,7 +133,7 @@ impl<'a> ContractConditionsHandler<'a> {
         quote!(
             #[kanitool::is_contract_generated(replace)]
             #[allow(dead_code, unused_variables, unused_mut)]
-            let mut #replace_ident = || #output #body;
+            let mut #replace_ident = kani_force_fn_once(|| #output #body);
         )
     }
 

tests/expected/function-contract/mutable-references/ensures_with_two_mut_refs_fail.expected

@@ -0,0 +1 @@
+error[E0501]: cannot borrow value as immutable because previous closure requires unique access

tests/expected/function-contract/mutable-references/ensures_with_two_mut_refs_fail.rs

@@ -0,0 +1,15 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+// kani-flags: -Z function-contracts
+
+#[kani::ensures(|result| **result == 42 && *result == n)]
+#[kani::modifies(n)]
+fn forty_two(n: &mut u8) -> &mut u8 {
+    *n = 42;
+    n
+}
+
+#[kani::proof_for_contract(forty_two)]
+fn forty_two_harness() {
+    forty_two(&mut kani::any());
+}

tests/expected/function-contract/mutable-references/return_mut_ref.expected

@@ -0,0 +1,11 @@
+assertion\
+	- Status: UNREACHABLE\
+	- Description: "attempt to divide by zero"\
+in function divide_by
+
+assertion\
+	- Status: SUCCESS\
+	- Description: "|result| **result <= a"\
+in function divide_by
+
+VERIFICATION:- SUCCESSFUL

tests/expected/function-contract/mutable-references/return_mut_ref.rs

@@ -0,0 +1,16 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+// kani-flags: -Z function-contracts
+
+#[kani::requires(*b > 0 && *b < a)]
+#[kani::ensures(|result| **result <= a)]
+#[kani::modifies(b)]
+fn divide_by(a: u8, b: &mut u8) -> &mut u8 {
+    *b = a / *b;
+    b
+}
+
+#[kani::proof_for_contract(divide_by)]
+fn divide_by_harness() {
+    divide_by(kani::any(), &mut kani::any());
+}