From 1fd83d6ee34e7fc7ef136030d92c6a28477c04dc Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Fri, 23 Oct 2026 21:32:44 +0000
Subject: [PATCH] Attach invariant! to the loop that contains it

`loop_header_of` took the first dominator of the marker block that has a
back edge, but a dominating loop header need not contain the marker. An
invariant written in an outer loop body after an inner loop was installed
at the inner loop's header, and one written after a loop outside every
loop was silently attached to that loop. Require the marker to lie in the
natural loop of one of the header's back edges.

Fixes #352

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01JTZrVqVEqBgdyunsY7yQ5P
---
 src/analyze/local_def.rs                      | 28 +++++++++++++++++--
 .../fail/loop_invariant_after_inner_loop.rs   | 22 +++++++++++++++
 .../pass/loop_invariant_after_inner_loop.rs   | 22 +++++++++++++++
 3 files changed, 70 insertions(+), 2 deletions(-)
 create mode 100644 tests/ui/fail/loop_invariant_after_inner_loop.rs
 create mode 100644 tests/ui/pass/loop_invariant_after_inner_loop.rs

diff --git a/src/analyze/local_def.rs b/src/analyze/local_def.rs
index 6a9cdf1b5..5672d42d2 100644
--- a/src/analyze/local_def.rs
+++ b/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`.
     ///
diff --git a/tests/ui/fail/loop_invariant_after_inner_loop.rs b/tests/ui/fail/loop_invariant_after_inner_loop.rs
new file mode 100644
index 000000000..c65e4e562
--- /dev/null
+++ b/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);
+}
diff --git a/tests/ui/pass/loop_invariant_after_inner_loop.rs b/tests/ui/pass/loop_invariant_after_inner_loop.rs
new file mode 100644
index 000000000..6de064bd2
--- /dev/null
+++ b/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);
+}