From 21770e6a4bf3eab6a16635b2a6c46f3836389c88 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Mon, 15 Jun 2026 16:39:00 +0000
Subject: [PATCH] Fix E0401 for invariant_context on trait methods

---
 tests/ui/fail/loop_invariant_trait.rs | 30 +++++++++++++++++++++++++++
 tests/ui/pass/loop_invariant_trait.rs | 30 +++++++++++++++++++++++++++
 thrust-macros/src/invariant.rs        |  6 +++++-
 3 files changed, 65 insertions(+), 1 deletion(-)
 create mode 100644 tests/ui/fail/loop_invariant_trait.rs
 create mode 100644 tests/ui/pass/loop_invariant_trait.rs

diff --git a/tests/ui/fail/loop_invariant_trait.rs b/tests/ui/fail/loop_invariant_trait.rs
new file mode 100644
index 00000000..cc7d95ed
--- /dev/null
+++ b/tests/ui/fail/loop_invariant_trait.rs
@@ -0,0 +1,30 @@
+//@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!() }
+
+#[thrust_macros::context]
+trait Foo {
+    #[thrust_macros::invariant_context]
+    fn run(&mut self) {
+        let mut x: i64 = 0;
+        while rand() == 0 {
+            thrust_macros::invariant!(|x: i64| x >= 1);
+            x += 1;
+        }
+        assert!(x >= 0);
+    }
+}
+
+struct Bar;
+impl thrust_models::Model for Bar {
+    type Ty = ();
+}
+impl Foo for Bar {}
+
+fn main() {
+    Bar.run();
+}
diff --git a/tests/ui/pass/loop_invariant_trait.rs b/tests/ui/pass/loop_invariant_trait.rs
new file mode 100644
index 00000000..8acf26ee
--- /dev/null
+++ b/tests/ui/pass/loop_invariant_trait.rs
@@ -0,0 +1,30 @@
+//@check-pass
+//@compile-flags: -C debug-assertions=off
+
+#[thrust_macros::requires(true)]
+#[thrust_macros::ensures(true)]
+#[thrust::trusted]
+fn rand() -> i64 { unimplemented!() }
+
+#[thrust_macros::context]
+trait Foo {
+    #[thrust_macros::invariant_context]
+    fn run(&mut self) {
+        let mut x: i64 = 0;
+        while rand() == 0 {
+            thrust_macros::invariant!(|x: i64| x >= 0);
+            x += 1;
+        }
+        assert!(x >= 0);
+    }
+}
+
+struct Bar;
+impl thrust_models::Model for Bar {
+    type Ty = ();
+}
+impl Foo for Bar {}
+
+fn main() {
+    Bar.run();
+}
diff --git a/thrust-macros/src/invariant.rs b/thrust-macros/src/invariant.rs
index 9f753069..761fa175 100644
--- a/thrust-macros/src/invariant.rs
+++ b/thrust-macros/src/invariant.rs
@@ -177,7 +177,11 @@ fn expand_invariant(
     }
     rewriter.visit_expr_mut(&mut body);
 
-    if crate::tokens_contain_ident(&closure.to_token_stream(), "Self") {
+    let self_used = crate::tokens_contain_ident(&closure.to_token_stream(), "Self")
+        || def_wheres
+            .iter()
+            .any(|pred| crate::tokens_contain_ident(&pred.to_token_stream(), "Self"));
+    if self_used {
         let Some(outer) = context.and_then(|context| context.outer.as_ref()) else {
             return Err(syn::Error::new_spanned(
                 closure,