From 6f0c0b5b97c6861a3113af85c5d94bc17ee3e306 Mon Sep 17 00:00:00 2001
From: Matias Scharager <matias.scharager@gmail.com>
Date: Mon, 1 Jul 2024 12:39:32 -0400
Subject: [PATCH] Function Contracts: Closure Type Inference (#3307)

The rust type inference for closures doesn't work in the particular use
case we are using it for ensures clauses. By creating a helper function,
we change the path the rust type inference takes and lets the type of
the closure be identified properly. This means type annotations are no
longer required within ensures clauses.

Resolves #3304

By submitting this pull request, I confirm that my contribution is made
under the terms of the Apache 2.0 and MIT licenses.
---
 library/kani/src/internal.rs                       |  7 +++++++
 .../kani_macros/src/sysroot/contracts/shared.rs    |  2 +-
 .../simple_ensures_pass_no_annotation.expected     |  6 ++++++
 .../simple_ensures_pass_no_annotation.rs           | 14 ++++++++++++++
 .../mutating_ensures_error.expected                |  1 +
 .../function-contracts/mutating_ensures_error.rs   | 12 ++++++++++++
 6 files changed, 41 insertions(+), 1 deletion(-)
 create mode 100644 tests/expected/function-contract/simple_ensures_pass_no_annotation.expected
 create mode 100644 tests/expected/function-contract/simple_ensures_pass_no_annotation.rs
 create mode 100644 tests/ui/function-contracts/mutating_ensures_error.expected
 create mode 100644 tests/ui/function-contracts/mutating_ensures_error.rs

diff --git a/library/kani/src/internal.rs b/library/kani/src/internal.rs
index a910c333b112..509f2cf51962 100644
--- a/library/kani/src/internal.rs
+++ b/library/kani/src/internal.rs
@@ -90,3 +90,10 @@ pub fn untracked_deref<T>(_: &T) -> T {
 #[doc(hidden)]
 #[rustc_diagnostic_item = "KaniInitContracts"]
 pub fn init_contracts() {}
+
+/// This should only be used within contracts. The intent is to
+/// perform type inference on a closure's argument
+#[doc(hidden)]
+pub fn apply_closure<T, U: Fn(&T) -> bool>(f: U, x: &T) -> bool {
+    f(x)
+}
diff --git a/library/kani_macros/src/sysroot/contracts/shared.rs b/library/kani_macros/src/sysroot/contracts/shared.rs
index 2682c0781661..1ab791d9a117 100644
--- a/library/kani_macros/src/sysroot/contracts/shared.rs
+++ b/library/kani_macros/src/sysroot/contracts/shared.rs
@@ -176,7 +176,7 @@ pub fn build_ensures(data: &ExprClosure) -> (TokenStream2, Expr) {
         .fold(quote!(), |collect, (ident, expr)| quote!(let #ident = #expr; #collect));
 
     let result: Ident = Ident::new(INTERNAL_RESULT_IDENT, Span::call_site());
-    (remembers_stmts, Expr::Verbatim(quote!((#expr)(&#result))))
+    (remembers_stmts, Expr::Verbatim(quote!(kani::internal::apply_closure(#expr, &#result))))
 }
 
 trait OldTrigger {
diff --git a/tests/expected/function-contract/simple_ensures_pass_no_annotation.expected b/tests/expected/function-contract/simple_ensures_pass_no_annotation.expected
new file mode 100644
index 000000000000..0779b6dc88f8
--- /dev/null
+++ b/tests/expected/function-contract/simple_ensures_pass_no_annotation.expected
@@ -0,0 +1,6 @@
+assertion\
+- Status: SUCCESS\
+- Description: "|result| (*result == x) | (*result == y)"\
+in function max
+
+VERIFICATION:- SUCCESSFUL
diff --git a/tests/expected/function-contract/simple_ensures_pass_no_annotation.rs b/tests/expected/function-contract/simple_ensures_pass_no_annotation.rs
new file mode 100644
index 000000000000..a3bf30e1c0f7
--- /dev/null
+++ b/tests/expected/function-contract/simple_ensures_pass_no_annotation.rs
@@ -0,0 +1,14 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+// kani-flags: -Zfunction-contracts
+
+#[kani::ensures(|result| (*result == x) | (*result == y))]
+fn max(x: u32, y: u32) -> u32 {
+    if x > y { x } else { y }
+}
+
+#[kani::proof_for_contract(max)]
+fn max_harness() {
+    let _ = Box::new(9_usize);
+    max(7, 6);
+}
diff --git a/tests/ui/function-contracts/mutating_ensures_error.expected b/tests/ui/function-contracts/mutating_ensures_error.expected
new file mode 100644
index 000000000000..4e9bb3984298
--- /dev/null
+++ b/tests/ui/function-contracts/mutating_ensures_error.expected
@@ -0,0 +1 @@
+cannot assign to `*_x`, as `Fn` closures cannot mutate their captured variables
diff --git a/tests/ui/function-contracts/mutating_ensures_error.rs b/tests/ui/function-contracts/mutating_ensures_error.rs
new file mode 100644
index 000000000000..2fc5f3c8d702
--- /dev/null
+++ b/tests/ui/function-contracts/mutating_ensures_error.rs
@@ -0,0 +1,12 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+// kani-flags: -Zfunction-contracts
+
+#[kani::ensures(|_| {*_x += 1; true})]
+fn unit(_x: &mut u32) {}
+
+#[kani::proof_for_contract(id)]
+fn harness() {
+    let mut x = kani::any();
+    unit(&mut x);
+}