annotate all ensures closures

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

@@ -34,7 +34,8 @@ impl<'a> ContractConditionsHandler<'a> {
                 )
             }
             ContractConditionsData::Ensures { attr } => {
-                let (remembers, ensures_clause) = build_ensures(attr);
+                let (remembers, ensures_clause) =
+                    build_ensures(attr, return_type_to_type(&self.annotated_fn.sig.output));
 
                 // The code that enforces the postconditions and cleans up the shallow
                 // argument copies (with `mem::forget`).

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

@@ -80,7 +80,8 @@ impl<'a> ContractConditionsHandler<'a> {
                 )
             }
             ContractConditionsData::Ensures { attr } => {
-                let (remembers, ensures_clause) = build_ensures(attr);
+                let (remembers, ensures_clause) =
+                    build_ensures(attr, return_type_to_type(&self.annotated_fn.sig.output));
                 let result = Ident::new(INTERNAL_RESULT_IDENT, Span::call_site());
                 quote!(
                     #remembers

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

@@ -11,9 +11,11 @@ use std::collections::HashMap;
 
 use proc_macro2::{Ident, Span, TokenStream as TokenStream2};
 use quote::{quote, ToTokens};
+use std::borrow::{Borrow, Cow};
 use std::hash::{DefaultHasher, Hash, Hasher};
 use syn::{
-    spanned::Spanned, visit_mut::VisitMut, Attribute, Expr, ExprCall, ExprClosure, ExprPath, Path,
+    spanned::Spanned, visit_mut::VisitMut, Attribute, Expr, ExprCall, ExprClosure, ExprPath,
+    PatType, Path, Token, Type, TypeReference,
 };
 
 use super::{ContractConditionsHandler, ContractFunctionState, INTERNAL_RESULT_IDENT};
@@ -165,7 +167,7 @@ pub fn try_as_result_assign_mut(stmt: &mut syn::Stmt) -> Option<&mut syn::LocalI
 /// of statements that instantiate these expressions as `let remember_kani_internal_x = old_expr;`
 /// where x is a unique hash. This is returned as the first return parameter. The second
 /// return parameter is the expression formed by passing in the result variable into the input closure.
-pub fn build_ensures(data: &ExprClosure) -> (TokenStream2, Expr) {
+pub fn build_ensures<'a>(data: &ExprClosure, return_type: Cow<'a, Type>) -> (TokenStream2, Expr) {
     let mut remembers_exprs = HashMap::new();
     let mut vis = OldVisitor { t: OldLifter::new(), remembers_exprs: &mut remembers_exprs };
     let mut expr = &mut data.clone();
@@ -175,6 +177,29 @@ pub fn build_ensures(data: &ExprClosure) -> (TokenStream2, Expr) {
         .iter()
         .fold(quote!(), |collect, (ident, expr)| quote!(let #ident = #expr; #collect));
 
+    *expr
+        .inputs
+        .first_mut()
+        .expect("Ensures closure should have the output to the function as an argument") =
+        syn::Pat::Type(PatType {
+            attrs: vec![],
+            pat: Box::new(
+                expr.inputs
+                    .first()
+                    .expect("Ensures closure should have the output to the function as an argument")
+                    .clone(),
+            ),
+            colon_token: Token![:](Span::call_site()),
+            ty: Box::new(Type::Reference(TypeReference {
+                and_token: Token![&](Span::call_site()),
+                lifetime: None,
+                mutability: None,
+                elem: Box::new(
+                    <Cow<'a, Type> as Borrow<Type>>::borrow(&return_type.to_owned()).clone(),
+                ),
+            })),
+        });
+
     let result: Ident = Ident::new(INTERNAL_RESULT_IDENT, Span::call_site());
     (remembers_stmts, Expr::Verbatim(quote!((#expr)(&#result))))
 }

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

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);
+}