Fix E0401 for invariant_context on trait methods

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();
+}

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();
+}

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,