Change ensures into closures (rust-lang#3207)

This change now separates the front facing "result" name and the
internal facing "result_kani_internal" ident where the user can specify
with the keyword "result" but then the system replaces this with the
internal representation.

If the user chooses to use a different variable name than result, this
now supports the syntax of
`#[kani::ensures(|result_var| expr)]`
where result_var can be any arbitrary name.

For example, the following test now works:
```
#[kani::ensures(|banana : &u32| *banana == a.wrapping_add(b))]
fn simple_addition(a: u32, b: u32) -> u32 {
    return a.wrapping_add(b);
}
```

In addition, the string "result_kani_internal" is now a global constant
and can be updated in a single place if needed.

Resolves rust-lang#2597 since the user can specify the variable name they want
within the ensures binding

An important change is that the result is now a pass by reference
instead, so as in the example an `&u32` instead of `u32`

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: Matias Scharager <[email protected]>

Author: pi314mm

library/kani/src/contracts.rs

@@ -51,14 +51,14 @@
 //! approximation of the result of division for instance could be this:
 //!
 //! ```
-//! #[ensures(result <= dividend)]
+//! #[ensures(|result : &u32| *result <= dividend)]
 //! ```
 //!
 //! This is called a postcondition and it also has access to the arguments and
 //! is expressed in regular Rust code. The same restrictions apply as did for
-//! [`requires`][macro@requires]. In addition to the arguments the postcondition
-//! also has access to the value returned from the function in a variable called
-//! `result`.
+//! [`requires`][macro@requires]. In addition to the postcondition is expressed
+//! as a closure where the value returned from the function is passed to this
+//! closure by reference.
 //!
 //! You may combine as many [`requires`][macro@requires] and
 //! [`ensures`][macro@ensures] attributes on a single function as you please.
@@ -67,7 +67,7 @@
 //!
 //! ```
 //! #[kani::requires(divisor != 0)]
-//! #[kani::ensures(result <= dividend)]
+//! #[kani::ensures(|result : &u32| *result <= dividend)]
 //! fn my_div(dividend: u32, divisor: u32) -> u32 {
 //!   dividend / divisor
 //! }

library/kani_macros/src/lib.rs

@@ -131,11 +131,11 @@ pub fn requires(attr: TokenStream, item: TokenStream) -> TokenStream {
 /// This is part of the function contract API, for more general information see
 /// the [module-level documentation](../kani/contracts/index.html).
 ///
-/// The contents of the attribute is a condition over the input values to the
-/// annotated function *and* its return value, accessible as a variable called
-/// `result`. All Rust syntax is supported, even calling other functions, but
-/// the computations must be side effect free, e.g. it cannot perform I/O or use
-/// mutable memory.
+/// The contents of the attribute is a closure that captures the input values to
+/// the annotated function and the input to the function is the return value of
+/// the function passed by reference. All Rust syntax is supported, even calling
+/// other functions, but the computations must be side effect free, e.g. it
+/// cannot perform I/O or use mutable memory.
 ///
 /// Kani requires each function that uses a contract (this attribute or
 /// [`requires`][macro@requires]) to have at least one designated

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

@@ -4,11 +4,14 @@
 //! Special way we handle the first time we encounter a contract attribute on a
 //! function.
 
-use proc_macro2::Span;
+use proc_macro2::{Ident, Span};
 use quote::quote;
 use syn::ItemFn;
 
-use super::{helpers::*, shared::identifier_for_generated_function, ContractConditionsHandler};
+use super::{
+    helpers::*, shared::identifier_for_generated_function, ContractConditionsHandler,
+    INTERNAL_RESULT_IDENT,
+};
 
 impl<'a> ContractConditionsHandler<'a> {
     /// The complex case. We are the first time a contract is handled on this function, so
@@ -80,6 +83,7 @@ impl<'a> ContractConditionsHandler<'a> {
             (quote!(#check_fn_name), quote!(#replace_fn_name))
         };
 
+        let result = Ident::new(INTERNAL_RESULT_IDENT, Span::call_site());
         self.output.extend(quote!(
             #[allow(dead_code, unused_variables)]
             #[kanitool::is_contract_generated(recursion_wrapper)]
@@ -89,9 +93,9 @@ impl<'a> ContractConditionsHandler<'a> {
                     #call_replace(#(#args),*)
                 } else {
                     unsafe { REENTRY = true };
-                    let result = #call_check(#(#also_args),*);
+                    let #result = #call_check(#(#also_args),*);
                     unsafe { REENTRY = false };
-                    result
+                    #result
                 }
             }
         ));

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

@@ -10,7 +10,7 @@ use syn::{Expr, FnArg, ItemFn, Token};
 use super::{
     helpers::*,
     shared::{make_unsafe_argument_copies, try_as_result_assign_mut},
-    ContractConditionsData, ContractConditionsHandler,
+    ContractConditionsData, ContractConditionsHandler, INTERNAL_RESULT_IDENT,
 };
 
 const WRAPPER_ARG_PREFIX: &str = "_wrapper_arg_";
@@ -46,14 +46,15 @@ impl<'a> ContractConditionsHandler<'a> {
                 assert!(matches!(
                     inner.pop(),
                     Some(syn::Stmt::Expr(syn::Expr::Path(pexpr), None))
-                        if pexpr.path.get_ident().map_or(false, |id| id == "result")
+                        if pexpr.path.get_ident().map_or(false, |id| id == INTERNAL_RESULT_IDENT)
                 ));
 
+                let result = Ident::new(INTERNAL_RESULT_IDENT, Span::call_site());
                 quote!(
                     #arg_copies
                     #(#inner)*
                     #exec_postconditions
-                    result
+                    #result
                 )
             }
             ContractConditionsData::Modifies { attr } => {
@@ -95,9 +96,10 @@ impl<'a> ContractConditionsHandler<'a> {
             } else {
                 quote!(#wrapper_name)
             };
+            let result = Ident::new(INTERNAL_RESULT_IDENT, Span::call_site());
             syn::parse_quote!(
-                let result : #return_type = #wrapper_call(#(#args),*);
-                result
+                let #result : #return_type = #wrapper_call(#(#args),*);
+                #result
             )
         } else {
             self.annotated_fn.block.stmts.clone()

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

@@ -7,12 +7,13 @@ use std::collections::{HashMap, HashSet};
 
 use proc_macro::{Diagnostic, TokenStream};
 use proc_macro2::{Ident, Span, TokenStream as TokenStream2};
-use syn::{spanned::Spanned, visit::Visit, visit_mut::VisitMut, Expr, ItemFn, Signature};
+use quote::quote;
+use syn::{spanned::Spanned, visit::Visit, visit_mut::VisitMut, Expr, ExprClosure, ItemFn};
 
 use super::{
     helpers::{chunks_by, is_token_stream_2_comma, matches_path},
     ContractConditionsData, ContractConditionsHandler, ContractConditionsType,
-    ContractFunctionState,
+    ContractFunctionState, INTERNAL_RESULT_IDENT,
 };
 
 impl<'a> TryFrom<&'a syn::Attribute> for ContractFunctionState {
@@ -81,7 +82,11 @@ impl<'a> ContractConditionsHandler<'a> {
                 ContractConditionsData::Requires { attr: syn::parse(attr)? }
             }
             ContractConditionsType::Ensures => {
-                ContractConditionsData::new_ensures(&annotated_fn.sig, syn::parse(attr)?)
+                let mut data: ExprClosure = syn::parse(attr)?;
+                let argument_names = rename_argument_occurrences(&annotated_fn.sig, &mut data);
+                let result: Ident = Ident::new(INTERNAL_RESULT_IDENT, Span::call_site());
+                let app: Expr = Expr::Verbatim(quote!((#data)(&#result)));
+                ContractConditionsData::Ensures { argument_names, attr: app }
             }
             ContractConditionsType::Modifies => {
                 ContractConditionsData::new_modifies(attr, &mut output)
@@ -92,16 +97,6 @@ impl<'a> ContractConditionsHandler<'a> {
     }
 }
 impl ContractConditionsData {
-    /// Constructs a [`Self::Ensures`] from the signature of the decorated
-    /// function and the contents of the decorating attribute.
-    ///
-    /// Renames the [`Ident`]s used in `attr` and stores the translation map in
-    /// `argument_names`.
-    fn new_ensures(sig: &Signature, mut attr: Expr) -> Self {
-        let argument_names = rename_argument_occurrences(sig, &mut attr);
-        ContractConditionsData::Ensures { argument_names, attr }
-    }
-
     /// Constructs a [`Self::Modifies`] from the contents of the decorating attribute.
     ///
     /// Responsible for parsing the attribute.
@@ -129,7 +124,10 @@ impl ContractConditionsData {
 /// - Creates new names for them;
 /// - Replaces all occurrences of those idents in `attrs` with the new names and;
 /// - Returns the mapping of old names to new names.
-fn rename_argument_occurrences(sig: &syn::Signature, attr: &mut Expr) -> HashMap<Ident, Ident> {
+fn rename_argument_occurrences(
+    sig: &syn::Signature,
+    attr: &mut ExprClosure,
+) -> HashMap<Ident, Ident> {
     let mut arg_ident_collector = ArgumentIdentCollector::new();
     arg_ident_collector.visit_signature(&sig);
 
@@ -144,7 +142,7 @@ fn rename_argument_occurrences(sig: &syn::Signature, attr: &mut Expr) -> HashMap
         .collect::<HashMap<_, _>>();
 
     let mut ident_rewriter = Renamer(&arg_idents);
-    ident_rewriter.visit_expr_mut(attr);
+    ident_rewriter.visit_expr_closure_mut(attr);
     arg_idents
 }
 

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

@@ -159,9 +159,9 @@
 //!         call_replace(fn args...)
 //!     } else {
 //!         unsafe { reentry = true };
-//!         let result = call_check(fn args...);
+//!         let result_kani_internal = call_check(fn args...);
 //!         unsafe { reentry = false };
-//!         result
+//!         result_kani_internal
 //!     }
 //! }
 //! ```
@@ -173,7 +173,7 @@
 //!
 //! ```
 //! #[kani::requires(divisor != 0)]
-//! #[kani::ensures(result <= dividend)]
+//! #[kani::ensures(|result : &u32| *result <= dividend)]
 //! fn div(dividend: u32, divisor: u32) -> u32 {
 //!     dividend / divisor
 //! }
@@ -186,31 +186,35 @@
 //! #[kanitool::replaced_with = "div_replace_965916"]
 //! fn div(dividend: u32, divisor: u32) -> u32 { dividend / divisor }
 //!
-//! #[allow(dead_code)]
-//! #[allow(unused_variables)]
-//! #[kanitool::is_contract_generated(check)]
-//! fn div_check_965916(dividend: u32, divisor: u32) -> u32 {
-//!     let dividend_renamed = kani::internal::untracked_deref(&dividend);
-//!     let divisor_renamed = kani::internal::untracked_deref(&divisor);
-//!     let result = { kani::assume(divisor != 0); { dividend / divisor } };
-//!     kani::assert(result <= dividend_renamed, "result <= dividend");
+//! #[allow(dead_code, unused_variables)]
+//! #[kanitool :: is_contract_generated(check)] fn
+//! div_check_b97df2(dividend : u32, divisor : u32) -> u32
+//! {
+//!     let dividend_renamed = kani::internal::untracked_deref(& dividend);
+//!     let divisor_renamed = kani::internal::untracked_deref(& divisor);
+//!     kani::assume(divisor != 0);
+//!     let result_kani_internal : u32 = div_wrapper_b97df2(dividend, divisor);
+//!     kani::assert(
+//!     (| result : & u32 | *result <= dividend_renamed)(& result_kani_internal),
+//!     stringify!(|result : &u32| *result <= dividend));
 //!     std::mem::forget(dividend_renamed);
 //!     std::mem::forget(divisor_renamed);
-//!     result
+//!     result_kani_internal
 //! }
 //!
-//! #[allow(dead_code)]
-//! #[allow(unused_variables)]
-//! #[kanitool::is_contract_generated(replace)]
-//! fn div_replace_965916(dividend: u32, divisor: u32) -> u32 {
-//!     kani::assert(divisor != 0, "divisor != 0");
-//!     let dividend_renamed = kani::internal::untracked_deref(&dividend);
-//!     let divisor_renamed = kani::internal::untracked_deref(&divisor);
-//!     let result = kani::any();
-//!     kani::assume(result <= dividend_renamed, "result <= dividend");
-//!     std::mem::forget(dividend_renamed);
+//! #[allow(dead_code, unused_variables)]
+//! #[kanitool :: is_contract_generated(replace)] fn
+//! div_replace_b97df2(dividend : u32, divisor : u32) -> u32
+//! {
+//!     let divisor_renamed = kani::internal::untracked_deref(& divisor);
+//!     let dividend_renamed = kani::internal::untracked_deref(& dividend);
+//!     kani::assert(divisor != 0, stringify! (divisor != 0));
+//!     let result_kani_internal : u32 = kani::any_modifies();
+//!     kani::assume(
+//!     (|result : & u32| *result <= dividend_renamed)(&result_kani_internal));
 //!     std::mem::forget(divisor_renamed);
-//!     result
+//!     std::mem::forget(dividend_renamed);
+//!     result_kani_internal
 //! }
 //!
 //! #[allow(dead_code)]
@@ -220,12 +224,12 @@
 //!     static mut REENTRY: bool = false;
 //!
 //!     if unsafe { REENTRY } {
-//!         div_replace_965916(dividend, divisor)
+//!         div_replace_b97df2(dividend, divisor)
 //!     } else {
 //!         unsafe { reentry = true };
-//!         let result = div_check_965916(dividend, divisor);
+//!         let result_kani_internal = div_check_b97df2(dividend, divisor);
 //!         unsafe { reentry = false };
-//!         result
+//!         result_kani_internal
 //!     }
 //! }
 //! ```
@@ -243,6 +247,8 @@ mod initialize;
 mod replace;
 mod shared;
 
+const INTERNAL_RESULT_IDENT: &str = "result_kani_internal";
+
 pub fn requires(attr: TokenStream, item: TokenStream) -> TokenStream {
     contract_main(attr, item, ContractConditionsType::Requires)
 }

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

@@ -3,13 +3,13 @@
 
 //! Logic used for generating the code that replaces a function with its contract.
 
-use proc_macro2::{Ident, TokenStream as TokenStream2};
+use proc_macro2::{Ident, Span, TokenStream as TokenStream2};
 use quote::quote;
 
 use super::{
     helpers::*,
     shared::{make_unsafe_argument_copies, try_as_result_assign},
-    ContractConditionsData, ContractConditionsHandler,
+    ContractConditionsData, ContractConditionsHandler, INTERNAL_RESULT_IDENT,
 };
 
 impl<'a> ContractConditionsHandler<'a> {
@@ -39,7 +39,8 @@ impl<'a> ContractConditionsHandler<'a> {
     fn ensure_bootstrapped_replace_body(&self) -> (Vec<syn::Stmt>, Vec<syn::Stmt>) {
         if self.is_first_emit() {
             let return_type = return_type_to_type(&self.annotated_fn.sig.output);
-            (vec![syn::parse_quote!(let result : #return_type = kani::any_modifies();)], vec![])
+            let result = Ident::new(INTERNAL_RESULT_IDENT, Span::call_site());
+            (vec![syn::parse_quote!(let #result : #return_type = kani::any_modifies();)], vec![])
         } else {
             let stmts = &self.annotated_fn.block.stmts;
             let idx = stmts
@@ -70,30 +71,33 @@ impl<'a> ContractConditionsHandler<'a> {
         match &self.condition_type {
             ContractConditionsData::Requires { attr } => {
                 let Self { attr_copy, .. } = self;
+                let result = Ident::new(INTERNAL_RESULT_IDENT, Span::call_site());
                 quote!(
                     kani::assert(#attr, stringify!(#attr_copy));
                     #(#before)*
                     #(#after)*
-                    result
+                    #result
                 )
             }
             ContractConditionsData::Ensures { attr, argument_names } => {
                 let (arg_copies, copy_clean) = make_unsafe_argument_copies(&argument_names);
+                let result = Ident::new(INTERNAL_RESULT_IDENT, Span::call_site());
                 quote!(
                     #arg_copies
                     #(#before)*
                     #(#after)*
                     kani::assume(#attr);
                     #copy_clean
-                    result
+                    #result
                 )
             }
             ContractConditionsData::Modifies { attr } => {
+                let result = Ident::new(INTERNAL_RESULT_IDENT, Span::call_site());
                 quote!(
                     #(#before)*
                     #(*unsafe { kani::internal::Pointer::assignable(#attr) } = kani::any_modifies();)*
                     #(#after)*
-                    result
+                    #result
                 )
             }
         }
@@ -126,7 +130,7 @@ impl<'a> ContractConditionsHandler<'a> {
     }
 }
 
-/// Is this statement `let result : <...> = kani::any_modifies();`.
+/// Is this statement `let result_kani_internal : <...> = kani::any_modifies();`.
 fn is_replace_return_havoc(stmt: &syn::Stmt) -> bool {
     let Some(syn::LocalInit { diverge: None, expr: e, .. }) = try_as_result_assign(stmt) else {
         return false;

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

@@ -13,7 +13,7 @@ use proc_macro2::{Ident, Span, TokenStream as TokenStream2};
 use quote::{quote, ToTokens};
 use syn::Attribute;
 
-use super::{ContractConditionsHandler, ContractFunctionState};
+use super::{ContractConditionsHandler, ContractFunctionState, INTERNAL_RESULT_IDENT};
 
 impl ContractFunctionState {
     /// Do we need to emit the `is_contract_generated` tag attribute on the
@@ -114,7 +114,7 @@ macro_rules! try_as_result_assign_pat {
                     ident: result_ident,
                     subpat: None,
                     ..
-                }) if result_ident == "result"
+                }) if result_ident == INTERNAL_RESULT_IDENT
             ) => init.$convert(),
             _ => None,
         }