Fold invariant_context into context

tests/ui/fail/loop_invariant_fn_param_at_entry.rs

@@ -3,7 +3,7 @@
 
 #[thrust_macros::requires(true)]
 #[thrust_macros::ensures(result.length == v.length + 2)]
-#[thrust_macros::invariant_context]
+#[thrust_macros::context]
 fn push_two(v: Vec<i64>) -> Vec<i64> {
     let mut w = v;
     let mut i = 0_i64;

tests/ui/fail/loop_invariant_fn_param_closure.rs

@@ -5,7 +5,7 @@
 // `f.at_entry()` yields `Closure<F>`. Here the invariant relates `acc` to the
 // entry closure's postcondition, from which the postcondition below is proven.
 #[thrust_macros::ensures((n > 0) ==> thrust_macros::post!(f(n - 1), result))]
-#[thrust_macros::invariant_context]
+#[thrust_macros::context]
 fn last_apply<F>(f: F, n: i64) -> i64
 where
     F: Fn(i64) -> i64,
@@ -25,7 +25,7 @@ where
 
 // A capture-free closure is null (singleton) sorted; comparing its identity
 // must collapse to a canonical value rather than ICE during clause building.
-#[thrust_macros::invariant_context]
+#[thrust_macros::context]
 fn unchanged<F>(mut f: F)
 where
     F: FnMut(i64) -> i64,

tests/ui/fail/loop_invariant_generic.rs

@@ -6,7 +6,7 @@
 #[thrust::trusted]
 fn rand() -> i64 { unimplemented!() }
 
-#[thrust_macros::invariant_context]
+#[thrust_macros::context]
 fn keep<T: Copy + PartialEq>(v: T) {
   let mut x = v;
   while rand() == 0 {

tests/ui/fail/loop_invariant_outer_param.rs

@@ -9,7 +9,7 @@ fn rand() -> bool {
 }
 
 #[thrust_macros::ensures(result == a)]
-#[thrust_macros::invariant_context]
+#[thrust_macros::context]
 fn keep_argument(a: i64) -> i64 {
     let mut v = a;
 

tests/ui/fail/loop_invariant_self.rs

@@ -13,7 +13,6 @@ impl thrust_models::Model for Counter {
 
 #[thrust_macros::context]
 impl Counter {
-    #[thrust_macros::invariant_context]
     fn run(self) {
         let mut c = self;
         let mut x = 1_i64;

tests/ui/fail/loop_invariant_self_receiver.rs

@@ -18,7 +18,6 @@ impl thrust_models::Model for Counter {
 
 #[thrust_macros::context]
 impl Counter {
-    #[thrust_macros::invariant_context]
     fn run(&mut self) -> i64 {
         let init = *self;
         while rand() {

tests/ui/fail/loop_invariant_trait.rs

@@ -8,7 +8,6 @@ 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 {

tests/ui/fail/loop_invariant_trait_self.rs

@@ -16,7 +16,6 @@ trait Gauge {
 
     fn update(&mut self) -> i32;
 
-    #[thrust_macros::invariant_context]
     fn run(&mut self) -> i32 {
         let mut state = 0;
         while rand() == 0 {

tests/ui/pass/loop_invariant_fn_param_at_entry.rs

@@ -3,7 +3,7 @@
 
 #[thrust_macros::requires(true)]
 #[thrust_macros::ensures(result.length == v.length + 2)]
-#[thrust_macros::invariant_context]
+#[thrust_macros::context]
 fn push_two(v: Vec<i64>) -> Vec<i64> {
     let mut w = v;
     let mut i = 0_i64;

tests/ui/pass/loop_invariant_fn_param_closure.rs

@@ -5,7 +5,7 @@
 // `f.at_entry()` yields `Closure<F>`. Here the invariant relates `acc` to the
 // entry closure's postcondition, from which the postcondition below is proven.
 #[thrust_macros::ensures((n > 0) ==> thrust_macros::post!(f(n - 1), result))]
-#[thrust_macros::invariant_context]
+#[thrust_macros::context]
 fn last_apply<F>(f: F, n: i64) -> i64
 where
     F: Fn(i64) -> i64,
@@ -25,7 +25,7 @@ where
 
 // A capture-free closure is null (singleton) sorted; comparing its identity
 // must collapse to a canonical value rather than ICE during clause building.
-#[thrust_macros::invariant_context]
+#[thrust_macros::context]
 fn unchanged<F>(mut f: F)
 where
     F: FnMut(i64) -> i64,

tests/ui/pass/loop_invariant_generic.rs

@@ -6,7 +6,7 @@
 #[thrust::trusted]
 fn rand() -> i64 { unimplemented!() }
 
-#[thrust_macros::invariant_context]
+#[thrust_macros::context]
 fn keep<T: Copy + PartialEq>(v: T) {
   let mut x = v;
   while rand() == 0 {

tests/ui/pass/loop_invariant_generic_closure.rs

@@ -9,7 +9,7 @@ fn rand() -> i64 { unimplemented!() }
 // A closure-typed generic param must not be given a `Model` bound: the
 // invariant only constrains the `Model`-typed `T`, and `keep` must still be
 // callable with a real closure.
-#[thrust_macros::invariant_context]
+#[thrust_macros::context]
 fn keep<F: Fn(i64) -> i64, T: Copy + PartialEq>(f: F, v: T) {
   let _ = f;
   let mut x = v;

tests/ui/pass/loop_invariant_outer_param.rs

@@ -9,7 +9,7 @@ fn rand() -> bool {
 }
 
 #[thrust_macros::ensures(result == a)]
-#[thrust_macros::invariant_context]
+#[thrust_macros::context]
 fn keep_argument(a: i64) -> i64 {
     let mut v = a;
 

tests/ui/pass/loop_invariant_self.rs

@@ -13,7 +13,6 @@ impl thrust_models::Model for Counter {
 
 #[thrust_macros::context]
 impl Counter {
-    #[thrust_macros::invariant_context]
     fn run(self) {
         let mut c = self;
         let mut x = 1_i64;

tests/ui/pass/loop_invariant_self_receiver.rs

@@ -18,7 +18,6 @@ impl thrust_models::Model for Counter {
 
 #[thrust_macros::context]
 impl Counter {
-    #[thrust_macros::invariant_context]
     fn run(&mut self) -> i64 {
         let init = *self;
         while rand() {

tests/ui/pass/loop_invariant_trait.rs

@@ -8,7 +8,6 @@ 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 {

tests/ui/pass/loop_invariant_trait_self.rs

@@ -16,7 +16,6 @@ trait Gauge {
 
     fn update(&mut self) -> i32;
 
-    #[thrust_macros::invariant_context]
     fn run(&mut self) -> i32 {
         let mut state = 0;
         while rand() == 0 {

thrust-macros/src/context.rs

@@ -1,36 +1,164 @@
 //! Expansion of `#[thrust_macros::context]`.
 //!
-//! Stamps each method in an `impl`/`trait` block with the enclosing header so
-//! method-level `requires`/`ensures` can recover the outer generics.
+//! Makes the enclosing context available to the specifications written inside an item.
+//!
+//! On a function, every `thrust_macros::invariant!(...)` in the body is rewritten into
+//! its context-carrying counterpart, carrying the host signature and, for a method, the
+//! enclosing `impl`/`trait` header, so an invariant may refer to generic- and
+//! `Self`-typed variables that the standalone macro cannot see. That also extends the
+//! function's where clause with the `Model` predicates for every in-scope type parameter
+//! (and for `Self` when used), since each injected marker call instantiates a
+//! `Model`-bounded formula function with the host's own generics.
+//!
+//! On an `impl`/`trait`, each method is stamped with the enclosing header — which is what
+//! method-level `requires`/`ensures` read to recover the outer generics — and with this
+//! attribute, so a method's body is threaded by an expansion of its own.
 
 use proc_macro::TokenStream;
-use quote::ToTokens as _;
+use proc_macro2::TokenStream as TokenStream2;
+use quote::{quote, ToTokens as _};
+use syn::{
+    parse::{Parse, ParseStream},
+    visit_mut::VisitMut,
+    Signature,
+};
 
-use crate::fn_outer_item::FnOuterItem;
+use crate::{fn_outer_item::FnOuterItem, spec::FnItemWithSignature};
 
 pub fn expand(item: TokenStream) -> TokenStream {
-    let mut outer_item = syn::parse_macro_input!(item as FnOuterItem);
-    let outer_header = outer_item.clone().into_header_only();
+    match syn::parse_macro_input!(item as ContextItem) {
+        ContextItem::Fn(func) => expand_fn(func),
+        ContextItem::Outer(outer_item) => expand_outer(outer_item),
+    }
+}
+
+/// An item `#[thrust_macros::context]` applies to.
+enum ContextItem {
+    Fn(FnItemWithSignature),
+    Outer(FnOuterItem),
+}
+
+impl Parse for ContextItem {
+    fn parse(input: ParseStream) -> syn::Result<Self> {
+        use syn::parse::discouraged::Speculative as _;
+
+        let fork = input.fork();
+        if let Ok(func) = fork.parse::<FnItemWithSignature>() {
+            input.advance_to(&fork);
+            return Ok(Self::Fn(func));
+        }
+
+        input.parse().map(Self::Outer)
+    }
+}
+
+/// Hands each method the enclosing header, and the attribute that puts it to use.
+fn expand_outer(mut outer_item: FnOuterItem) -> TokenStream {
+    let header = outer_item.clone().into_header_only();
+    let method_attrs: [syn::Attribute; 2] = [
+        syn::parse_quote!(#[thrust::_outer_context(#header)]),
+        syn::parse_quote!(#[::thrust_macros::context]),
+    ];
     match &mut outer_item {
         FnOuterItem::ItemImpl(item_impl) => {
             for item in &mut item_impl.items {
                 let syn::ImplItem::Fn(item) = item else {
                     continue;
                 };
-                item.attrs
-                    .push(syn::parse_quote!(#[thrust::_outer_context(#outer_header)]));
+                item.attrs.extend(method_attrs.clone());
             }
         }
         FnOuterItem::ItemTrait(item_trait) => {
             for item in &mut item_trait.items {
                 let syn::TraitItem::Fn(item) = item else {
                     continue;
                 };
-                item.attrs
-                    .push(syn::parse_quote!(#[thrust::_outer_context(#outer_header)]));
+                item.attrs.extend(method_attrs.clone());
             }
         }
     }
-
     outer_item.into_token_stream().into()
 }
+
+/// Rewrites each `invariant!` in the body into its context-carrying counterpart and
+/// extends the where clause with the `Model` predicates those calls need. A body naming
+/// no invariant — or a trait method that has no body at all — is left as it is.
+fn expand_fn(mut func: FnItemWithSignature) -> TokenStream {
+    let outer = match crate::extract_outer_context(func.attrs()) {
+        Ok(outer) => outer,
+        Err(e) => return e.to_compile_error().into(),
+    };
+
+    let host_sig = func.sig().clone();
+    let mut injector = ContextInjector {
+        sig: &host_sig,
+        outer: outer.as_ref(),
+        injected: false,
+        self_used: false,
+    };
+    if let Some(body) = func.block_mut() {
+        injector.visit_block_mut(body);
+    }
+    if !injector.injected {
+        return func.into_token_stream().into();
+    }
+
+    let type_lowering = match &outer {
+        Some(outer) => crate::FormulaFnTypeLowering::with_outer_context(&host_sig, outer),
+        None => crate::FormulaFnTypeLowering::new(&host_sig),
+    };
+    let mut predicates = type_lowering.model_where_predicates();
+    if injector.self_used {
+        predicates.extend(type_lowering.model_where_predicates_for(&quote::format_ident!("Self")));
+    }
+    if !predicates.is_empty() {
+        func.sig_mut()
+            .generics
+            .make_where_clause()
+            .predicates
+            .extend(predicates);
+    }
+
+    func.into_token_stream().into()
+}
+
+struct ContextInjector<'a> {
+    sig: &'a Signature,
+    outer: Option<&'a FnOuterItem>,
+    injected: bool,
+    self_used: bool,
+}
+
+impl ContextInjector<'_> {
+    fn inject_context(&self, closure: &TokenStream2) -> TokenStream2 {
+        let sig = self.sig;
+        let outer_attr = self
+            .outer
+            .map(|outer| quote!(#[thrust::_outer_context(#outer)]));
+
+        quote! {
+            #outer_attr
+            #sig;
+            #closure
+        }
+    }
+}
+
+impl VisitMut for ContextInjector<'_> {
+    fn visit_macro_mut(&mut self, mac: &mut syn::Macro) {
+        if !is_invariant_macro(&mac.path) {
+            return;
+        }
+        self.injected = true;
+        if crate::tokens_contain_ident(&mac.tokens, "Self") {
+            self.self_used = true;
+        }
+        mac.tokens = self.inject_context(&mac.tokens);
+        mac.path = syn::parse_quote!(::thrust_macros::_invariant_with_context);
+    }
+}
+
+fn is_invariant_macro(path: &syn::Path) -> bool {
+    // TODO: identify the macro precisely
+    path.segments.last().is_some_and(|s| s.ident == "invariant")
+}

thrust-macros/src/invariant.rs

@@ -8,7 +8,7 @@
 //! - `invariant!(|x: i64| x >= 1)` takes a bare predicate closure and only sees
 //!   concrete types.
 //! - `_invariant_with_context!(..)` additionally carries the enclosing generic
-//!   context. It is never written by hand: `#[thrust_macros::invariant_context]`
+//!   context. It is never written by hand: `#[thrust_macros::context]`
 //!   rewrites each `invariant!` it finds into this form, pasting the host
 //!   function's signature (and, in methods, a `#[thrust::_outer_context(..)]`
 //!   attribute carrying the enclosing `impl`/`trait` header) ahead of the
@@ -57,7 +57,7 @@ pub fn expand(input: TokenStream) -> TokenStream {
 }
 
 /// Expands `_invariant_with_context!(#outer_attr #sig; CLOSURE)`, the form
-/// `#[thrust_macros::invariant_context]` rewrites each `invariant!` into.
+/// `#[thrust_macros::context]` rewrites each `invariant!` into.
 pub fn expand_with_context(input: TokenStream) -> TokenStream {
     struct WithContext {
         context: Context,
@@ -90,7 +90,7 @@ pub fn expand_with_context(input: TokenStream) -> TokenStream {
 }
 
 /// The enclosing context threaded into an invariant by
-/// `#[thrust_macros::invariant_context]`: the host function signature and, for a
+/// `#[thrust_macros::context]`: the host function signature and, for a
 /// method, its `impl`/`trait` header. A standalone `invariant!` has none.
 struct Context {
     sig: Signature,