Attach invariant! to the loop that contains it

src/analyze/local_def.rs

@@ -768,14 +768,17 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
 
     /// Walks up the dominator tree from the marker block to the innermost
     /// enclosing loop header: the first dominator that needs its own
-    /// precondition (in-degree >= 2) and has a back edge.
+    /// precondition (in-degree >= 2) and has a back edge whose natural loop
+    /// contains the marker block.
     fn loop_header_of(body: &Body<'_>, marker_bb: BasicBlock) -> Option<BasicBlock> {
         let doms = body.basic_blocks.dominators();
         let preds = body.basic_blocks.predecessors();
         let mut cur = Some(marker_bb);
         while let Some(bb) = cur {
             if analyze::basic_block::needs_own_precondition(body, bb)
-                && preds[bb].iter().any(|&p| doms.dominates(bb, p))
+                && preds[bb].iter().any(|&p| {
+                    doms.dominates(bb, p) && Self::natural_loop_contains(body, bb, p, marker_bb)
+                })
             {
                 return Some(bb);
             }
@@ -784,6 +787,27 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         None
     }
 
+    /// Whether `bb` belongs to the natural loop of the back edge `latch -> header`:
+    /// the blocks that reach `latch` without passing through `header`, and `header` itself.
+    fn natural_loop_contains(
+        body: &Body<'_>,
+        header: BasicBlock,
+        latch: BasicBlock,
+        bb: BasicBlock,
+    ) -> bool {
+        let preds = body.basic_blocks.predecessors();
+        let mut visited = DenseBitSet::new_empty(body.basic_blocks.len());
+        visited.insert(header);
+        let mut stack = vec![latch];
+        while let Some(cur) = stack.pop() {
+            if !visited.insert(cur) {
+                continue;
+            }
+            stack.extend(preds[cur].iter().copied());
+        }
+        visited.contains(bb)
+    }
+
     /// Resolves the live local matching a source variable name at the given
     /// basic block, among the locals that are parameters of `bty`.
     ///

tests/ui/fail/loop_invariant_after_inner_loop.rs

@@ -0,0 +1,22 @@
+//@error-in-other-file: Unsat
+//@compile-flags: -C debug-assertions=off
+
+#[thrust_macros::requires(true)]
+#[thrust_macros::ensures(true)]
+#[thrust::trusted]
+fn rand() -> i64 { unimplemented!() }
+
+fn main() {
+  let mut total = 0_i64;
+  let mut i = 0_i64;
+  while rand() == 0 {
+    let mut j = 0_i64;
+    while j < 3 {
+      total += 1;
+      j += 1;
+    }
+    i += 1;
+    thrust_macros::invariant!(|total: i64, i: i64| total == 2 * i);
+  }
+  assert!(total == 3 * i);
+}

tests/ui/pass/loop_invariant_after_inner_loop.rs

@@ -0,0 +1,22 @@
+//@check-pass
+//@compile-flags: -C debug-assertions=off
+
+#[thrust_macros::requires(true)]
+#[thrust_macros::ensures(true)]
+#[thrust::trusted]
+fn rand() -> i64 { unimplemented!() }
+
+fn main() {
+  let mut total = 0_i64;
+  let mut i = 0_i64;
+  while rand() == 0 {
+    let mut j = 0_i64;
+    while j < 3 {
+      total += 1;
+      j += 1;
+    }
+    i += 1;
+    thrust_macros::invariant!(|total: i64, i: i64| total == 3 * i);
+  }
+  assert!(total == 3 * i);
+}