Fix the bug: Loop contracts are not composable with function contracts (#3979)

This PR fixes the bug of loop contracts are not composable with function
contracts.
Previously, using loop-contract inside a function with contract may
result in unwinding the loop instead of contracting it.

Resolves #3910

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: Carolyn Zech <cmzech@amazon.com>

Author: thanhnguyen-aws

docs/src/reference/experimental/loop-contracts.md

@@ -139,13 +139,36 @@ In proof path 1, we prove properties inside the loop and at last check that the
 In proof path 2, we prove properties after leaving the loop. As we leave the loop only when the loop guard is violated, the post condition of the loop can be expressed as
 `!guard && inv`, which is `x <= 1 && x >= 1` in the example. The postcondition implies `x == 1`—the property we want to prove at the end of `simple_loop_with_loop_contracts`.
 
+## Loop contracts inside functions with contracts 
+Kani supports using loop contracts together with function contracts, as demonstrated in the following example:
+``` Rust
+#![feature(proc_macro_hygiene)]
+#![feature(stmt_expr_attributes)]
+
+#[kani::requires(i>=2)]
+#[kani::ensures(|ret| *ret == 2)]
+pub fn has_loop(mut i: u16) -> u16 {
+    #[kani::loop_invariant(i>=2)]
+    while i > 2 {
+        i = i - 1
+    }
+    i
+}
+
+#[kani::proof_for_contract(has_loop)]
+fn contract_proof() {
+    let i: u16 = kani::any();
+    let j = has_loop(i);
+}
+```
 
+When loop contracts and function contracts are both enabled (by flags `-Z loop-contracts -Z function-contracts`), 
+Kani automatically contracts (instead of unwinds) all loops in the functions that we want to prove contracts for.
 ## Limitations
 
 Loop contracts comes with the following limitations.
 
-1. Only `while` loops are supported. The other three kinds of loops are not supported: [`loop` loops](https://doc.rust-lang.org/reference/expressions/loop-expr.html#infinite-loops)
-   , [`while let` loops](https://doc.rust-lang.org/reference/expressions/loop-expr.html#predicate-pattern-loops), and [`for` loops](https://doc.rust-lang.org/reference/expressions/loop-expr.html#iterator-loops).
+1. `while` loops and `loop` loops are supported. The other kinds of loops are not supported: [`while let` loops](https://doc.rust-lang.org/reference/expressions/loop-expr.html#predicate-pattern-loops), and [`for` loops](https://doc.rust-lang.org/reference/expressions/loop-expr.html#iterator-loops).
 2. Kani infers *loop modifies* with alias analysis. Loop modifies are those variables we assume to be arbitrary in the inductive hypothesis, and should cover all memory locations that are written to during 
    the execution of the loops. A proof will fail if the inferred loop modifies misses some targets written in the loops.
    We observed this happens when some fields of structs are modified by some other functions called in the loops.

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

@@ -103,39 +103,10 @@ impl TransformPass for LoopContractPass {
                     let run = Instance::resolve(self.run_contract_fn.unwrap(), args).unwrap();
                     (true, run.body().unwrap())
                 } else {
-                    let mut new_body = MutableBody::from(body);
-                    let mut contain_loop_contracts: bool = false;
-
-                    // Visit basic blocks in control flow order (BFS).
-                    let mut visited: HashSet<BasicBlockIdx> = HashSet::new();
-                    let mut queue: VecDeque<BasicBlockIdx> = VecDeque::new();
-                    // Visit blocks in loops only when there is no blocks in queue.
-                    let mut loop_queue: VecDeque<BasicBlockIdx> = VecDeque::new();
-                    queue.push_back(0);
-
-                    while let Some(bb_idx) = queue.pop_front().or_else(|| loop_queue.pop_front()) {
-                        visited.insert(bb_idx);
-
-                        let terminator = new_body.blocks()[bb_idx].terminator.clone();
-
-                        let is_loop_head = self.transform_bb(tcx, &mut new_body, bb_idx);
-                        contain_loop_contracts |= is_loop_head;
-
-                        // Add successors of the current basic blocks to
-                        // the visiting queue.
-                        for to_visit in terminator.successors() {
-                            if !visited.contains(&to_visit) {
-                                if is_loop_head {
-                                    loop_queue.push_back(to_visit);
-                                } else {
-                                    queue.push_back(to_visit)
-                                };
-                            }
-                        }
-                    }
-                    (contain_loop_contracts, new_body.into())
+                    self.transform_body_with_loop(tcx, body)
                 }
             }
+            RigidTy::Closure(_, _) => self.transform_body_with_loop(tcx, body),
             _ => {
                 /* static variables case */
                 (false, body)
@@ -194,6 +165,43 @@ impl LoopContractPass {
         ))
     }
 
+    /// This function transform the function body as described in fn transform.
+    /// It is the core of fn transform, and is separated just to avoid code repetition.
+    fn transform_body_with_loop(&mut self, tcx: TyCtxt, body: Body) -> (bool, Body) {
+        let mut new_body = MutableBody::from(body);
+        let mut contain_loop_contracts: bool = false;
+
+        // Visit basic blocks in control flow order (BFS).
+        let mut visited: HashSet<BasicBlockIdx> = HashSet::new();
+        let mut queue: VecDeque<BasicBlockIdx> = VecDeque::new();
+        // Visit blocks in loops only when there is no blocks in queue.
+        let mut loop_queue: VecDeque<BasicBlockIdx> = VecDeque::new();
+        queue.push_back(0);
+
+        while let Some(bb_idx) = queue.pop_front().or_else(|| loop_queue.pop_front()) {
+            visited.insert(bb_idx);
+
+            let terminator = new_body.blocks()[bb_idx].terminator.clone();
+
+            let is_loop_head = self.transform_bb(tcx, &mut new_body, bb_idx);
+            contain_loop_contracts |= is_loop_head;
+
+            // Add successors of the current basic blocks to
+            // the visiting queue.
+            for to_visit in terminator.successors() {
+                if !visited.contains(&to_visit) {
+                    if is_loop_head {
+                        loop_queue.push_back(to_visit);
+                    } else {
+                        queue.push_back(to_visit)
+                    };
+                }
+            }
+        }
+
+        (contain_loop_contracts, new_body.into())
+    }
+
     /// Transform loops with contracts from
     ///    ```ignore
     ///    bb_idx: {
@@ -316,7 +324,7 @@ impl LoopContractPass {
                 for stmt in &new_body.blocks()[bb_idx].statements {
                     if let StatementKind::Assign(place, rvalue) = &stmt.kind {
                         match rvalue {
-                            Rvalue::Ref(_,_,rplace) => {
+                            Rvalue::Ref(_,_,rplace) | Rvalue::CopyForDeref(rplace) => {
                                 if supported_vars.contains(&rplace.local) {
                                     supported_vars.push(place.local);
                                 } }

tests/expected/loop-contract/contract_proof_function_with_loop.expected

@@ -0,0 +1,13 @@
+has_loop::{closure#2}::{closure#1}.loop_invariant_step.1\
+	 - Status: SUCCESS\
+	 - Description: "Check invariant after step for loop has_loop::{closure#2}::{closure#1}.0"\
+in function has_loop::{closure#2}::{closure#1}
+
+
+
+has_loop::{closure#2}::{closure#1}.precondition_instance.1\
+	 - Status: SUCCESS\
+	 - Description: "free argument must be NULL or valid pointer"\
+in function has_loop::{closure#2}::{closure#1}
+
+VERIFICATION:- SUCCESSFUL

tests/expected/loop-contract/contract_proof_function_with_loop.rs

@@ -0,0 +1,25 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+// kani-flags: -Z loop-contracts -Z function-contracts
+
+//! Test that we can prove a function contract and a loop contract in tandem.
+
+#![feature(proc_macro_hygiene)]
+#![feature(stmt_expr_attributes)]
+
+#[kani::requires(i>=2)]
+#[kani::ensures(|ret| *ret == 2)]
+pub fn has_loop(mut i: u16) -> u16 {
+    #[kani::loop_invariant(i>=2)]
+    while i > 2 {
+        i = i - 1
+    }
+    i
+}
+
+#[kani::proof_for_contract(has_loop)]
+fn contract_proof() {
+    let i: u16 = kani::any();
+    let j = has_loop(i);
+}

tests/expected/loop-contract/function_call_with_loop.expected

@@ -0,0 +1,11 @@
+has_loop::{closure#3}::{closure#0}.loop_invariant_base.1\
+	 - Status: SUCCESS\
+	 - Description: "Check invariant before entry for loop has_loop::{closure#3}::{closure#0}.0"\
+in function has_loop::{closure#3}::{closure#0}
+
+has_loop::{closure#3}::{closure#0}.precondition_instance.5\
+	 - Status: SUCCESS\
+	 - Description: "free called for new[] object"\
+in function has_loop::{closure#3}::{closure#0}
+
+Failed Checks: i>=2

tests/expected/loop-contract/function_call_with_loop.rs

@@ -0,0 +1,25 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+// kani-flags: -Z loop-contracts -Z function-contracts
+
+//! Calling a function that contains loops and test that using a #[kani::proof] harness fails because the function's precondition gets asserted.
+
+#![feature(proc_macro_hygiene)]
+#![feature(stmt_expr_attributes)]
+
+#[kani::requires(i>=2)]
+#[kani::ensures(|ret| *ret == 2)]
+pub fn has_loop(mut i: u16) -> u16 {
+    #[kani::loop_invariant(i>=2)]
+    while i > 2 {
+        i = i - 1
+    }
+    i
+}
+
+#[kani::proof]
+fn call_has_loop() {
+    let i: u16 = kani::any();
+    let j = has_loop(i);
+}

tests/expected/loop-contract/function_with_loop_no_assertion.expected

@@ -0,0 +1,8 @@
+has_loop.loop_invariant_base.1\
+	 - Status: SUCCESS\
+	 - Description: "Check invariant before entry for loop has_loop.0"\
+in function has_loop
+
+
+
+VERIFICATION:- SUCCESSFUL

tests/expected/loop-contract/function_with_loop_no_assertion.rs

@@ -0,0 +1,25 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+// kani-flags: -Z loop-contracts -Z function-contracts --no-assert-contracts
+
+//Call a function with loop without checking the contract.
+
+#![feature(proc_macro_hygiene)]
+#![feature(stmt_expr_attributes)]
+
+#[kani::requires(i>=2)]
+#[kani::ensures(|ret| *ret == 2)]
+pub fn has_loop(mut i: u16) -> u16 {
+    #[kani::loop_invariant(i>=2)]
+    while i > 2 {
+        i = i - 1
+    }
+    i
+}
+
+#[kani::proof]
+fn contract_proof() {
+    let i: u16 = kani::any();
+    let j = has_loop(i);
+}