Fix contract expansion for old (#3491)

Fixes the macro expansion for contracts to properly place history
expressions.

## Problem
Before this PR, instantiations of "remembers variables" (i.e., the
variables that save the value before the function executes) were always
put *above* any statements from previous macro expansions. For example,
for this code (from #3359):

```rust
#[kani::requires(val < i32::MAX)]
#[kani::ensures(|result| *result == old(val + 1))]
pub fn next(mut val: i32) -> i32 {
    val + 1
}
```
Kani would first expand the `requires` attribute and insert
`kani::assume(val < i32::MAX)`. The expansion of `ensures` would then
put the remembers variables first, generating this:

```
let remember_kani_internal_1e725538cd5566b8 = val + 1;
kani::assume(val < i32::MAX);
```
which causes an integer overflow because we don't restrict the value of
`val` before adding 1. Instead, we want:

```
kani::assume(val < i32::MAX);
let remember_kani_internal_1e725538cd5566b8 = val + 1;
```

## Solution
The solution is to insert the remembers variables immediately after
preconditions--that way, they respect the preconditions but are still
declared before the function under contract executes. When we're
expanding an `ensures` clause, we iterate through each of the
already-generated statements, find the position where the preconditions
end, then insert the remembers variables there. For instance:

```
kani::assume(x < 100);
kani::assume(y < 10);
kani::assume(x + y < 105);
<-- remembers variables go here -->
let _wrapper_arg = ...
```
---

Resolves #3359

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_macros/src/sysroot/contracts/check.rs

@@ -9,8 +9,8 @@ use std::mem;
 use syn::{parse_quote, Block, Expr, FnArg, Local, LocalInit, Pat, PatIdent, ReturnType, Stmt};
 
 use super::{
-    helpers::*, shared::build_ensures, ContractConditionsData, ContractConditionsHandler,
-    INTERNAL_RESULT_IDENT,
+    helpers::*, shared::build_ensures, ClosureType, ContractConditionsData,
+    ContractConditionsHandler, INTERNAL_RESULT_IDENT,
 };
 
 const WRAPPER_ARG: &str = "_wrapper_arg";
@@ -38,9 +38,14 @@ impl<'a> ContractConditionsHandler<'a> {
                 );
 
                 let return_expr = body_stmts.pop();
+
+                let (assumes, rest_of_body) =
+                    split_for_remembers(&body_stmts[..], ClosureType::Check);
+
                 quote!({
+                    #(#assumes)*
                     #remembers
-                    #(#body_stmts)*
+                    #(#rest_of_body)*
                     #exec_postconditions
                     #return_expr
                 })

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

@@ -4,10 +4,14 @@
 //! Functions that operate third party data structures with no logic that is
 //! specific to Kani and contracts.
 
+use crate::attr_impl::contracts::ClosureType;
 use proc_macro2::{Ident, Span};
 use std::borrow::Cow;
 use syn::spanned::Spanned;
-use syn::{parse_quote, Attribute, Expr, ExprBlock, Local, LocalInit, PatIdent, Stmt};
+use syn::{
+    parse_quote, Attribute, Expr, ExprBlock, ExprCall, ExprPath, Local, LocalInit, PatIdent, Path,
+    Stmt,
+};
 
 /// If an explicit return type was provided it is returned, otherwise `()`.
 pub fn return_type_to_type(return_type: &syn::ReturnType) -> Cow<syn::Type> {
@@ -169,6 +173,53 @@ pub fn chunks_by<'a, T, C: Default + Extend<T>>(
     })
 }
 
