From 68415fff960977f8080abbd2fe2f5b1cc9a5c53e Mon Sep 17 00:00:00 2001
From: Qinheping Hu <qinhh@amazon.com>
Date: Fri, 25 Oct 2024 14:16:00 -0500
Subject: [PATCH] Fix loop contracts transformation when loops in branching
 (#3640)

In loop contracts transformation, we want to visit non-loop blocks
before visiting loop blocks. So we pop `loop_queue` only when `queue` is
empty. We should use `or_else` instead of `or` which always pop the
front of `loop_queue`.

By submitting this pull request, I confirm that my contribution is made
under the terms of the Apache 2.0 and MIT licenses.
---
 .../src/kani_middle/transform/loop_contracts.rs        | 10 ++++++----
 tests/expected/loop-contract/multiple_loops.rs         | 10 ++++++----
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/kani-compiler/src/kani_middle/transform/loop_contracts.rs b/kani-compiler/src/kani_middle/transform/loop_contracts.rs
index 8038ef784970..887c07ed4792 100644
--- a/kani-compiler/src/kani_middle/transform/loop_contracts.rs
+++ b/kani-compiler/src/kani_middle/transform/loop_contracts.rs
@@ -114,7 +114,7 @@ impl TransformPass for LoopContractPass {
                     let mut loop_queue: VecDeque<BasicBlockIdx> = VecDeque::new();
                     queue.push_back(0);
 
-                    while let Some(bb_idx) = queue.pop_front().or(loop_queue.pop_front()) {
+                    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();
@@ -126,9 +126,11 @@ impl TransformPass for LoopContractPass {
                         // the visiting queue.
                         for to_visit in terminator.successors() {
                             if !visited.contains(&to_visit) {
-                                let target_queue =
-                                    if is_loop_head { &mut loop_queue } else { &mut queue };
-                                target_queue.push_back(to_visit);
+                                if is_loop_head {
+                                    loop_queue.push_back(to_visit);
+                                } else {
+                                    queue.push_back(to_visit)
+                                };
                             }
                         }
                     }
diff --git a/tests/expected/loop-contract/multiple_loops.rs b/tests/expected/loop-contract/multiple_loops.rs
index e56af83dca0f..8baf1f251b4c 100644
--- a/tests/expected/loop-contract/multiple_loops.rs
+++ b/tests/expected/loop-contract/multiple_loops.rs
@@ -11,12 +11,14 @@
 fn multiple_loops() {
     let mut x: u8 = kani::any_where(|i| *i >= 10);
 
-    #[kani::loop_invariant(x >= 5)]
-    while x > 5 {
-        x = x - 1;
+    if x != 20 {
+        #[kani::loop_invariant(x >= 5)]
+        while x > 5 {
+            x = x - 1;
+        }
     }
 
-    assert!(x == 5);
+    assert!(x == 5 || x == 20);
 
     #[kani::loop_invariant(x >= 2)]
     while x > 2 {