Fix the bug: Loop contracts are not composable with function contracts

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

@@ -103,40 +103,13 @@ 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),
             _ => {
+                println!("Instance defid: {:?}", instance.def.name());
+                println!("Instance ty: {:?}", instance.ty().kind().rigid().unwrap());
                 /* static variables case */
                 (false, body)
             }
@@ -194,6 +167,41 @@ impl LoopContractPass {
         ))
     }
 
+    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: {
@@ -310,7 +318,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
+
+//! Proving contract for function with loop
+
+#![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
+
+#![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);
+}