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,169 @@
 //! 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, in a method, the
+//!   enclosing `impl`/`trait` header), so an invariant may refer to generic- and
+//!   `Self`-typed variables that the standalone macro cannot see.
+//! - On an `impl`/`trait`, each method is stamped with the enclosing header so
+//!   method-level `requires`/`ensures` can recover the outer generics, and its body is
+//!   threaded as above.
+//!
+//! Threading a body 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.
 
 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;
 
 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();
+    let item = syn::parse_macro_input!(item as ContextItem);
+    match item {
+        ContextItem::Fn(mut item_fn) => {
+            let outer = match crate::extract_outer_context(&item_fn.attrs) {
+                Ok(outer) => outer,
+                Err(e) => return e.to_compile_error().into(),
+            };
+            thread_into_body(&mut item_fn.sig, &mut item_fn.block, outer.as_ref());
+            item_fn.into_token_stream().into()
+        }
+        ContextItem::Outer(outer_item) => expand_outer(outer_item).into(),
+    }
+}
+
+/// An item `#[thrust_macros::context]` applies to.
+enum ContextItem {
+    Fn(syn::ItemFn),
+    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(item_fn) = fork.parse::<syn::ItemFn>() {
+            input.advance_to(&fork);
+            return Ok(Self::Fn(item_fn));
+        }
+
+        input.parse().map(Self::Outer)
+    }
+}
+
+fn expand_outer(mut outer_item: FnOuterItem) -> TokenStream2 {
+    let header = outer_item.clone().into_header_only();
+    let header_attr: syn::Attribute = syn::parse_quote!(#[thrust::_outer_context(#header)]);
     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.push(header_attr.clone());
+                thread_into_body(&mut item.sig, &mut item.block, Some(&header));
             }
         }
         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.push(header_attr.clone());
+                if let Some(default) = &mut item.default {
+                    thread_into_body(&mut item.sig, default, Some(&header));
+                }
             }
         }
     }
+    outer_item.into_token_stream()
+}
+
+/// Rewrites each spec macro in `body` into its context-carrying counterpart, carrying
+/// `sig` and the `outer` header of the `impl`/`trait` the function is a method of, and
+/// extends `sig`'s where clause with the `Model` predicates those calls need. A body
+/// naming no spec macro is left as it is.
+fn thread_into_body(sig: &mut Signature, body: &mut syn::Block, outer: Option<&FnOuterItem>) {
+    let host_sig = sig.clone();
+    let mut injector = ContextInjector {
+        sig: &host_sig,
+        outer,
+        injected: false,
+        self_used: false,
+    };
+    injector.visit_block_mut(body);
+    if !injector.injected {
+        return;
+    }
 
-    outer_item.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() {
+        sig.generics
+            .make_where_clause()
+            .predicates
+            .extend(predicates);
+    }
+}
+
+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) {
+        let Some(with_context) = context_carrying_form(&mac.path) else {
+            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 = with_context;
+    }
+}
+
+/// The context-carrying counterpart of a spec macro that takes a formula over live
+/// variables, or `None` for any other macro.
+fn context_carrying_form(path: &syn::Path) -> Option<syn::Path> {
+    // TODO: identify the macro precisely
+    match path.segments.last()?.ident.to_string().as_str() {
+        "invariant" => Some(syn::parse_quote!(::thrust_macros::_invariant_with_context)),
+        _ => None,
+    }
 }

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,