+/// Splits `stmts` into (preconditions, rest).
+/// For example, ClosureType::Check assumes preconditions, so given this sequence of statements:
+/// ```ignore
+/// kani::assume(.. precondition_1);
+/// kani::assume(.. precondition_2);
+/// let _wrapper_arg = (ptr as * const _,);
+/// ...
+/// ```
+/// This function would return the two kani::assume statements in the former slice
+/// and the remaining statements in the latter.
+/// The flow for ClosureType::Replace is the same, except preconditions are asserted rather than assumed.
+///
+/// The caller can use the returned tuple to insert remembers statements after `preconditions` and before `rest`.
+/// Inserting the remembers statements after `preconditions` ensures that they are bound by the preconditions.
+/// To understand why this is important, take the following example:
+/// ```ignore
+/// #[kani::requires(x < u32::MAX)]
+/// #[kani::ensures(|result| old(x + 1) == *result)]
+/// fn add_one(x: u32) -> u32 {...}
+/// ```
+/// If the `old(x + 1)` statement didn't respect the precondition's upper bound on `x`, Kani would encounter an integer overflow.
+///
+/// Inserting the remembers statements before `rest` ensures that they are declared before the original function executes,
+/// so that they will store historical, pre-computation values as intended.
+pub fn split_for_remembers(stmts: &[Stmt], closure_type: ClosureType) -> (&[Stmt], &[Stmt]) {
+    let mut pos = 0;
+
+    let check_str = match closure_type {
+        ClosureType::Check => "assume",
+        ClosureType::Replace => "assert",
+    };
+
+    for stmt in stmts {
+        if let Stmt::Expr(Expr::Call(ExprCall { func, .. }), _) = stmt {
+            if let Expr::Path(ExprPath { path: Path { segments, .. }, .. }) = func.as_ref() {
+                let first_two_idents =
+                    segments.iter().take(2).map(|sgmt| sgmt.ident.to_string()).collect::<Vec<_>>();
+
+                if first_two_idents == vec!["kani", check_str] {
+                    pos += 1;
+                }
+            }
+        }
+    }
+    stmts.split_at(pos)
+}
+
 macro_rules! assert_spanned_err {
     ($condition:expr, $span_source:expr, $msg:expr, $($args:expr),+) => {
         if !$condition {

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

@@ -514,6 +514,13 @@ enum ContractConditionsData {
     },
 }
 
+/// Which function are we currently generating?
+#[derive(Copy, Clone, Eq, PartialEq)]
+enum ClosureType {
+    Check,
+    Replace,
+}
+
 impl<'a> ContractConditionsHandler<'a> {
     /// Handle the contract state and return the generated code
     fn dispatch_on(mut self, state: ContractFunctionState) -> TokenStream2 {

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

@@ -11,7 +11,7 @@ use syn::Stmt;
 use super::{
     helpers::*,
     shared::{build_ensures, try_as_result_assign},
-    ContractConditionsData, ContractConditionsHandler, INTERNAL_RESULT_IDENT,
+    ClosureType, ContractConditionsData, ContractConditionsHandler, INTERNAL_RESULT_IDENT,
 };
 
 impl<'a> ContractConditionsHandler<'a> {
@@ -84,9 +84,13 @@ impl<'a> ContractConditionsHandler<'a> {
             ContractConditionsData::Ensures { attr } => {
                 let (remembers, ensures_clause) = build_ensures(attr);
                 let result = Ident::new(INTERNAL_RESULT_IDENT, Span::call_site());
+
+                let (asserts, rest_of_before) = split_for_remembers(before, ClosureType::Replace);
+
                 quote!({
+                    #(#asserts)*
                     #remembers
-                    #(#before)*
+                    #(#rest_of_before)*
                     #(#after)*
                     kani::assume(#ensures_clause);
                     #result

tests/expected/function-contract/history/respects-preconditions/ensures_before_requires.expected

@@ -0,0 +1,5 @@
+next\
+     - Status: SUCCESS\
+	- Description: "attempt to add with overflow"
+
+VERIFICATION:- SUCCESSFUL

tests/expected/function-contract/history/respects-preconditions/ensures_before_requires.rs

@@ -0,0 +1,18 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+// kani-flags: -Zfunction-contracts
+
+// Demonstrate that when the ensures contract is before the requires contract,
+// the history expression respects the upper bound on x, so x + 1 doesn't overflow
+// This example is taken from https://github.com/model-checking/kani/issues/3359
+
+#[kani::ensures(|result| *result == old(val + 1))]
+#[kani::requires(val < i32::MAX)]
+pub fn next(val: i32) -> i32 {
+    val + 1
+}
+
+#[kani::proof_for_contract(next)]
+pub fn check_next() {
+    let _ = next(kani::any());
+}

tests/expected/function-contract/history/respects-preconditions/modifies.expected

@@ -0,0 +1,5 @@
+modify\
+	 - Status: SUCCESS\
+	 - Description: "attempt to add with overflow"
+
+VERIFICATION:- SUCCESSFUL

tests/expected/function-contract/history/respects-preconditions/modifies.rs

@@ -0,0 +1,29 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+// kani-flags: -Zfunction-contracts
+
+// Demonstrate the the history expression respects preconditions
+// with multiple interleaved preconditions, modifies contracts, and history expressions
+
+#[derive(kani::Arbitrary)]
+struct Point<X, Y> {
+    x: X,
+    y: Y,
+}
+
+#[kani::requires(ptr.x < 100)]
+#[kani::ensures(|result| old(ptr.x + 1) == ptr.x)]
+#[kani::modifies(&mut ptr.x)]
+#[kani::ensures(|result| old(ptr.y - 1) == ptr.y)]
+#[kani::modifies(&mut ptr.y)]
+#[kani::requires(ptr.y > 0)]
+fn modify(ptr: &mut Point<u32, u32>) {
+    ptr.x += 1;
+    ptr.y -= 1;
+}
+
+#[kani::proof_for_contract(modify)]
+fn main() {
+    let mut p: Point<u32, u32> = kani::any();
+    modify(&mut p);
+}

tests/expected/function-contract/history/respects-preconditions/requires_before_ensures.expected

@@ -0,0 +1,5 @@
+next\
+     - Status: SUCCESS\
+	- Description: "attempt to add with overflow"
+
+VERIFICATION:- SUCCESSFUL

tests/expected/function-contract/history/respects-preconditions/requires_before_ensures.rs

@@ -0,0 +1,17 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+// kani-flags: -Zfunction-contracts
+
+// Demonstrate that when the requires contract is before the ensures contract, the history expression respects the upper bound on x, so x + 1 doesn't overflow
+// This example is taken from https://github.com/model-checking/kani/issues/3359
+
+#[kani::requires(val < i32::MAX)]
+#[kani::ensures(|result| *result == old(val + 1))]
+pub fn next(val: i32) -> i32 {
+    val + 1
+}
+
+#[kani::proof_for_contract(next)]
+pub fn check_next() {
+    let _ = next(kani::any());
+}