From 5705d6d2bddcb8df6f2519250209fcfefe0183b0 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 1 Jun 2026 07:04:06 +0000 Subject: [PATCH 01/13] Support Self and generic variables in loop invariants MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `#[thrust_macros::invariant_context]` — a new attribute that threads the surrounding generic context (and, in methods, `Self`) into every `thrust_macros::invariant!(...)` inside the annotated function, impl, or trait. Each call is replaced in-place by an expanded `#[thrust::formula_fn]` plus a marker call, with the threaded generics declared on the formula function and instantiated via turbofish. In methods, `Self` is re-declared as a synthetic type parameter and instantiated with the real `Self` (legal in expression position), so an invariant may refer to generic- and `Self`-typed variables that the standalone `invariant!` macro cannot see. `thrust_macros::context` and the standalone `thrust_macros::invariant!` macro are unchanged; `invariant_context` rewrites the macro tokens before `invariant!` would otherwise expand, so the simple `invariant!` path keeps working as-is for concrete invariants. Adds pass/fail UI test pairs for the generic and `Self` cases. Note that a struct used as an invariant variable must model structurally (its `Model::Ty` must mirror its fields, e.g. a one-field tuple struct models as a one-tuple); a mismatching model trips an existing subtyping panic unrelated to this change. https://claude.ai/code/session_01WB28auaD8dSQrckqBwJWBt --- tests/ui/fail/loop_invariant_generic.rs | 22 ++ tests/ui/fail/loop_invariant_self.rs | 29 +++ tests/ui/pass/loop_invariant_generic.rs | 22 ++ tests/ui/pass/loop_invariant_self.rs | 29 +++ thrust-macros/src/invariant_context.rs | 303 ++++++++++++++++++++++++ thrust-macros/src/lib.rs | 10 + 6 files changed, 415 insertions(+) create mode 100644 tests/ui/fail/loop_invariant_generic.rs create mode 100644 tests/ui/fail/loop_invariant_self.rs create mode 100644 tests/ui/pass/loop_invariant_generic.rs create mode 100644 tests/ui/pass/loop_invariant_self.rs create mode 100644 thrust-macros/src/invariant_context.rs diff --git a/tests/ui/fail/loop_invariant_generic.rs b/tests/ui/fail/loop_invariant_generic.rs new file mode 100644 index 0000000..5342e97 --- /dev/null +++ b/tests/ui/fail/loop_invariant_generic.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!() } + +#[thrust_macros::invariant_context] +fn keep(v: T) { + let mut x = v; + while rand() == 0 { + thrust_macros::invariant!(|v: T| v == v); + x = v; + } + assert!(x == v); +} + +fn main() { + keep(0_i64); + keep(true); +} diff --git a/tests/ui/fail/loop_invariant_self.rs b/tests/ui/fail/loop_invariant_self.rs new file mode 100644 index 0000000..04f775b --- /dev/null +++ b/tests/ui/fail/loop_invariant_self.rs @@ -0,0 +1,29 @@ +//@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!() } + +struct Counter(i64); +impl thrust_models::Model for Counter { + type Ty = (thrust_models::model::Int,); +} + +#[thrust_macros::invariant_context] +impl Counter { + fn run(self) { + let mut c = self; + let mut x = 1_i64; + while rand() == 0 { + thrust_macros::invariant!(|x: i64, c: Self| x >= 2 && c == c); + x = x + 1; + c = Counter(0); + } + let _last = c; + assert!(x >= 1); + } +} + +fn main() { Counter(0).run(); } diff --git a/tests/ui/pass/loop_invariant_generic.rs b/tests/ui/pass/loop_invariant_generic.rs new file mode 100644 index 0000000..7e31b00 --- /dev/null +++ b/tests/ui/pass/loop_invariant_generic.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!() } + +#[thrust_macros::invariant_context] +fn keep(v: T) { + let mut x = v; + while rand() == 0 { + thrust_macros::invariant!(|x: T, v: T| x == v); + x = v; + } + assert!(x == v); +} + +fn main() { + keep(0_i64); + keep(true); +} diff --git a/tests/ui/pass/loop_invariant_self.rs b/tests/ui/pass/loop_invariant_self.rs new file mode 100644 index 0000000..31cfb49 --- /dev/null +++ b/tests/ui/pass/loop_invariant_self.rs @@ -0,0 +1,29 @@ +//@check-pass +//@compile-flags: -C debug-assertions=off + +#[thrust_macros::requires(true)] +#[thrust_macros::ensures(true)] +#[thrust::trusted] +fn rand() -> i64 { unimplemented!() } + +struct Counter(i64); +impl thrust_models::Model for Counter { + type Ty = (thrust_models::model::Int,); +} + +#[thrust_macros::invariant_context] +impl Counter { + fn run(self) { + let mut c = self; + let mut x = 1_i64; + while rand() == 0 { + thrust_macros::invariant!(|x: i64, c: Self| x >= 1 && c == c); + x = x + 1; + c = Counter(0); + } + let _last = c; + assert!(x >= 1); + } +} + +fn main() { Counter(0).run(); } diff --git a/thrust-macros/src/invariant_context.rs b/thrust-macros/src/invariant_context.rs new file mode 100644 index 0000000..1d3d4ab --- /dev/null +++ b/thrust-macros/src/invariant_context.rs @@ -0,0 +1,303 @@ +//! Expansion of `#[thrust_macros::invariant_context]`. +//! +//! Threads the surrounding generic context (and, in methods, `Self`) into +//! every `thrust_macros::invariant!(...)` call inside the annotated item. Each +//! such call is replaced in-place by an expanded `#[thrust::formula_fn]` plus +//! a marker call, with the threaded generics declared on the formula function +//! and instantiated via turbofish. In methods, `Self` is re-declared as a +//! synthetic type parameter and instantiated with the real `Self` the same +//! way, so an invariant may refer to generic- and `Self`-typed variables that +//! the standalone `invariant!` macro cannot see. +//! +//! The host function's where clause gains `T: Model` and +//! `::Ty: PartialEq` predicates for every in-scope type parameter +//! (and for `Self` when used), as the generated marker call instantiates a +//! `Model`-bounded function. + +use std::sync::atomic::{AtomicUsize, Ordering}; + +use proc_macro::TokenStream; +use proc_macro2::{TokenStream as TokenStream2, TokenTree}; +use quote::{format_ident, quote, ToTokens}; +use syn::{parse_macro_input, visit_mut::VisitMut, FnArg, GenericParam, Generics, WherePredicate}; + +use crate::{fn_params_with_model_ty, FnOuterItem}; + +static COUNTER: AtomicUsize = AtomicUsize::new(0); + +pub(super) fn expand(item: TokenStream) -> TokenStream { + let mut item = parse_macro_input!(item as syn::Item); + + match &mut item { + syn::Item::Fn(item_fn) => { + let generics = item_fn.sig.generics.clone(); + let mut rewriter = InvariantRewriter::new(&generics, None); + rewriter.visit_block_mut(&mut item_fn.block); + if rewriter.found { + inject_model_bounds(&mut item_fn.sig.generics, None, rewriter.self_used); + } + } + syn::Item::Impl(item_impl) => { + let outer = FnOuterItem::ItemImpl(item_impl.clone()).into_header_only(); + for impl_item in &mut item_impl.items { + let syn::ImplItem::Fn(method) = impl_item else { + continue; + }; + let generics = method.sig.generics.clone(); + let mut rewriter = InvariantRewriter::new(&generics, Some(&outer)); + rewriter.visit_block_mut(&mut method.block); + if rewriter.found { + inject_model_bounds(&mut method.sig.generics, Some(&outer), rewriter.self_used); + } + } + } + syn::Item::Trait(item_trait) => { + let outer = FnOuterItem::ItemTrait(item_trait.clone()).into_header_only(); + for trait_item in &mut item_trait.items { + let syn::TraitItem::Fn(method) = trait_item else { + continue; + }; + let Some(block) = &mut method.default else { + continue; + }; + let generics = method.sig.generics.clone(); + let mut rewriter = InvariantRewriter::new(&generics, Some(&outer)); + rewriter.visit_block_mut(block); + if rewriter.found { + inject_model_bounds(&mut method.sig.generics, Some(&outer), rewriter.self_used); + } + } + } + _ => { + return syn::Error::new_spanned( + &item, + "#[thrust_macros::invariant_context] expects a function, impl block, or trait \ + definition", + ) + .to_compile_error() + .into(); + } + } + + item.into_token_stream().into() +} + +struct InvariantRewriter { + /// In-scope generic params (the function's own, plus the outer impl/trait's + /// own for a method). + params: Vec, + /// Combined where-predicates from the function and (for a method) the + /// outer impl/trait. `Model`/`PartialEq` bounds for `params` are added + /// when constructing the formula function so each turbofish argument + /// satisfies them. + wheres: Vec, + is_method: bool, + found: bool, + self_used: bool, +} + +impl InvariantRewriter { + fn new(fn_generics: &Generics, outer: Option<&FnOuterItem>) -> Self { + let mut params: Vec = fn_generics.params.iter().cloned().collect(); + let mut wheres: Vec = fn_generics + .where_clause + .as_ref() + .map(|wc| wc.predicates.iter().cloned().collect()) + .unwrap_or_default(); + if let Some(outer) = outer { + params.extend(outer.generics().params.iter().cloned()); + if let Some(wc) = &outer.generics().where_clause { + wheres.extend(wc.predicates.iter().cloned()); + } + } + Self { + params, + wheres, + is_method: outer.is_some(), + found: false, + self_used: false, + } + } + + /// Builds the expansion for a single `invariant!` macro call. + fn expand_invariant(&mut self, mac: &syn::Macro) -> syn::Expr { + let closure: syn::ExprClosure = match syn::parse2(mac.tokens.clone()) { + Ok(c) => c, + Err(e) => return syn::Expr::Verbatim(e.to_compile_error()), + }; + + let mut fn_params: Vec = Vec::new(); + for param in &closure.inputs { + let syn::Pat::Type(pt) = param else { + let err = syn::Error::new_spanned( + param, + "invariant closure parameters must have explicit types, e.g. `|x: i64| ...`", + ); + return syn::Expr::Verbatim(err.to_compile_error()); + }; + let pat = &pt.pat; + let ty = &pt.ty; + fn_params.push(syn::parse_quote!(#pat: #ty)); + } + + let mut def_params: Vec = Vec::new(); + let mut def_wheres: Vec = + self.wheres.iter().map(|w| w.to_token_stream()).collect(); + let mut turbofish_args: Vec = Vec::new(); + + // `Self` in a method context: rewrite to a synthetic generic, then + // pass the real `Self` via turbofish (legal in expression position). + let uses_self = self.is_method && tokens_contain_self(&closure.to_token_stream()); + if uses_self { + self.self_used = true; + let synth: syn::Ident = format_ident!("__ThrustSelf"); + for param in &mut fn_params { + SelfRewriter { synth: &synth }.visit_fn_arg_mut(param); + } + def_params.push(quote!(#synth)); + def_wheres.push(quote!(#synth: thrust_models::Model)); + def_wheres.push(quote!(<#synth as thrust_models::Model>::Ty: PartialEq)); + turbofish_args.push(quote!(Self)); + } + + for param in &self.params { + def_params.push(param.to_token_stream()); + match param { + GenericParam::Type(tp) => { + let ident = &tp.ident; + def_wheres.push(quote!(#ident: thrust_models::Model)); + def_wheres.push(quote!(<#ident as thrust_models::Model>::Ty: PartialEq)); + turbofish_args.push(ident.to_token_stream()); + } + GenericParam::Const(cp) => turbofish_args.push(cp.ident.to_token_stream()), + GenericParam::Lifetime(_) => {} + } + } + + let model_ty_params = fn_params_with_model_ty(&fn_params); + let body = &closure.body; + + let id = COUNTER.fetch_add(1, Ordering::Relaxed); + let name = format_ident!("_thrust_invariant_{}", id); + + let def_generics = if def_params.is_empty() { + quote!() + } else { + quote!(<#(#def_params),*>) + }; + let where_clause = if def_wheres.is_empty() { + quote!() + } else { + quote!(where #(#def_wheres),*) + }; + let turbofish = if turbofish_args.is_empty() { + quote!() + } else { + quote!(::<#(#turbofish_args),*>) + }; + + syn::parse_quote!({ + #[allow(unused_variables)] + #[allow(non_snake_case)] + #[thrust::formula_fn] + fn #name #def_generics(#model_ty_params) -> bool #where_clause { + #body + } + + thrust_models::__invariant_marker(#name #turbofish) + }) + } +} + +impl VisitMut for InvariantRewriter { + fn visit_stmt_mut(&mut self, stmt: &mut syn::Stmt) { + syn::visit_mut::visit_stmt_mut(self, stmt); + if let syn::Stmt::Macro(stmt_macro) = stmt { + if is_invariant_macro(&stmt_macro.mac.path) { + self.found = true; + let expanded = self.expand_invariant(&stmt_macro.mac); + *stmt = syn::Stmt::Expr(expanded, stmt_macro.semi_token); + } + } + } + + fn visit_expr_mut(&mut self, expr: &mut syn::Expr) { + syn::visit_mut::visit_expr_mut(self, expr); + if let syn::Expr::Macro(expr_macro) = expr { + if is_invariant_macro(&expr_macro.mac.path) { + self.found = true; + *expr = self.expand_invariant(&expr_macro.mac); + } + } + } +} + +/// Rewrites a bare `Self` type path to a synthetic type parameter, so the type +/// can be named inside a nested formula function. Qualified paths such as +/// `Self::Assoc` are left untouched (and are not supported in invariants). +struct SelfRewriter<'a> { + synth: &'a syn::Ident, +} + +impl VisitMut for SelfRewriter<'_> { + fn visit_path_mut(&mut self, path: &mut syn::Path) { + syn::visit_mut::visit_path_mut(self, path); + if path.leading_colon.is_none() + && path.segments.len() == 1 + && path.segments[0].ident == "Self" + { + path.segments[0].ident = self.synth.clone(); + } + } +} + +/// Adds `T: Model` and `::Ty: PartialEq` bounds for every type +/// parameter in scope to a function's where clause. Each marker call generated +/// by the rewriter instantiates a `Model`-bounded formula function, so the +/// host function must itself satisfy those bounds. When an invariant names +/// `Self`, the rewriter instantiates the formula function with `Self`, so the +/// same bounds are added for `Self` (`with_self`). +fn inject_model_bounds(generics: &mut Generics, outer: Option<&FnOuterItem>, with_self: bool) { + let mut tys: Vec = generics + .params + .iter() + .filter_map(|p| match p { + GenericParam::Type(tp) => Some(tp.ident.to_token_stream()), + _ => None, + }) + .collect(); + if let Some(outer) = outer { + for param in &outer.generics().params { + if let GenericParam::Type(tp) = param { + tys.push(tp.ident.to_token_stream()); + } + } + } + if with_self { + tys.push(quote!(Self)); + } + if tys.is_empty() { + return; + } + let where_clause = generics.make_where_clause(); + for ty in tys { + where_clause + .predicates + .push(syn::parse_quote!(#ty: thrust_models::Model)); + where_clause + .predicates + .push(syn::parse_quote!(<#ty as thrust_models::Model>::Ty: PartialEq)); + } +} + +fn is_invariant_macro(path: &syn::Path) -> bool { + path.segments.last().is_some_and(|s| s.ident == "invariant") +} + +fn tokens_contain_self(tokens: &TokenStream2) -> bool { + tokens.clone().into_iter().any(|tt| match tt { + TokenTree::Ident(ident) => ident == "Self", + TokenTree::Group(group) => tokens_contain_self(&group.stream()), + _ => false, + }) +} diff --git a/thrust-macros/src/lib.rs b/thrust-macros/src/lib.rs index 9f62b3a..5be5be4 100644 --- a/thrust-macros/src/lib.rs +++ b/thrust-macros/src/lib.rs @@ -4,6 +4,7 @@ use proc_macro2::TokenStream as TokenStream2; mod context; mod fn_outer_item; mod invariant; +mod invariant_context; mod spec; #[proc_macro_attribute] @@ -30,6 +31,15 @@ pub fn invariant(input: TokenStream) -> TokenStream { invariant::expand(input) } +/// Threads the surrounding generic context (and, in methods, `Self`) into +/// every `thrust_macros::invariant!(...)` inside the annotated function, impl, +/// or trait, so an invariant may refer to generic- and `Self`-typed variables +/// that the standalone `invariant!` macro cannot see. +#[proc_macro_attribute] +pub fn invariant_context(_attr: TokenStream, item: TokenStream) -> TokenStream { + invariant_context::expand(item) +} + #[proc_macro_attribute] pub fn predicate(_attr: TokenStream, item: TokenStream) -> TokenStream { spec::expand_predicate(item) From 395cae4372ea3f11db61ace3a431cfd065bd57e2 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 2 Jun 2026 15:40:18 +0000 Subject: [PATCH 02/13] thrust-macros: make invariant_context inject context, invariant! expand it Previously invariant_context re-implemented the whole invariant! expansion (closure parsing, formula_fn synthesis, marker call), duplicating the logic already in invariant.rs. Now the responsibilities are split, mirroring the existing requires/ensures _outer_context convention: - invariant_context only *injects* the enclosing context into each invariant!(CLOSURE) call, rewriting it into a context-carrying form built from existing syntax (a fn header carrying the in-scope generics/where clause, tagged for methods with #[thrust::_outer_context(impl/trait header)]). It still extends the host fn's where clause with the Model/PartialEq bounds, which a function-like macro cannot do itself. - invariant! now performs the single, shared expansion for both the standalone concrete form and the threaded form, re-declaring the threaded generics (and synthetic Self) on the formula function and instantiating via turbofish. No new bracket syntax is introduced; the context rides on ordinary Rust fn header syntax and the existing _outer_context attribute. --- thrust-macros/src/invariant.rs | 250 ++++++++++++++++++--- thrust-macros/src/invariant_context.rs | 294 +++++++------------------ 2 files changed, 304 insertions(+), 240 deletions(-) diff --git a/thrust-macros/src/invariant.rs b/thrust-macros/src/invariant.rs index 8edf62e..24bdd24 100644 --- a/thrust-macros/src/invariant.rs +++ b/thrust-macros/src/invariant.rs @@ -1,54 +1,248 @@ //! Expansion of `thrust_macros::invariant!`. //! -//! Expands a closure with concrete parameter types into a -//! `#[thrust::formula_fn]` over `Model::Ty` parameters and a marker call +//! Expands a predicate closure with explicit parameter types into a +//! `#[thrust::formula_fn]` over `Model::Ty` parameters plus a marker call //! referencing it. +//! +//! A standalone `invariant!(|x: i64| x >= 1)` only sees concrete types. To +//! refer to generic- or `Self`-typed variables, the enclosing item must be +//! annotated with `#[thrust_macros::invariant_context]`, which rewrites each +//! call into a context-carrying form that threads the in-scope generics (and, +//! in methods, `Self`) into the macro input: +//! +//! ```ignore +//! invariant!( +//! #[thrust::_outer_context(impl Foo where ..)] // methods only +//! fn _ctx() where .. { +//! |x: T, v: T| x == v +//! } +//! ) +//! ``` +//! +//! This macro re-declares the threaded generics (shadowing the enclosing ones) +//! on the formula function and instantiates it via turbofish; in methods, +//! `Self` is re-declared as a synthetic type parameter and instantiated with +//! the real `Self` (legal in expression position). use std::sync::atomic::{AtomicUsize, Ordering}; use proc_macro::TokenStream; -use quote::{format_ident, quote}; -use syn::{parse_macro_input, FnArg}; +use proc_macro2::{TokenStream as TokenStream2, TokenTree}; +use quote::{format_ident, quote, ToTokens}; +use syn::{ + parse::{Parse, ParseStream}, + parse_macro_input, + visit_mut::VisitMut, + FnArg, GenericParam, WherePredicate, +}; -use crate::fn_params_with_model_ty; +use crate::FnOuterItem; static COUNTER: AtomicUsize = AtomicUsize::new(0); pub fn expand(input: TokenStream) -> TokenStream { - let closure = parse_macro_input!(input as syn::ExprClosure); - - let mut fn_params: Vec = Vec::new(); - for param in &closure.inputs { - let syn::Pat::Type(pt) = param else { - return syn::Error::new_spanned( - param, - "invariant closure parameters must have explicit types, e.g. `|x: i64| ...`", - ) - .to_compile_error() - .into(); + let invariant = parse_macro_input!(input as Invariant); + invariant.expand().into_token_stream().into() +} + +/// The parsed `invariant!` input: the predicate closure plus the enclosing +/// generic context threaded in by `#[thrust_macros::invariant_context]`. +struct Invariant { + closure: syn::ExprClosure, + /// In-scope generic params (the enclosing function's own, plus the outer + /// `impl`/`trait`'s for a method). Empty for a standalone `invariant!`. + params: Vec, + /// Combined where-predicates from the function and (for a method) the + /// outer `impl`/`trait`. + wheres: Vec, + /// Whether `Self` is nameable at the call site (i.e. we are in a method). + is_method: bool, +} + +impl Parse for Invariant { + fn parse(input: ParseStream) -> syn::Result { + // Context-carrying form injected by `invariant_context`: a function + // header whose generics/where clause carry the enclosing context and + // whose body is the predicate closure. It always starts with the + // `#[thrust::_outer_context(..)]` attribute (methods) or `fn` (free + // functions); a standalone predicate closure starts with neither. + if input.peek(syn::Token![#]) || input.peek(syn::Token![fn]) { + let ctx_fn: syn::ItemFn = input.parse()?; + return Self::from_context_fn(ctx_fn); + } + + Ok(Self { + closure: input.parse()?, + params: Vec::new(), + wheres: Vec::new(), + is_method: false, + }) + } +} + +impl Invariant { + fn from_context_fn(ctx_fn: syn::ItemFn) -> syn::Result { + let outer = extract_outer_context(&ctx_fn.attrs)?; + + let closure = match ctx_fn.block.stmts.as_slice() { + [syn::Stmt::Expr(syn::Expr::Closure(closure), _)] => closure.clone(), + _ => { + return Err(syn::Error::new_spanned( + &ctx_fn.block, + "invariant context body must be a single predicate closure", + )) + } }; - let pat = &pt.pat; - let ty = &pt.ty; - fn_params.push(syn::parse_quote!(#pat: #ty)); + + let mut params: Vec = ctx_fn.sig.generics.params.iter().cloned().collect(); + let mut wheres: Vec = ctx_fn + .sig + .generics + .where_clause + .as_ref() + .map(|wc| wc.predicates.iter().cloned().collect()) + .unwrap_or_default(); + if let Some(outer) = &outer { + params.extend(outer.generics().params.iter().cloned()); + if let Some(wc) = &outer.generics().where_clause { + wheres.extend(wc.predicates.iter().cloned()); + } + } + + Ok(Self { + closure, + params, + wheres, + is_method: outer.is_some(), + }) } - let model_ty_params = fn_params_with_model_ty(&fn_params); - let body = &closure.body; + fn expand(&self) -> syn::Expr { + let mut fn_params: Vec = Vec::new(); + for param in &self.closure.inputs { + let syn::Pat::Type(pt) = param else { + let err = syn::Error::new_spanned( + param, + "invariant closure parameters must have explicit types, e.g. `|x: i64| ...`", + ); + return syn::Expr::Verbatim(err.to_compile_error()); + }; + let pat = &pt.pat; + let ty = &pt.ty; + fn_params.push(syn::parse_quote!(#pat: #ty)); + } - let id = COUNTER.fetch_add(1, Ordering::Relaxed); - let name = format_ident!("_thrust_invariant_{}", id); + let mut def_params: Vec = Vec::new(); + let mut def_wheres: Vec = + self.wheres.iter().map(|w| w.to_token_stream()).collect(); + let mut turbofish_args: Vec = Vec::new(); - quote! { - { + // `Self` in a method context: rewrite it to a synthetic generic, then + // pass the real `Self` via turbofish (legal in expression position). + let uses_self = self.is_method && tokens_contain_self(&self.closure.to_token_stream()); + if uses_self { + let synth: syn::Ident = format_ident!("__ThrustSelf"); + for param in &mut fn_params { + SelfRewriter { synth: &synth }.visit_fn_arg_mut(param); + } + def_params.push(quote!(#synth)); + def_wheres.push(quote!(#synth: thrust_models::Model)); + def_wheres.push(quote!(<#synth as thrust_models::Model>::Ty: PartialEq)); + turbofish_args.push(quote!(Self)); + } + + for param in &self.params { + def_params.push(param.to_token_stream()); + match param { + GenericParam::Type(tp) => { + let ident = &tp.ident; + def_wheres.push(quote!(#ident: thrust_models::Model)); + def_wheres.push(quote!(<#ident as thrust_models::Model>::Ty: PartialEq)); + turbofish_args.push(ident.to_token_stream()); + } + GenericParam::Const(cp) => turbofish_args.push(cp.ident.to_token_stream()), + GenericParam::Lifetime(_) => {} + } + } + + let model_ty_params = crate::fn_params_with_model_ty(&fn_params); + let body = &self.closure.body; + + let id = COUNTER.fetch_add(1, Ordering::Relaxed); + let name = format_ident!("_thrust_invariant_{}", id); + + let def_generics = if def_params.is_empty() { + quote!() + } else { + quote!(<#(#def_params),*>) + }; + let where_clause = if def_wheres.is_empty() { + quote!() + } else { + quote!(where #(#def_wheres),*) + }; + let turbofish = if turbofish_args.is_empty() { + quote!() + } else { + quote!(::<#(#turbofish_args),*>) + }; + + syn::parse_quote!({ #[allow(unused_variables)] #[allow(non_snake_case)] #[thrust::formula_fn] - fn #name(#model_ty_params) -> bool { + fn #name #def_generics(#model_ty_params) -> bool #where_clause { #body } - thrust_models::__invariant_marker(#name) + thrust_models::__invariant_marker(#name #turbofish) + }) + } +} + +/// Reads the optional `#[thrust::_outer_context(..)]` attribute threaded onto a +/// context-carrying invariant by `invariant_context`. +fn extract_outer_context(attrs: &[syn::Attribute]) -> syn::Result> { + let outer_context_path: syn::Path = syn::parse_quote!(thrust::_outer_context); + let mut outer_context = None; + for attr in attrs { + if attr.path() != &outer_context_path { + continue; + } + if outer_context.is_some() { + return Err(syn::Error::new_spanned( + attr, + "multiple _outer_context attributes found; expected at most one", + )); } + outer_context = Some(attr.parse_args()?); } - .into() + Ok(outer_context) +} + +/// Rewrites a bare `Self` type path to a synthetic type parameter, so the type +/// can be named inside a nested formula function. Qualified paths such as +/// `Self::Assoc` are left untouched (and are not supported in invariants). +struct SelfRewriter<'a> { + synth: &'a syn::Ident, +} + +impl VisitMut for SelfRewriter<'_> { + fn visit_path_mut(&mut self, path: &mut syn::Path) { + syn::visit_mut::visit_path_mut(self, path); + if path.leading_colon.is_none() + && path.segments.len() == 1 + && path.segments[0].ident == "Self" + { + path.segments[0].ident = self.synth.clone(); + } + } +} + +fn tokens_contain_self(tokens: &TokenStream2) -> bool { + tokens.clone().into_iter().any(|tt| match tt { + TokenTree::Ident(ident) => ident == "Self", + TokenTree::Group(group) => tokens_contain_self(&group.stream()), + _ => false, + }) } diff --git a/thrust-macros/src/invariant_context.rs b/thrust-macros/src/invariant_context.rs index 1d3d4ab..ac46c59 100644 --- a/thrust-macros/src/invariant_context.rs +++ b/thrust-macros/src/invariant_context.rs @@ -1,41 +1,37 @@ //! Expansion of `#[thrust_macros::invariant_context]`. //! //! Threads the surrounding generic context (and, in methods, `Self`) into -//! every `thrust_macros::invariant!(...)` call inside the annotated item. Each -//! such call is replaced in-place by an expanded `#[thrust::formula_fn]` plus -//! a marker call, with the threaded generics declared on the formula function -//! and instantiated via turbofish. In methods, `Self` is re-declared as a -//! synthetic type parameter and instantiated with the real `Self` the same -//! way, so an invariant may refer to generic- and `Self`-typed variables that -//! the standalone `invariant!` macro cannot see. +//! every `thrust_macros::invariant!(...)` call inside the annotated item, so an +//! invariant may refer to generic- and `Self`-typed variables that the +//! standalone `invariant!` macro cannot see. //! -//! The host function's where clause gains `T: Model` and +//! This macro does **not** expand the invariants itself: it only injects the +//! context, by rewriting each call into the context-carrying form that +//! `invariant!` knows how to expand (a `fn` header carrying the in-scope +//! generics/where clause, plus a `#[thrust::_outer_context(..)]` attribute +//! carrying the enclosing `impl`/`trait` header for methods). +//! +//! It also extends the host function's where clause with `T: Model` and //! `::Ty: PartialEq` predicates for every in-scope type parameter -//! (and for `Self` when used), as the generated marker call instantiates a -//! `Model`-bounded function. - -use std::sync::atomic::{AtomicUsize, Ordering}; +//! (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 proc_macro2::{TokenStream as TokenStream2, TokenTree}; -use quote::{format_ident, quote, ToTokens}; -use syn::{parse_macro_input, visit_mut::VisitMut, FnArg, GenericParam, Generics, WherePredicate}; - -use crate::{fn_params_with_model_ty, FnOuterItem}; +use quote::{quote, ToTokens}; +use syn::{visit_mut::VisitMut, GenericParam, Generics}; -static COUNTER: AtomicUsize = AtomicUsize::new(0); +use crate::FnOuterItem; pub(super) fn expand(item: TokenStream) -> TokenStream { - let mut item = parse_macro_input!(item as syn::Item); + let mut item = syn::parse_macro_input!(item as syn::Item); match &mut item { syn::Item::Fn(item_fn) => { let generics = item_fn.sig.generics.clone(); - let mut rewriter = InvariantRewriter::new(&generics, None); - rewriter.visit_block_mut(&mut item_fn.block); - if rewriter.found { - inject_model_bounds(&mut item_fn.sig.generics, None, rewriter.self_used); - } + let mut injector = InvariantInjector::new(&generics, None); + injector.visit_block_mut(&mut item_fn.block); + injector.inject_model_bounds(&mut item_fn.sig.generics); } syn::Item::Impl(item_impl) => { let outer = FnOuterItem::ItemImpl(item_impl.clone()).into_header_only(); @@ -44,11 +40,9 @@ pub(super) fn expand(item: TokenStream) -> TokenStream { continue; }; let generics = method.sig.generics.clone(); - let mut rewriter = InvariantRewriter::new(&generics, Some(&outer)); - rewriter.visit_block_mut(&mut method.block); - if rewriter.found { - inject_model_bounds(&mut method.sig.generics, Some(&outer), rewriter.self_used); - } + let mut injector = InvariantInjector::new(&generics, Some(&outer)); + injector.visit_block_mut(&mut method.block); + injector.inject_model_bounds(&mut method.sig.generics); } } syn::Item::Trait(item_trait) => { @@ -61,11 +55,9 @@ pub(super) fn expand(item: TokenStream) -> TokenStream { continue; }; let generics = method.sig.generics.clone(); - let mut rewriter = InvariantRewriter::new(&generics, Some(&outer)); - rewriter.visit_block_mut(block); - if rewriter.found { - inject_model_bounds(&mut method.sig.generics, Some(&outer), rewriter.self_used); - } + let mut injector = InvariantInjector::new(&generics, Some(&outer)); + injector.visit_block_mut(block); + injector.inject_model_bounds(&mut method.sig.generics); } } _ => { @@ -82,212 +74,90 @@ pub(super) fn expand(item: TokenStream) -> TokenStream { item.into_token_stream().into() } -struct InvariantRewriter { - /// In-scope generic params (the function's own, plus the outer impl/trait's - /// own for a method). - params: Vec, - /// Combined where-predicates from the function and (for a method) the - /// outer impl/trait. `Model`/`PartialEq` bounds for `params` are added - /// when constructing the formula function so each turbofish argument - /// satisfies them. - wheres: Vec, - is_method: bool, - found: bool, +struct InvariantInjector<'a> { + /// The enclosing function's own generics. + fn_generics: &'a Generics, + /// The enclosing `impl`/`trait` header, for a method. + outer: Option<&'a FnOuterItem>, + /// Whether any rewritten invariant references `Self`. self_used: bool, } -impl InvariantRewriter { - fn new(fn_generics: &Generics, outer: Option<&FnOuterItem>) -> Self { - let mut params: Vec = fn_generics.params.iter().cloned().collect(); - let mut wheres: Vec = fn_generics - .where_clause - .as_ref() - .map(|wc| wc.predicates.iter().cloned().collect()) - .unwrap_or_default(); - if let Some(outer) = outer { - params.extend(outer.generics().params.iter().cloned()); - if let Some(wc) = &outer.generics().where_clause { - wheres.extend(wc.predicates.iter().cloned()); - } - } +impl<'a> InvariantInjector<'a> { + fn new(fn_generics: &'a Generics, outer: Option<&'a FnOuterItem>) -> Self { Self { - params, - wheres, - is_method: outer.is_some(), - found: false, + fn_generics, + outer, self_used: false, } } - /// Builds the expansion for a single `invariant!` macro call. - fn expand_invariant(&mut self, mac: &syn::Macro) -> syn::Expr { - let closure: syn::ExprClosure = match syn::parse2(mac.tokens.clone()) { - Ok(c) => c, - Err(e) => return syn::Expr::Verbatim(e.to_compile_error()), - }; - - let mut fn_params: Vec = Vec::new(); - for param in &closure.inputs { - let syn::Pat::Type(pt) = param else { - let err = syn::Error::new_spanned( - param, - "invariant closure parameters must have explicit types, e.g. `|x: i64| ...`", - ); - return syn::Expr::Verbatim(err.to_compile_error()); - }; - let pat = &pt.pat; - let ty = &pt.ty; - fn_params.push(syn::parse_quote!(#pat: #ty)); - } - - let mut def_params: Vec = Vec::new(); - let mut def_wheres: Vec = - self.wheres.iter().map(|w| w.to_token_stream()).collect(); - let mut turbofish_args: Vec = Vec::new(); - - // `Self` in a method context: rewrite to a synthetic generic, then - // pass the real `Self` via turbofish (legal in expression position). - let uses_self = self.is_method && tokens_contain_self(&closure.to_token_stream()); - if uses_self { - self.self_used = true; - let synth: syn::Ident = format_ident!("__ThrustSelf"); - for param in &mut fn_params { - SelfRewriter { synth: &synth }.visit_fn_arg_mut(param); + /// Rewrites the closure inside `invariant!(CLOSURE)` into the + /// context-carrying form that `invariant!` expands using the threaded + /// context: a `fn` header carrying the in-scope generics/where clause, + /// tagged (for methods) with the enclosing `impl`/`trait` header. + fn inject_context(&self, closure: &TokenStream2) -> TokenStream2 { + let generics = self.fn_generics; + let where_clause = &self.fn_generics.where_clause; + let outer_attr = self + .outer + .map(|outer| quote!(#[thrust::_outer_context(#outer)])); + + quote! { + #outer_attr + fn _thrust_invariant_context #generics () #where_clause { + #closure } - def_params.push(quote!(#synth)); - def_wheres.push(quote!(#synth: thrust_models::Model)); - def_wheres.push(quote!(<#synth as thrust_models::Model>::Ty: PartialEq)); - turbofish_args.push(quote!(Self)); } - - for param in &self.params { - def_params.push(param.to_token_stream()); - match param { - GenericParam::Type(tp) => { - let ident = &tp.ident; - def_wheres.push(quote!(#ident: thrust_models::Model)); - def_wheres.push(quote!(<#ident as thrust_models::Model>::Ty: PartialEq)); - turbofish_args.push(ident.to_token_stream()); - } - GenericParam::Const(cp) => turbofish_args.push(cp.ident.to_token_stream()), - GenericParam::Lifetime(_) => {} - } - } - - let model_ty_params = fn_params_with_model_ty(&fn_params); - let body = &closure.body; - - let id = COUNTER.fetch_add(1, Ordering::Relaxed); - let name = format_ident!("_thrust_invariant_{}", id); - - let def_generics = if def_params.is_empty() { - quote!() - } else { - quote!(<#(#def_params),*>) - }; - let where_clause = if def_wheres.is_empty() { - quote!() - } else { - quote!(where #(#def_wheres),*) - }; - let turbofish = if turbofish_args.is_empty() { - quote!() - } else { - quote!(::<#(#turbofish_args),*>) - }; - - syn::parse_quote!({ - #[allow(unused_variables)] - #[allow(non_snake_case)] - #[thrust::formula_fn] - fn #name #def_generics(#model_ty_params) -> bool #where_clause { - #body - } - - thrust_models::__invariant_marker(#name #turbofish) - }) } -} -impl VisitMut for InvariantRewriter { - fn visit_stmt_mut(&mut self, stmt: &mut syn::Stmt) { - syn::visit_mut::visit_stmt_mut(self, stmt); - if let syn::Stmt::Macro(stmt_macro) = stmt { - if is_invariant_macro(&stmt_macro.mac.path) { - self.found = true; - let expanded = self.expand_invariant(&stmt_macro.mac); - *stmt = syn::Stmt::Expr(expanded, stmt_macro.semi_token); - } + /// Adds `T: Model` and `::Ty: PartialEq` bounds for every type + /// parameter in scope (the host function's own, the outer `impl`/`trait`'s, + /// and `Self` when an invariant names it) to the host's where clause. + fn inject_model_bounds(&self, generics: &mut Generics) { + let mut tys: Vec = type_param_idents(self.fn_generics); + if let Some(outer) = self.outer { + tys.extend(type_param_idents(outer.generics())); } - } - - fn visit_expr_mut(&mut self, expr: &mut syn::Expr) { - syn::visit_mut::visit_expr_mut(self, expr); - if let syn::Expr::Macro(expr_macro) = expr { - if is_invariant_macro(&expr_macro.mac.path) { - self.found = true; - *expr = self.expand_invariant(&expr_macro.mac); - } + if self.self_used { + tys.push(quote!(Self)); + } + if tys.is_empty() { + return; + } + let where_clause = generics.make_where_clause(); + for ty in tys { + where_clause + .predicates + .push(syn::parse_quote!(#ty: thrust_models::Model)); + where_clause + .predicates + .push(syn::parse_quote!(<#ty as thrust_models::Model>::Ty: PartialEq)); } } } -/// Rewrites a bare `Self` type path to a synthetic type parameter, so the type -/// can be named inside a nested formula function. Qualified paths such as -/// `Self::Assoc` are left untouched (and are not supported in invariants). -struct SelfRewriter<'a> { - synth: &'a syn::Ident, -} - -impl VisitMut for SelfRewriter<'_> { - fn visit_path_mut(&mut self, path: &mut syn::Path) { - syn::visit_mut::visit_path_mut(self, path); - if path.leading_colon.is_none() - && path.segments.len() == 1 - && path.segments[0].ident == "Self" - { - path.segments[0].ident = self.synth.clone(); +impl VisitMut for InvariantInjector<'_> { + fn visit_macro_mut(&mut self, mac: &mut syn::Macro) { + if !is_invariant_macro(&mac.path) { + return; + } + if self.outer.is_some() && tokens_contain_self(&mac.tokens) { + self.self_used = true; } + mac.tokens = self.inject_context(&mac.tokens); } } -/// Adds `T: Model` and `::Ty: PartialEq` bounds for every type -/// parameter in scope to a function's where clause. Each marker call generated -/// by the rewriter instantiates a `Model`-bounded formula function, so the -/// host function must itself satisfy those bounds. When an invariant names -/// `Self`, the rewriter instantiates the formula function with `Self`, so the -/// same bounds are added for `Self` (`with_self`). -fn inject_model_bounds(generics: &mut Generics, outer: Option<&FnOuterItem>, with_self: bool) { - let mut tys: Vec = generics +fn type_param_idents(generics: &Generics) -> Vec { + generics .params .iter() .filter_map(|p| match p { GenericParam::Type(tp) => Some(tp.ident.to_token_stream()), _ => None, }) - .collect(); - if let Some(outer) = outer { - for param in &outer.generics().params { - if let GenericParam::Type(tp) = param { - tys.push(tp.ident.to_token_stream()); - } - } - } - if with_self { - tys.push(quote!(Self)); - } - if tys.is_empty() { - return; - } - let where_clause = generics.make_where_clause(); - for ty in tys { - where_clause - .predicates - .push(syn::parse_quote!(#ty: thrust_models::Model)); - where_clause - .predicates - .push(syn::parse_quote!(<#ty as thrust_models::Model>::Ty: PartialEq)); - } + .collect() } fn is_invariant_macro(path: &syn::Path) -> bool { From cea08a5956a14976dc6cb2dec9204b052bfdee91 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 2 Jun 2026 16:00:24 +0000 Subject: [PATCH 03/13] thrust-macros: route threaded invariants through a dedicated macro Rather than having invariant! accept two different input shapes (a bare closure, or a context-carrying fn header), introduce a dedicated internal macro thrust_macros::_invariant_with_context! for the threaded form. - invariant! now only accepts a bare predicate closure. - _invariant_with_context! only accepts the context-carrying fn header. - invariant_context rewrites each invariant!(CLOSURE) into _invariant_with_context!( { CLOSURE }) by swapping the macro path as well as injecting the context tokens. Both macros share Invariant::expand, so there is still a single expansion implementation. The leading underscore follows the existing _requires_ensures convention for macros emitted by other macros and not written by hand. --- thrust-macros/src/invariant.rs | 99 +++++++++++++------------- thrust-macros/src/invariant_context.rs | 18 ++--- thrust-macros/src/lib.rs | 12 +++- 3 files changed, 69 insertions(+), 60 deletions(-) diff --git a/thrust-macros/src/invariant.rs b/thrust-macros/src/invariant.rs index 24bdd24..de070ac 100644 --- a/thrust-macros/src/invariant.rs +++ b/thrust-macros/src/invariant.rs @@ -1,52 +1,70 @@ -//! Expansion of `thrust_macros::invariant!`. +//! Expansion of `thrust_macros::invariant!` and its context-carrying sibling +//! `thrust_macros::_invariant_with_context!`. //! -//! Expands a predicate closure with explicit parameter types into a +//! Both expand a predicate closure with explicit parameter types into a //! `#[thrust::formula_fn]` over `Model::Ty` parameters plus a marker call -//! referencing it. +//! referencing it; they share [`Invariant::expand`] and differ only in input: //! -//! A standalone `invariant!(|x: i64| x >= 1)` only sees concrete types. To -//! refer to generic- or `Self`-typed variables, the enclosing item must be -//! annotated with `#[thrust_macros::invariant_context]`, which rewrites each -//! call into a context-carrying form that threads the in-scope generics (and, -//! in methods, `Self`) into the macro input: +//! - `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]` +//! rewrites each `invariant!` it finds into this form, threading the in-scope +//! generics (and, in methods, `Self`) in as an ordinary `fn` header whose body +//! is the predicate closure: //! -//! ```ignore -//! invariant!( -//! #[thrust::_outer_context(impl Foo where ..)] // methods only -//! fn _ctx() where .. { -//! |x: T, v: T| x == v -//! } -//! ) -//! ``` +//! ```ignore +//! _invariant_with_context!( +//! #[thrust::_outer_context(impl Foo where ..)] // methods only +//! fn _ctx() where .. { +//! |x: T, v: T| x == v +//! } +//! ) +//! ``` //! -//! This macro re-declares the threaded generics (shadowing the enclosing ones) -//! on the formula function and instantiates it via turbofish; in methods, -//! `Self` is re-declared as a synthetic type parameter and instantiated with -//! the real `Self` (legal in expression position). +//! The threaded generics (shadowing the enclosing ones) are re-declared on the +//! formula function and instantiated via turbofish; in methods, `Self` is +//! re-declared as a synthetic type parameter and instantiated with the real +//! `Self` (legal in expression position). use std::sync::atomic::{AtomicUsize, Ordering}; use proc_macro::TokenStream; use proc_macro2::{TokenStream as TokenStream2, TokenTree}; use quote::{format_ident, quote, ToTokens}; -use syn::{ - parse::{Parse, ParseStream}, - parse_macro_input, - visit_mut::VisitMut, - FnArg, GenericParam, WherePredicate, -}; +use syn::{parse_macro_input, visit_mut::VisitMut, FnArg, GenericParam, WherePredicate}; use crate::FnOuterItem; static COUNTER: AtomicUsize = AtomicUsize::new(0); +/// Expands `invariant!(CLOSURE)`: a bare predicate closure with no threaded +/// context. pub fn expand(input: TokenStream) -> TokenStream { - let invariant = parse_macro_input!(input as Invariant); - invariant.expand().into_token_stream().into() + let closure = parse_macro_input!(input as syn::ExprClosure); + Invariant { + closure, + params: Vec::new(), + wheres: Vec::new(), + is_method: false, + } + .expand() + .into_token_stream() + .into() +} + +/// Expands `_invariant_with_context!(CONTEXT_FN)`, the form +/// `#[thrust_macros::invariant_context]` rewrites each `invariant!` into. +pub fn expand_with_context(input: TokenStream) -> TokenStream { + let ctx_fn = parse_macro_input!(input as syn::ItemFn); + match Invariant::from_context_fn(ctx_fn) { + Ok(invariant) => invariant.expand().into_token_stream().into(), + Err(e) => e.to_compile_error().into(), + } } -/// The parsed `invariant!` input: the predicate closure plus the enclosing -/// generic context threaded in by `#[thrust_macros::invariant_context]`. +/// The closure to expand plus the enclosing generic context threaded in by +/// `#[thrust_macros::invariant_context]`. struct Invariant { closure: syn::ExprClosure, /// In-scope generic params (the enclosing function's own, plus the outer @@ -59,27 +77,6 @@ struct Invariant { is_method: bool, } -impl Parse for Invariant { - fn parse(input: ParseStream) -> syn::Result { - // Context-carrying form injected by `invariant_context`: a function - // header whose generics/where clause carry the enclosing context and - // whose body is the predicate closure. It always starts with the - // `#[thrust::_outer_context(..)]` attribute (methods) or `fn` (free - // functions); a standalone predicate closure starts with neither. - if input.peek(syn::Token![#]) || input.peek(syn::Token![fn]) { - let ctx_fn: syn::ItemFn = input.parse()?; - return Self::from_context_fn(ctx_fn); - } - - Ok(Self { - closure: input.parse()?, - params: Vec::new(), - wheres: Vec::new(), - is_method: false, - }) - } -} - impl Invariant { fn from_context_fn(ctx_fn: syn::ItemFn) -> syn::Result { let outer = extract_outer_context(&ctx_fn.attrs)?; diff --git a/thrust-macros/src/invariant_context.rs b/thrust-macros/src/invariant_context.rs index ac46c59..81cbe09 100644 --- a/thrust-macros/src/invariant_context.rs +++ b/thrust-macros/src/invariant_context.rs @@ -6,10 +6,11 @@ //! standalone `invariant!` macro cannot see. //! //! This macro does **not** expand the invariants itself: it only injects the -//! context, by rewriting each call into the context-carrying form that -//! `invariant!` knows how to expand (a `fn` header carrying the in-scope -//! generics/where clause, plus a `#[thrust::_outer_context(..)]` attribute -//! carrying the enclosing `impl`/`trait` header for methods). +//! context, by rewriting each `invariant!(..)` call into a +//! `thrust_macros::_invariant_with_context!(..)` call carrying that context (a +//! `fn` header carrying the in-scope generics/where clause, plus a +//! `#[thrust::_outer_context(..)]` attribute carrying the enclosing +//! `impl`/`trait` header for methods). //! //! It also extends the host function's where clause with `T: Model` and //! `::Ty: PartialEq` predicates for every in-scope type parameter @@ -92,10 +93,10 @@ impl<'a> InvariantInjector<'a> { } } - /// Rewrites the closure inside `invariant!(CLOSURE)` into the - /// context-carrying form that `invariant!` expands using the threaded - /// context: a `fn` header carrying the in-scope generics/where clause, - /// tagged (for methods) with the enclosing `impl`/`trait` header. + /// Builds the context-carrying argument for `_invariant_with_context!` from + /// the closure inside an `invariant!(CLOSURE)` call: a `fn` header carrying + /// the in-scope generics/where clause, tagged (for methods) with the + /// enclosing `impl`/`trait` header. fn inject_context(&self, closure: &TokenStream2) -> TokenStream2 { let generics = self.fn_generics; let where_clause = &self.fn_generics.where_clause; @@ -146,6 +147,7 @@ impl VisitMut for InvariantInjector<'_> { self.self_used = true; } mac.tokens = self.inject_context(&mac.tokens); + mac.path = syn::parse_quote!(::thrust_macros::_invariant_with_context); } } diff --git a/thrust-macros/src/lib.rs b/thrust-macros/src/lib.rs index 5be5be4..c4eed46 100644 --- a/thrust-macros/src/lib.rs +++ b/thrust-macros/src/lib.rs @@ -31,10 +31,20 @@ pub fn invariant(input: TokenStream) -> TokenStream { invariant::expand(input) } +/// Context-carrying counterpart of `invariant!`, emitted by +/// `#[thrust_macros::invariant_context]`. Not intended to be written by hand: +/// it takes a `fn` header carrying the threaded generics/where clause whose +/// body is the predicate closure (see [`invariant`]). +#[proc_macro] +pub fn _invariant_with_context(input: TokenStream) -> TokenStream { + invariant::expand_with_context(input) +} + /// Threads the surrounding generic context (and, in methods, `Self`) into /// every `thrust_macros::invariant!(...)` inside the annotated function, impl, /// or trait, so an invariant may refer to generic- and `Self`-typed variables -/// that the standalone `invariant!` macro cannot see. +/// that the standalone `invariant!` macro cannot see. Each such call is +/// rewritten into `thrust_macros::_invariant_with_context!`. #[proc_macro_attribute] pub fn invariant_context(_attr: TokenStream, item: TokenStream) -> TokenStream { invariant_context::expand(item) From a46eeef96c3465196fd96f7788982a923f1faa97 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 2 Jun 2026 16:07:55 +0000 Subject: [PATCH 04/13] thrust-macros: model invariant context as Option Replace the Invariant struct, which represented a standalone invariant as empty param/where vecs plus is_method=false, with an explicit Option<&Context> argument to a shared expand_invariant function. None now unambiguously means "no threaded context" (standalone invariant!), distinct from a contextful invariant that merely has no generics. --- thrust-macros/src/invariant.rs | 202 +++++++++++++++++---------------- 1 file changed, 103 insertions(+), 99 deletions(-) diff --git a/thrust-macros/src/invariant.rs b/thrust-macros/src/invariant.rs index de070ac..d4a00e3 100644 --- a/thrust-macros/src/invariant.rs +++ b/thrust-macros/src/invariant.rs @@ -3,7 +3,7 @@ //! //! Both expand a predicate closure with explicit parameter types into a //! `#[thrust::formula_fn]` over `Model::Ty` parameters plus a marker call -//! referencing it; they share [`Invariant::expand`] and differ only in input: +//! referencing it; they share [`expand_invariant`] and differ only in input: //! //! - `invariant!(|x: i64| x >= 1)` takes a bare predicate closure and only sees //! concrete types. @@ -42,33 +42,26 @@ static COUNTER: AtomicUsize = AtomicUsize::new(0); /// context. pub fn expand(input: TokenStream) -> TokenStream { let closure = parse_macro_input!(input as syn::ExprClosure); - Invariant { - closure, - params: Vec::new(), - wheres: Vec::new(), - is_method: false, - } - .expand() - .into_token_stream() - .into() + expand_invariant(&closure, None).into_token_stream().into() } /// Expands `_invariant_with_context!(CONTEXT_FN)`, the form /// `#[thrust_macros::invariant_context]` rewrites each `invariant!` into. pub fn expand_with_context(input: TokenStream) -> TokenStream { let ctx_fn = parse_macro_input!(input as syn::ItemFn); - match Invariant::from_context_fn(ctx_fn) { - Ok(invariant) => invariant.expand().into_token_stream().into(), + match Context::from_context_fn(ctx_fn) { + Ok((closure, context)) => expand_invariant(&closure, Some(&context)) + .into_token_stream() + .into(), Err(e) => e.to_compile_error().into(), } } -/// The closure to expand plus the enclosing generic context threaded in by -/// `#[thrust_macros::invariant_context]`. -struct Invariant { - closure: syn::ExprClosure, - /// In-scope generic params (the enclosing function's own, plus the outer - /// `impl`/`trait`'s for a method). Empty for a standalone `invariant!`. +/// The enclosing generic context threaded into an invariant by +/// `#[thrust_macros::invariant_context]`. A standalone `invariant!` has none. +struct Context { + /// In-scope generic params: the enclosing function's own, plus the outer + /// `impl`/`trait`'s for a method. params: Vec, /// Combined where-predicates from the function and (for a method) the /// outer `impl`/`trait`. @@ -77,8 +70,10 @@ struct Invariant { is_method: bool, } -impl Invariant { - fn from_context_fn(ctx_fn: syn::ItemFn) -> syn::Result { +impl Context { + /// Parses the `fn` header that carries the threaded context, returning the + /// predicate closure from its body alongside the recovered context. + fn from_context_fn(ctx_fn: syn::ItemFn) -> syn::Result<(syn::ExprClosure, Self)> { let outer = extract_outer_context(&ctx_fn.attrs)?; let closure = match ctx_fn.block.stmts.as_slice() { @@ -106,95 +101,104 @@ impl Invariant { } } - Ok(Self { + Ok(( closure, - params, - wheres, - is_method: outer.is_some(), - }) + Self { + params, + wheres, + is_method: outer.is_some(), + }, + )) } +} - fn expand(&self) -> syn::Expr { - let mut fn_params: Vec = Vec::new(); - for param in &self.closure.inputs { - let syn::Pat::Type(pt) = param else { - let err = syn::Error::new_spanned( - param, - "invariant closure parameters must have explicit types, e.g. `|x: i64| ...`", - ); - return syn::Expr::Verbatim(err.to_compile_error()); - }; - let pat = &pt.pat; - let ty = &pt.ty; - fn_params.push(syn::parse_quote!(#pat: #ty)); - } +/// Expands a predicate closure into a `#[thrust::formula_fn]` plus a marker +/// call. With `context`, the threaded generics (and, in methods, `Self`) are +/// re-declared on the formula function and instantiated via turbofish. +fn expand_invariant(closure: &syn::ExprClosure, context: Option<&Context>) -> syn::Expr { + let mut fn_params: Vec = Vec::new(); + for param in &closure.inputs { + let syn::Pat::Type(pt) = param else { + let err = syn::Error::new_spanned( + param, + "invariant closure parameters must have explicit types, e.g. `|x: i64| ...`", + ); + return syn::Expr::Verbatim(err.to_compile_error()); + }; + let pat = &pt.pat; + let ty = &pt.ty; + fn_params.push(syn::parse_quote!(#pat: #ty)); + } - let mut def_params: Vec = Vec::new(); - let mut def_wheres: Vec = - self.wheres.iter().map(|w| w.to_token_stream()).collect(); - let mut turbofish_args: Vec = Vec::new(); - - // `Self` in a method context: rewrite it to a synthetic generic, then - // pass the real `Self` via turbofish (legal in expression position). - let uses_self = self.is_method && tokens_contain_self(&self.closure.to_token_stream()); - if uses_self { - let synth: syn::Ident = format_ident!("__ThrustSelf"); - for param in &mut fn_params { - SelfRewriter { synth: &synth }.visit_fn_arg_mut(param); - } - def_params.push(quote!(#synth)); - def_wheres.push(quote!(#synth: thrust_models::Model)); - def_wheres.push(quote!(<#synth as thrust_models::Model>::Ty: PartialEq)); - turbofish_args.push(quote!(Self)); + let mut def_params: Vec = Vec::new(); + let mut def_wheres: Vec = context + .iter() + .flat_map(|ctx| ctx.wheres.iter()) + .map(|w| w.to_token_stream()) + .collect(); + let mut turbofish_args: Vec = Vec::new(); + + // `Self` in a method context: rewrite it to a synthetic generic, then pass + // the real `Self` via turbofish (legal in expression position). + let uses_self = + context.is_some_and(|ctx| ctx.is_method) && tokens_contain_self(&closure.to_token_stream()); + if uses_self { + let synth: syn::Ident = format_ident!("__ThrustSelf"); + for param in &mut fn_params { + SelfRewriter { synth: &synth }.visit_fn_arg_mut(param); } + def_params.push(quote!(#synth)); + def_wheres.push(quote!(#synth: thrust_models::Model)); + def_wheres.push(quote!(<#synth as thrust_models::Model>::Ty: PartialEq)); + turbofish_args.push(quote!(Self)); + } - for param in &self.params { - def_params.push(param.to_token_stream()); - match param { - GenericParam::Type(tp) => { - let ident = &tp.ident; - def_wheres.push(quote!(#ident: thrust_models::Model)); - def_wheres.push(quote!(<#ident as thrust_models::Model>::Ty: PartialEq)); - turbofish_args.push(ident.to_token_stream()); - } - GenericParam::Const(cp) => turbofish_args.push(cp.ident.to_token_stream()), - GenericParam::Lifetime(_) => {} + for param in context.iter().flat_map(|ctx| ctx.params.iter()) { + def_params.push(param.to_token_stream()); + match param { + GenericParam::Type(tp) => { + let ident = &tp.ident; + def_wheres.push(quote!(#ident: thrust_models::Model)); + def_wheres.push(quote!(<#ident as thrust_models::Model>::Ty: PartialEq)); + turbofish_args.push(ident.to_token_stream()); } + GenericParam::Const(cp) => turbofish_args.push(cp.ident.to_token_stream()), + GenericParam::Lifetime(_) => {} } + } - let model_ty_params = crate::fn_params_with_model_ty(&fn_params); - let body = &self.closure.body; - - let id = COUNTER.fetch_add(1, Ordering::Relaxed); - let name = format_ident!("_thrust_invariant_{}", id); - - let def_generics = if def_params.is_empty() { - quote!() - } else { - quote!(<#(#def_params),*>) - }; - let where_clause = if def_wheres.is_empty() { - quote!() - } else { - quote!(where #(#def_wheres),*) - }; - let turbofish = if turbofish_args.is_empty() { - quote!() - } else { - quote!(::<#(#turbofish_args),*>) - }; - - syn::parse_quote!({ - #[allow(unused_variables)] - #[allow(non_snake_case)] - #[thrust::formula_fn] - fn #name #def_generics(#model_ty_params) -> bool #where_clause { - #body - } + let model_ty_params = crate::fn_params_with_model_ty(&fn_params); + let body = &closure.body; + + let id = COUNTER.fetch_add(1, Ordering::Relaxed); + let name = format_ident!("_thrust_invariant_{}", id); + + let def_generics = if def_params.is_empty() { + quote!() + } else { + quote!(<#(#def_params),*>) + }; + let where_clause = if def_wheres.is_empty() { + quote!() + } else { + quote!(where #(#def_wheres),*) + }; + let turbofish = if turbofish_args.is_empty() { + quote!() + } else { + quote!(::<#(#turbofish_args),*>) + }; + + syn::parse_quote!({ + #[allow(unused_variables)] + #[allow(non_snake_case)] + #[thrust::formula_fn] + fn #name #def_generics(#model_ty_params) -> bool #where_clause { + #body + } - thrust_models::__invariant_marker(#name #turbofish) - }) - } + thrust_models::__invariant_marker(#name #turbofish) + }) } /// Reads the optional `#[thrust::_outer_context(..)]` attribute threaded onto a From 109d9d4cfc587ca6030f08ab293c649dfd4cf004 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 2 Jun 2026 16:17:21 +0000 Subject: [PATCH 05/13] thrust-macros: share extract_outer_context via crate root invariant.rs and spec.rs each had their own copy of the _outer_context attribute reader. Hoist the shared attribute-reading logic into the crate root (alongside FnOuterItem and the other shared helpers); spec's wrapper now calls it and keeps only its method-specific mentions_self diagnostic. --- thrust-macros/src/invariant.rs | 24 +----------------------- thrust-macros/src/spec.rs | 17 +---------------- 2 files changed, 2 insertions(+), 39 deletions(-) diff --git a/thrust-macros/src/invariant.rs b/thrust-macros/src/invariant.rs index d4a00e3..0ae5db0 100644 --- a/thrust-macros/src/invariant.rs +++ b/thrust-macros/src/invariant.rs @@ -34,8 +34,6 @@ use proc_macro2::{TokenStream as TokenStream2, TokenTree}; use quote::{format_ident, quote, ToTokens}; use syn::{parse_macro_input, visit_mut::VisitMut, FnArg, GenericParam, WherePredicate}; -use crate::FnOuterItem; - static COUNTER: AtomicUsize = AtomicUsize::new(0); /// Expands `invariant!(CLOSURE)`: a bare predicate closure with no threaded @@ -74,7 +72,7 @@ impl Context { /// Parses the `fn` header that carries the threaded context, returning the /// predicate closure from its body alongside the recovered context. fn from_context_fn(ctx_fn: syn::ItemFn) -> syn::Result<(syn::ExprClosure, Self)> { - let outer = extract_outer_context(&ctx_fn.attrs)?; + let outer = crate::extract_outer_context(&ctx_fn.attrs)?; let closure = match ctx_fn.block.stmts.as_slice() { [syn::Stmt::Expr(syn::Expr::Closure(closure), _)] => closure.clone(), @@ -201,26 +199,6 @@ fn expand_invariant(closure: &syn::ExprClosure, context: Option<&Context>) -> sy }) } -/// Reads the optional `#[thrust::_outer_context(..)]` attribute threaded onto a -/// context-carrying invariant by `invariant_context`. -fn extract_outer_context(attrs: &[syn::Attribute]) -> syn::Result> { - let outer_context_path: syn::Path = syn::parse_quote!(thrust::_outer_context); - let mut outer_context = None; - for attr in attrs { - if attr.path() != &outer_context_path { - continue; - } - if outer_context.is_some() { - return Err(syn::Error::new_spanned( - attr, - "multiple _outer_context attributes found; expected at most one", - )); - } - outer_context = Some(attr.parse_args()?); - } - Ok(outer_context) -} - /// Rewrites a bare `Self` type path to a synthetic type parameter, so the type /// can be named inside a nested formula function. Qualified paths such as /// `Self::Assoc` are left untouched (and are not supported in invariants). diff --git a/thrust-macros/src/spec.rs b/thrust-macros/src/spec.rs index e049b33..98bd51f 100644 --- a/thrust-macros/src/spec.rs +++ b/thrust-macros/src/spec.rs @@ -235,22 +235,7 @@ fn extract_requires_ensures(func: &mut FnItemWithSignature) -> syn::Result<(syn: } fn extract_outer_context(func: &FnItemWithSignature) -> syn::Result> { - let outer_context_path: syn::Path = syn::parse_quote!(thrust::_outer_context); - let mut outer_context = None; - for attr in func.attrs() { - if attr.path() != &outer_context_path { - continue; - } - - let item = attr.parse_args()?; - if outer_context.is_some() { - return Err(syn::Error::new_spanned( - attr, - "multiple _outer_context attributes found; expected at most one", - )); - } - outer_context = Some(item); - } + let outer_context = crate::extract_outer_context(func.attrs())?; if mentions_self(func.sig()) && outer_context.is_none() { return Err(syn::Error::new_spanned( func.sig().ident.clone(), From a08e1259536cd6825fdc1a48f729ce996e9e0adf Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 2 Jun 2026 16:23:27 +0000 Subject: [PATCH 06/13] thrust-macros: invariant_context accepts only a function item invariant_context previously walked impl/trait blocks itself, duplicating the method-iteration that #[thrust_macros::context] already does. Restrict it to a single function: a method recovers its enclosing impl/trait context from the #[thrust::_outer_context(..)] attribute stamped by #[thrust_macros::context], the same mechanism requires/ensures use. Update the loop_invariant_self tests to combine #[context] on the impl with #[invariant_context] on the method accordingly. --- tests/ui/fail/loop_invariant_self.rs | 3 +- tests/ui/pass/loop_invariant_self.rs | 3 +- thrust-macros/src/invariant_context.rs | 74 +++++++++----------------- thrust-macros/src/lib.rs | 12 +++-- 4 files changed, 38 insertions(+), 54 deletions(-) diff --git a/tests/ui/fail/loop_invariant_self.rs b/tests/ui/fail/loop_invariant_self.rs index 04f775b..5ff43b6 100644 --- a/tests/ui/fail/loop_invariant_self.rs +++ b/tests/ui/fail/loop_invariant_self.rs @@ -11,8 +11,9 @@ impl thrust_models::Model for Counter { type Ty = (thrust_models::model::Int,); } -#[thrust_macros::invariant_context] +#[thrust_macros::context] impl Counter { + #[thrust_macros::invariant_context] fn run(self) { let mut c = self; let mut x = 1_i64; diff --git a/tests/ui/pass/loop_invariant_self.rs b/tests/ui/pass/loop_invariant_self.rs index 31cfb49..f24809d 100644 --- a/tests/ui/pass/loop_invariant_self.rs +++ b/tests/ui/pass/loop_invariant_self.rs @@ -11,8 +11,9 @@ impl thrust_models::Model for Counter { type Ty = (thrust_models::model::Int,); } -#[thrust_macros::invariant_context] +#[thrust_macros::context] impl Counter { + #[thrust_macros::invariant_context] fn run(self) { let mut c = self; let mut x = 1_i64; diff --git a/thrust-macros/src/invariant_context.rs b/thrust-macros/src/invariant_context.rs index 81cbe09..1d7a148 100644 --- a/thrust-macros/src/invariant_context.rs +++ b/thrust-macros/src/invariant_context.rs @@ -1,10 +1,23 @@ //! Expansion of `#[thrust_macros::invariant_context]`. //! //! Threads the surrounding generic context (and, in methods, `Self`) into -//! every `thrust_macros::invariant!(...)` call inside the annotated item, so an -//! invariant may refer to generic- and `Self`-typed variables that the +//! every `thrust_macros::invariant!(...)` call inside the annotated function, so +//! an invariant may refer to generic- and `Self`-typed variables that the //! standalone `invariant!` macro cannot see. //! +//! It is applied to a single function. A method recovers its enclosing +//! `impl`/`trait` header from the `#[thrust::_outer_context(..)]` attribute +//! stamped by `#[thrust_macros::context]`, the same mechanism `requires`/ +//! `ensures` use: +//! +//! ```ignore +//! #[thrust_macros::context] +//! impl Foo { +//! #[thrust_macros::invariant_context] +//! fn f(&self) { .. } +//! } +//! ``` +//! //! This macro does **not** expand the invariants itself: it only injects the //! context, by rewriting each `invariant!(..)` call into a //! `thrust_macros::_invariant_with_context!(..)` call carrying that context (a @@ -25,54 +38,19 @@ use syn::{visit_mut::VisitMut, GenericParam, Generics}; use crate::FnOuterItem; pub(super) fn expand(item: TokenStream) -> TokenStream { - let mut item = syn::parse_macro_input!(item as syn::Item); + let mut item_fn = syn::parse_macro_input!(item as syn::ItemFn); - match &mut item { - syn::Item::Fn(item_fn) => { - let generics = item_fn.sig.generics.clone(); - let mut injector = InvariantInjector::new(&generics, None); - injector.visit_block_mut(&mut item_fn.block); - injector.inject_model_bounds(&mut item_fn.sig.generics); - } - syn::Item::Impl(item_impl) => { - let outer = FnOuterItem::ItemImpl(item_impl.clone()).into_header_only(); - for impl_item in &mut item_impl.items { - let syn::ImplItem::Fn(method) = impl_item else { - continue; - }; - let generics = method.sig.generics.clone(); - let mut injector = InvariantInjector::new(&generics, Some(&outer)); - injector.visit_block_mut(&mut method.block); - injector.inject_model_bounds(&mut method.sig.generics); - } - } - syn::Item::Trait(item_trait) => { - let outer = FnOuterItem::ItemTrait(item_trait.clone()).into_header_only(); - for trait_item in &mut item_trait.items { - let syn::TraitItem::Fn(method) = trait_item else { - continue; - }; - let Some(block) = &mut method.default else { - continue; - }; - let generics = method.sig.generics.clone(); - let mut injector = InvariantInjector::new(&generics, Some(&outer)); - injector.visit_block_mut(block); - injector.inject_model_bounds(&mut method.sig.generics); - } - } - _ => { - return syn::Error::new_spanned( - &item, - "#[thrust_macros::invariant_context] expects a function, impl block, or trait \ - definition", - ) - .to_compile_error() - .into(); - } - } + let outer = match crate::extract_outer_context(&item_fn.attrs) { + Ok(outer) => outer, + Err(e) => return e.to_compile_error().into(), + }; + + let generics = item_fn.sig.generics.clone(); + let mut injector = InvariantInjector::new(&generics, outer.as_ref()); + injector.visit_block_mut(&mut item_fn.block); + injector.inject_model_bounds(&mut item_fn.sig.generics); - item.into_token_stream().into() + item_fn.into_token_stream().into() } struct InvariantInjector<'a> { diff --git a/thrust-macros/src/lib.rs b/thrust-macros/src/lib.rs index c4eed46..c23b043 100644 --- a/thrust-macros/src/lib.rs +++ b/thrust-macros/src/lib.rs @@ -41,10 +41,14 @@ pub fn _invariant_with_context(input: TokenStream) -> TokenStream { } /// Threads the surrounding generic context (and, in methods, `Self`) into -/// every `thrust_macros::invariant!(...)` inside the annotated function, impl, -/// or trait, so an invariant may refer to generic- and `Self`-typed variables -/// that the standalone `invariant!` macro cannot see. Each such call is -/// rewritten into `thrust_macros::_invariant_with_context!`. +/// every `thrust_macros::invariant!(...)` inside the annotated function, so an +/// invariant may refer to generic- and `Self`-typed variables that the +/// standalone `invariant!` macro cannot see. Each such call is rewritten into +/// `thrust_macros::_invariant_with_context!`. +/// +/// Applied to a single function; a method recovers its enclosing +/// `impl`/`trait` context from `#[thrust_macros::context]` on the surrounding +/// block, as `requires`/`ensures` do. #[proc_macro_attribute] pub fn invariant_context(_attr: TokenStream, item: TokenStream) -> TokenStream { invariant_context::expand(item) From 2774c11683eae24b67ff9599316966ea4c0611c6 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 2 Jun 2026 16:31:07 +0000 Subject: [PATCH 07/13] thrust-macros: share Model-predicate construction with spec The [T: Model, ::Ty: PartialEq] predicate pair was built by hand in three places (twice in spec, once in invariant_context). Hoist it into a shared crate-root helper model_predicates(ty); each caller keeps its own which-types-to-bound logic but no longer open-codes the predicate pair. --- thrust-macros/src/invariant_context.rs | 7 +------ thrust-macros/src/spec.rs | 7 ++----- 2 files changed, 3 insertions(+), 11 deletions(-) diff --git a/thrust-macros/src/invariant_context.rs b/thrust-macros/src/invariant_context.rs index 1d7a148..e8c5db0 100644 --- a/thrust-macros/src/invariant_context.rs +++ b/thrust-macros/src/invariant_context.rs @@ -106,12 +106,7 @@ impl<'a> InvariantInjector<'a> { } let where_clause = generics.make_where_clause(); for ty in tys { - where_clause - .predicates - .push(syn::parse_quote!(#ty: thrust_models::Model)); - where_clause - .predicates - .push(syn::parse_quote!(<#ty as thrust_models::Model>::Ty: PartialEq)); + where_clause.predicates.extend(crate::model_predicates(&ty)); } } } diff --git a/thrust-macros/src/spec.rs b/thrust-macros/src/spec.rs index 98bd51f..141b407 100644 --- a/thrust-macros/src/spec.rs +++ b/thrust-macros/src/spec.rs @@ -600,9 +600,7 @@ fn model_where_predicates( let mut predicates: Vec = Vec::new(); for param in &generic_type_params { - let ident = ¶m.ident; - predicates.push(syn::parse_quote!(#ident: thrust_models::Model)); - predicates.push(syn::parse_quote!(<#ident as thrust_models::Model>::Ty: PartialEq)); + predicates.extend(crate::model_predicates(¶m.ident)); } struct Visitor { @@ -641,8 +639,7 @@ fn model_where_predicates( } visitor.visit_return_type(&func.sig().output); for tp in visitor.generic_paths { - predicates.push(syn::parse_quote!(#tp: thrust_models::Model)); - predicates.push(syn::parse_quote!(<#tp as thrust_models::Model>::Ty: PartialEq)); + predicates.extend(crate::model_predicates(&tp)); } predicates From 43dd5969f8515f5aa5632a7fe82eabde46bc6518 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 2 Jun 2026 16:36:06 +0000 Subject: [PATCH 08/13] thrust-macros: exclude closure-typed params from invariant Model bounds spec skips Fn/FnMut/FnOnce-bounded type params when adding Model predicates (closures have no Model instance), but the invariant path bounded every type param unconditionally. A generic invariant inside a function with a closure- typed generic param therefore generated F: Model and failed at the call site with "the trait bound {closure}: Model is not satisfied". Share spec's Fn-bound check as a crate-root helper has_fn_bound and apply it in both the formula function (invariant.rs) and the host where clause (invariant_context.rs). Add a regression test. --- .../ui/pass/loop_invariant_generic_closure.rs | 26 +++++++++++++++++++ thrust-macros/src/invariant.rs | 17 ++++++------ thrust-macros/src/invariant_context.rs | 6 ++++- thrust-macros/src/spec.rs | 15 +---------- 4 files changed, 41 insertions(+), 23 deletions(-) create mode 100644 tests/ui/pass/loop_invariant_generic_closure.rs diff --git a/tests/ui/pass/loop_invariant_generic_closure.rs b/tests/ui/pass/loop_invariant_generic_closure.rs new file mode 100644 index 0000000..73668eb --- /dev/null +++ b/tests/ui/pass/loop_invariant_generic_closure.rs @@ -0,0 +1,26 @@ +//@check-pass +//@compile-flags: -C debug-assertions=off + +#[thrust_macros::requires(true)] +#[thrust_macros::ensures(true)] +#[thrust::trusted] +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] +fn keep i64, T: Copy + PartialEq>(f: F, v: T) { + let _ = f; + let mut x = v; + while rand() == 0 { + thrust_macros::invariant!(|x: T, v: T| x == v); + x = v; + } + assert!(x == v); +} + +fn main() { + keep(|x| x + 1, 0_i64); + keep(|x| x, true); +} diff --git a/thrust-macros/src/invariant.rs b/thrust-macros/src/invariant.rs index 0ae5db0..1935bb3 100644 --- a/thrust-macros/src/invariant.rs +++ b/thrust-macros/src/invariant.rs @@ -129,10 +129,10 @@ fn expand_invariant(closure: &syn::ExprClosure, context: Option<&Context>) -> sy } let mut def_params: Vec = Vec::new(); - let mut def_wheres: Vec = context + let mut def_wheres: Vec = context .iter() .flat_map(|ctx| ctx.wheres.iter()) - .map(|w| w.to_token_stream()) + .cloned() .collect(); let mut turbofish_args: Vec = Vec::new(); @@ -146,8 +146,7 @@ fn expand_invariant(closure: &syn::ExprClosure, context: Option<&Context>) -> sy SelfRewriter { synth: &synth }.visit_fn_arg_mut(param); } def_params.push(quote!(#synth)); - def_wheres.push(quote!(#synth: thrust_models::Model)); - def_wheres.push(quote!(<#synth as thrust_models::Model>::Ty: PartialEq)); + def_wheres.extend(crate::model_predicates(&synth)); turbofish_args.push(quote!(Self)); } @@ -155,10 +154,12 @@ fn expand_invariant(closure: &syn::ExprClosure, context: Option<&Context>) -> sy def_params.push(param.to_token_stream()); match param { GenericParam::Type(tp) => { - let ident = &tp.ident; - def_wheres.push(quote!(#ident: thrust_models::Model)); - def_wheres.push(quote!(<#ident as thrust_models::Model>::Ty: PartialEq)); - turbofish_args.push(ident.to_token_stream()); + // A closure-typed param has no `Model` instance; re-declare and + // pass it through, but do not bound it. + if !crate::has_fn_bound(&tp.bounds) { + def_wheres.extend(crate::model_predicates(&tp.ident)); + } + turbofish_args.push(tp.ident.to_token_stream()); } GenericParam::Const(cp) => turbofish_args.push(cp.ident.to_token_stream()), GenericParam::Lifetime(_) => {} diff --git a/thrust-macros/src/invariant_context.rs b/thrust-macros/src/invariant_context.rs index e8c5db0..fe9d8c4 100644 --- a/thrust-macros/src/invariant_context.rs +++ b/thrust-macros/src/invariant_context.rs @@ -124,12 +124,16 @@ impl VisitMut for InvariantInjector<'_> { } } +/// Type-parameter idents that need a `Model` bound: every type param except +/// closure-typed ones, which have no `Model` instance. fn type_param_idents(generics: &Generics) -> Vec { generics .params .iter() .filter_map(|p| match p { - GenericParam::Type(tp) => Some(tp.ident.to_token_stream()), + GenericParam::Type(tp) if !crate::has_fn_bound(&tp.bounds) => { + Some(tp.ident.to_token_stream()) + } _ => None, }) .collect() diff --git a/thrust-macros/src/spec.rs b/thrust-macros/src/spec.rs index 141b407..9bdd787 100644 --- a/thrust-macros/src/spec.rs +++ b/thrust-macros/src/spec.rs @@ -562,19 +562,6 @@ fn model_where_predicates( } } - impl GenericTypeParam { - fn has_fn_bound(&self) -> bool { - self.bounds.iter().any(|b| { - let TypeParamBound::Trait(tb) = b else { - return false; - }; - tb.path.segments.last().is_some_and(|s| { - matches!(s.ident.to_string().as_str(), "Fn" | "FnOnce" | "FnMut") - }) - }) - } - } - let mut generic_type_params: Vec = Vec::new(); for param in &func.sig().generics.params { let GenericParam::Type(tp) = param else { @@ -596,7 +583,7 @@ fn model_where_predicates( }); } } - generic_type_params.retain(|p| !p.has_fn_bound()); + generic_type_params.retain(|p| !crate::has_fn_bound(&p.bounds)); let mut predicates: Vec = Vec::new(); for param in &generic_type_params { From c9e17ce77595475bf79a195c5d22259edbfece84 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 3 Jun 2026 14:09:12 +0000 Subject: [PATCH 09/13] thrust-macros: reuse model_where_predicates; carry host signature verbatim Address three points: 1. Reuse the whole model_where_predicates (not just has_fn_bound). Move it to the crate root, retargeted at &Signature, and call it from both the formula function (invariant) and the host where clause (invariant_context). Each path adds Self handling on top (synthetic param vs Self: Model) but no longer open-codes generic collection or the Fn-bound exclusion. 2. Rename InvariantInjector to ContextInjector. 3. Change the _invariant_with_context input to '#outer_attr #signature; #closure'. ContextInjector now pastes the host function signature verbatim, and the expansion parses it back, so model_where_predicates sees the real signature (its generics, where clause, and parameter/return projections). --- thrust-macros/src/invariant.rs | 154 +++++++++++++------------ thrust-macros/src/invariant_context.rs | 81 +++++-------- thrust-macros/src/spec.rs | 98 +--------------- 3 files changed, 113 insertions(+), 220 deletions(-) diff --git a/thrust-macros/src/invariant.rs b/thrust-macros/src/invariant.rs index 1935bb3..0c30505 100644 --- a/thrust-macros/src/invariant.rs +++ b/thrust-macros/src/invariant.rs @@ -9,20 +9,20 @@ //! concrete types. //! - `_invariant_with_context!(..)` additionally carries the enclosing generic //! context. It is never written by hand: `#[thrust_macros::invariant_context]` -//! rewrites each `invariant!` it finds into this form, threading the in-scope -//! generics (and, in methods, `Self`) in as an ordinary `fn` header whose body -//! is the predicate closure: +//! 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 +//! closure: //! //! ```ignore //! _invariant_with_context!( //! #[thrust::_outer_context(impl Foo where ..)] // methods only -//! fn _ctx() where .. { -//! |x: T, v: T| x == v -//! } +//! fn f(..) -> .. where ..; // host signature, as-is +//! |x: T, v: T| x == v //! ) //! ``` //! -//! The threaded generics (shadowing the enclosing ones) are re-declared on the +//! The in-scope generics (shadowing the enclosing ones) are re-declared on the //! formula function and instantiated via turbofish; in methods, `Self` is //! re-declared as a synthetic type parameter and instantiated with the real //! `Self` (legal in expression position). @@ -32,7 +32,14 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use proc_macro::TokenStream; use proc_macro2::{TokenStream as TokenStream2, TokenTree}; use quote::{format_ident, quote, ToTokens}; -use syn::{parse_macro_input, visit_mut::VisitMut, FnArg, GenericParam, WherePredicate}; +use syn::{ + parse::{Parse, ParseStream}, + parse_macro_input, + visit_mut::VisitMut, + FnArg, GenericParam, Signature, WherePredicate, +}; + +use crate::FnOuterItem; static COUNTER: AtomicUsize = AtomicUsize::new(0); @@ -43,75 +50,70 @@ pub fn expand(input: TokenStream) -> TokenStream { expand_invariant(&closure, None).into_token_stream().into() } -/// Expands `_invariant_with_context!(CONTEXT_FN)`, the form +/// Expands `_invariant_with_context!(#outer_attr #sig; CLOSURE)`, the form /// `#[thrust_macros::invariant_context]` rewrites each `invariant!` into. pub fn expand_with_context(input: TokenStream) -> TokenStream { - let ctx_fn = parse_macro_input!(input as syn::ItemFn); - match Context::from_context_fn(ctx_fn) { - Ok((closure, context)) => expand_invariant(&closure, Some(&context)) - .into_token_stream() - .into(), - Err(e) => e.to_compile_error().into(), + let WithContext { closure, context } = parse_macro_input!(input as WithContext); + expand_invariant(&closure, Some(&context)) + .into_token_stream() + .into() +} + +/// The parsed `_invariant_with_context!` input. +struct WithContext { + context: Context, + closure: syn::ExprClosure, +} + +impl Parse for WithContext { + fn parse(input: ParseStream) -> syn::Result { + let attrs = input.call(syn::Attribute::parse_outer)?; + let outer = crate::extract_outer_context(&attrs)?; + let sig: Signature = input.parse()?; + input.parse::()?; + let closure: syn::ExprClosure = input.parse()?; + Ok(Self { + context: Context { sig, outer }, + closure, + }) } } -/// The enclosing generic context threaded into an invariant by -/// `#[thrust_macros::invariant_context]`. A standalone `invariant!` has none. +/// The enclosing context threaded into an invariant by +/// `#[thrust_macros::invariant_context]`: the host function signature and, for a +/// method, its `impl`/`trait` header. A standalone `invariant!` has none. struct Context { - /// In-scope generic params: the enclosing function's own, plus the outer - /// `impl`/`trait`'s for a method. - params: Vec, - /// Combined where-predicates from the function and (for a method) the - /// outer `impl`/`trait`. - wheres: Vec, - /// Whether `Self` is nameable at the call site (i.e. we are in a method). - is_method: bool, + sig: Signature, + outer: Option, } impl Context { - /// Parses the `fn` header that carries the threaded context, returning the - /// predicate closure from its body alongside the recovered context. - fn from_context_fn(ctx_fn: syn::ItemFn) -> syn::Result<(syn::ExprClosure, Self)> { - let outer = crate::extract_outer_context(&ctx_fn.attrs)?; - - let closure = match ctx_fn.block.stmts.as_slice() { - [syn::Stmt::Expr(syn::Expr::Closure(closure), _)] => closure.clone(), - _ => { - return Err(syn::Error::new_spanned( - &ctx_fn.block, - "invariant context body must be a single predicate closure", - )) - } - }; + fn is_method(&self) -> bool { + self.outer.is_some() + } - let mut params: Vec = ctx_fn.sig.generics.params.iter().cloned().collect(); - let mut wheres: Vec = ctx_fn - .sig + /// The generic params in scope: the host signature's own, plus the outer + /// `impl`/`trait`'s for a method. + fn params(&self) -> impl Iterator { + self.sig .generics - .where_clause - .as_ref() - .map(|wc| wc.predicates.iter().cloned().collect()) - .unwrap_or_default(); - if let Some(outer) = &outer { - params.extend(outer.generics().params.iter().cloned()); - if let Some(wc) = &outer.generics().where_clause { - wheres.extend(wc.predicates.iter().cloned()); - } - } + .params + .iter() + .chain(self.outer.iter().flat_map(|o| o.generics().params.iter())) + } - Ok(( - closure, - Self { - params, - wheres, - is_method: outer.is_some(), - }, - )) + /// The where-predicates in scope, from the host signature and (for a method) + /// the outer `impl`/`trait`. + fn where_predicates(&self) -> impl Iterator { + fn preds(g: &syn::Generics) -> impl Iterator { + g.where_clause.iter().flat_map(|wc| wc.predicates.iter()) + } + preds(&self.sig.generics).chain(self.outer.iter().flat_map(|o| preds(o.generics()))) } } /// Expands a predicate closure into a `#[thrust::formula_fn]` plus a marker -/// call. With `context`, the threaded generics (and, in methods, `Self`) are +/// call. With `context`, the in-scope generics (and, in methods, `Self`) are /// re-declared on the formula function and instantiated via turbofish. fn expand_invariant(closure: &syn::ExprClosure, context: Option<&Context>) -> syn::Expr { let mut fn_params: Vec = Vec::new(); @@ -129,17 +131,18 @@ fn expand_invariant(closure: &syn::ExprClosure, context: Option<&Context>) -> sy } let mut def_params: Vec = Vec::new(); + let mut turbofish_args: Vec = Vec::new(); + // The where-predicates the host already declares, carried verbatim. let mut def_wheres: Vec = context - .iter() - .flat_map(|ctx| ctx.wheres.iter()) + .into_iter() + .flat_map(Context::where_predicates) .cloned() .collect(); - let mut turbofish_args: Vec = Vec::new(); // `Self` in a method context: rewrite it to a synthetic generic, then pass // the real `Self` via turbofish (legal in expression position). let uses_self = - context.is_some_and(|ctx| ctx.is_method) && tokens_contain_self(&closure.to_token_stream()); + context.is_some_and(Context::is_method) && tokens_contain_self(&closure.to_token_stream()); if uses_self { let synth: syn::Ident = format_ident!("__ThrustSelf"); for param in &mut fn_params { @@ -150,22 +153,25 @@ fn expand_invariant(closure: &syn::ExprClosure, context: Option<&Context>) -> sy turbofish_args.push(quote!(Self)); } - for param in context.iter().flat_map(|ctx| ctx.params.iter()) { + // Re-declare the in-scope generics on the formula function and pass them + // through the turbofish. + for param in context.into_iter().flat_map(Context::params) { def_params.push(param.to_token_stream()); match param { - GenericParam::Type(tp) => { - // A closure-typed param has no `Model` instance; re-declare and - // pass it through, but do not bound it. - if !crate::has_fn_bound(&tp.bounds) { - def_wheres.extend(crate::model_predicates(&tp.ident)); - } - turbofish_args.push(tp.ident.to_token_stream()); - } + GenericParam::Type(tp) => turbofish_args.push(tp.ident.to_token_stream()), GenericParam::Const(cp) => turbofish_args.push(cp.ident.to_token_stream()), GenericParam::Lifetime(_) => {} } } + // `Model` bounds for those generics (closure-typed params are skipped). + if let Some(context) = context { + def_wheres.extend(crate::model_where_predicates( + &context.sig, + context.outer.as_ref(), + )); + } + let model_ty_params = crate::fn_params_with_model_ty(&fn_params); let body = &closure.body; diff --git a/thrust-macros/src/invariant_context.rs b/thrust-macros/src/invariant_context.rs index fe9d8c4..9efd03e 100644 --- a/thrust-macros/src/invariant_context.rs +++ b/thrust-macros/src/invariant_context.rs @@ -19,21 +19,21 @@ //! ``` //! //! This macro does **not** expand the invariants itself: it only injects the -//! context, by rewriting each `invariant!(..)` call into a -//! `thrust_macros::_invariant_with_context!(..)` call carrying that context (a -//! `fn` header carrying the in-scope generics/where clause, plus a +//! context, by rewriting each `invariant!(CLOSURE)` call into a +//! `thrust_macros::_invariant_with_context!(#outer_attr #signature; CLOSURE)` +//! call, pasting the host function's signature (and, for methods, a //! `#[thrust::_outer_context(..)]` attribute carrying the enclosing -//! `impl`/`trait` header for methods). +//! `impl`/`trait` header) so the expansion can recover the in-scope context. //! -//! It also extends the host function's where clause with `T: Model` and -//! `::Ty: PartialEq` predicates for every in-scope type parameter +//! It also extends the host function's where clause with the `Model` predicates +//! (see [`crate::model_where_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 proc_macro2::{TokenStream as TokenStream2, TokenTree}; use quote::{quote, ToTokens}; -use syn::{visit_mut::VisitMut, GenericParam, Generics}; +use syn::{visit_mut::VisitMut, Generics, Signature}; use crate::FnOuterItem; @@ -45,73 +45,65 @@ pub(super) fn expand(item: TokenStream) -> TokenStream { Err(e) => return e.to_compile_error().into(), }; - let generics = item_fn.sig.generics.clone(); - let mut injector = InvariantInjector::new(&generics, outer.as_ref()); + let sig = item_fn.sig.clone(); + let mut injector = ContextInjector::new(&sig, outer.as_ref()); injector.visit_block_mut(&mut item_fn.block); injector.inject_model_bounds(&mut item_fn.sig.generics); item_fn.into_token_stream().into() } -struct InvariantInjector<'a> { - /// The enclosing function's own generics. - fn_generics: &'a Generics, +struct ContextInjector<'a> { + /// The host function's signature, pasted verbatim into each rewritten call. + sig: &'a Signature, /// The enclosing `impl`/`trait` header, for a method. outer: Option<&'a FnOuterItem>, /// Whether any rewritten invariant references `Self`. self_used: bool, } -impl<'a> InvariantInjector<'a> { - fn new(fn_generics: &'a Generics, outer: Option<&'a FnOuterItem>) -> Self { +impl<'a> ContextInjector<'a> { + fn new(sig: &'a Signature, outer: Option<&'a FnOuterItem>) -> Self { Self { - fn_generics, + sig, outer, self_used: false, } } /// Builds the context-carrying argument for `_invariant_with_context!` from - /// the closure inside an `invariant!(CLOSURE)` call: a `fn` header carrying - /// the in-scope generics/where clause, tagged (for methods) with the - /// enclosing `impl`/`trait` header. + /// the closure inside an `invariant!(CLOSURE)` call: the host signature + /// pasted as-is, tagged (for methods) with the enclosing `impl`/`trait` + /// header, followed by `; CLOSURE`. fn inject_context(&self, closure: &TokenStream2) -> TokenStream2 { - let generics = self.fn_generics; - let where_clause = &self.fn_generics.where_clause; + let sig = self.sig; let outer_attr = self .outer .map(|outer| quote!(#[thrust::_outer_context(#outer)])); quote! { #outer_attr - fn _thrust_invariant_context #generics () #where_clause { - #closure - } + #sig; + #closure } } - /// Adds `T: Model` and `::Ty: PartialEq` bounds for every type - /// parameter in scope (the host function's own, the outer `impl`/`trait`'s, - /// and `Self` when an invariant names it) to the host's where clause. + /// Extends the host where clause with the `Model` predicates for every + /// in-scope type parameter (and for `Self` when an invariant names it), so + /// the host can instantiate the `Model`-bounded formula functions. fn inject_model_bounds(&self, generics: &mut Generics) { - let mut tys: Vec = type_param_idents(self.fn_generics); - if let Some(outer) = self.outer { - tys.extend(type_param_idents(outer.generics())); - } + let mut predicates = crate::model_where_predicates(self.sig, self.outer); if self.self_used { - tys.push(quote!(Self)); + predicates.extend(crate::model_predicates("e!(Self))); } - if tys.is_empty() { + if predicates.is_empty() { return; } - let where_clause = generics.make_where_clause(); - for ty in tys { - where_clause.predicates.extend(crate::model_predicates(&ty)); - } + generics.make_where_clause().predicates.extend(predicates); } } -impl VisitMut for InvariantInjector<'_> { +impl VisitMut for ContextInjector<'_> { fn visit_macro_mut(&mut self, mac: &mut syn::Macro) { if !is_invariant_macro(&mac.path) { return; @@ -124,21 +116,6 @@ impl VisitMut for InvariantInjector<'_> { } } -/// Type-parameter idents that need a `Model` bound: every type param except -/// closure-typed ones, which have no `Model` instance. -fn type_param_idents(generics: &Generics) -> Vec { - generics - .params - .iter() - .filter_map(|p| match p { - GenericParam::Type(tp) if !crate::has_fn_bound(&tp.bounds) => { - Some(tp.ident.to_token_stream()) - } - _ => None, - }) - .collect() -} - fn is_invariant_macro(path: &syn::Path) -> bool { path.segments.last().is_some_and(|s| s.ident == "invariant") } diff --git a/thrust-macros/src/spec.rs b/thrust-macros/src/spec.rs index 9bdd787..1faf6a3 100644 --- a/thrust-macros/src/spec.rs +++ b/thrust-macros/src/spec.rs @@ -10,8 +10,7 @@ use proc_macro::TokenStream; use proc_macro2::TokenStream as TokenStream2; use quote::{format_ident, quote, ToTokens}; use syn::{ - parse_macro_input, punctuated::Punctuated, FnArg, GenericParam, Generics, TypeParamBound, - WherePredicate, + parse_macro_input, punctuated::Punctuated, FnArg, GenericParam, Generics, WherePredicate, }; use crate::{fn_outer_item::FnOuterItem, fn_params_with_model_ty}; @@ -31,7 +30,7 @@ pub fn expand_predicate(item: TokenStream) -> TokenStream { let model_ty_params = fn_params_with_model_ty(&func.sig().inputs); let model_ret = fn_return_ty_with_model_ty(&func.sig().output); - let model_preds = model_where_predicates(&func, outer_context.as_ref()); + let model_preds = crate::model_where_predicates(func.sig(), outer_context.as_ref()); let extended_where = extended_where_clause(&func, &model_preds); let sig = quote! { @@ -310,7 +309,8 @@ impl ExpandedTokens { } fn extended_where_clause(&self) -> TokenStream2 { - let model_preds = model_where_predicates(&self.func, self.outer_context.as_ref()); + let model_preds = + crate::model_where_predicates(self.func.sig(), self.outer_context.as_ref()); extended_where_clause(&self.func, &model_preds) } @@ -542,96 +542,6 @@ fn rewrite_inputs_for_call( (quote!(#(#rewritten),*), quote!(#(#call_args),*)) } -/// Returns `T: thrust_models::Model` predicates for every type param that does not -/// already carry an `Fn`, `FnOnce`, or `FnMut` bound. -fn model_where_predicates( - func: &FnItemWithSignature, - outer_context: Option<&FnOuterItem>, -) -> Vec { - struct GenericTypeParam { - ident: syn::Ident, - bounds: Vec, - } - - impl From for GenericTypeParam { - fn from(tp: syn::TypeParam) -> Self { - Self { - ident: tp.ident, - bounds: tp.bounds.into_iter().collect(), - } - } - } - - let mut generic_type_params: Vec = Vec::new(); - for param in &func.sig().generics.params { - let GenericParam::Type(tp) = param else { - continue; - }; - generic_type_params.push(tp.clone().into()); - } - if let Some(outer_item) = outer_context { - for param in &outer_item.generics().params { - let GenericParam::Type(tp) = param else { - continue; - }; - generic_type_params.push(tp.clone().into()); - } - if let FnOuterItem::ItemTrait(outer_item) = &outer_item { - generic_type_params.push(GenericTypeParam { - ident: format_ident!("Self"), - bounds: outer_item.supertraits.iter().cloned().collect(), - }); - } - } - generic_type_params.retain(|p| !crate::has_fn_bound(&p.bounds)); - - let mut predicates: Vec = Vec::new(); - for param in &generic_type_params { - predicates.extend(crate::model_predicates(¶m.ident)); - } - - struct Visitor { - generic_type_params: Vec, - generic_paths: Vec, - } - - impl syn::visit::Visit<'_> for Visitor { - fn visit_type_path(&mut self, tp: &syn::TypePath) { - for param in &self.generic_type_params { - if let Some(qself) = &tp.qself { - let param = ¶m.ident; - let param_ty: syn::Type = syn::parse_quote!(#param); - if *qself.ty == param_ty { - self.generic_paths.push(tp.clone()); - } - } - if tp.path.segments.len() > 1 - && tp.path.segments.first().unwrap().ident == param.ident - && tp.qself.is_none() - { - self.generic_paths.push(tp.clone()); - } - } - syn::visit::visit_type_path(self, tp); - } - } - - let mut visitor = Visitor { - generic_type_params, - generic_paths: Vec::new(), - }; - use syn::visit::Visit as _; - for arg in &func.sig().inputs { - visitor.visit_fn_arg(arg); - } - visitor.visit_return_type(&func.sig().output); - for tp in visitor.generic_paths { - predicates.extend(crate::model_predicates(&tp)); - } - - predicates -} - /// Builds `where , `. /// Returns an empty token stream when both sets are empty. fn extended_where_clause( From b47972c9938a5723c90607793d677583db521153 Mon Sep 17 00:00:00 2001 From: coord_e Date: Thu, 4 Jun 2026 00:05:47 +0900 Subject: [PATCH 10/13] style fix --- thrust-macros/src/invariant.rs | 97 ++++++++++---------------- thrust-macros/src/invariant_context.rs | 85 ++++++---------------- 2 files changed, 59 insertions(+), 123 deletions(-) diff --git a/thrust-macros/src/invariant.rs b/thrust-macros/src/invariant.rs index 0c30505..baba89d 100644 --- a/thrust-macros/src/invariant.rs +++ b/thrust-macros/src/invariant.rs @@ -30,7 +30,7 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use proc_macro::TokenStream; -use proc_macro2::{TokenStream as TokenStream2, TokenTree}; +use proc_macro2::TokenStream as TokenStream2; use quote::{format_ident, quote, ToTokens}; use syn::{ parse::{Parse, ParseStream}, @@ -53,32 +53,31 @@ pub fn expand(input: TokenStream) -> TokenStream { /// Expands `_invariant_with_context!(#outer_attr #sig; CLOSURE)`, the form /// `#[thrust_macros::invariant_context]` rewrites each `invariant!` into. pub fn expand_with_context(input: TokenStream) -> TokenStream { + struct WithContext { + context: Context, + closure: syn::ExprClosure, + } + + impl Parse for WithContext { + fn parse(input: ParseStream) -> syn::Result { + let attrs = input.call(syn::Attribute::parse_outer)?; + let outer = crate::extract_outer_context(&attrs)?; + let sig: Signature = input.parse()?; + input.parse::()?; + let closure: syn::ExprClosure = input.parse()?; + Ok(Self { + context: Context { sig, outer }, + closure, + }) + } + } + let WithContext { closure, context } = parse_macro_input!(input as WithContext); expand_invariant(&closure, Some(&context)) .into_token_stream() .into() } -/// The parsed `_invariant_with_context!` input. -struct WithContext { - context: Context, - closure: syn::ExprClosure, -} - -impl Parse for WithContext { - fn parse(input: ParseStream) -> syn::Result { - let attrs = input.call(syn::Attribute::parse_outer)?; - let outer = crate::extract_outer_context(&attrs)?; - let sig: Signature = input.parse()?; - input.parse::()?; - let closure: syn::ExprClosure = input.parse()?; - Ok(Self { - context: Context { sig, outer }, - closure, - }) - } -} - /// The enclosing context threaded into an invariant by /// `#[thrust_macros::invariant_context]`: the host function signature and, for a /// method, its `impl`/`trait` header. A standalone `invariant!` has none. @@ -88,13 +87,9 @@ struct Context { } impl Context { - fn is_method(&self) -> bool { - self.outer.is_some() - } - /// The generic params in scope: the host signature's own, plus the outer /// `impl`/`trait`'s for a method. - fn params(&self) -> impl Iterator { + fn generic_params(&self) -> impl Iterator { self.sig .generics .params @@ -132,18 +127,30 @@ fn expand_invariant(closure: &syn::ExprClosure, context: Option<&Context>) -> sy let mut def_params: Vec = Vec::new(); let mut turbofish_args: Vec = Vec::new(); - // The where-predicates the host already declares, carried verbatim. + for param in context.into_iter().flat_map(Context::generic_params) { + def_params.push(param.to_token_stream()); + match param { + GenericParam::Type(tp) => turbofish_args.push(tp.ident.to_token_stream()), + GenericParam::Const(cp) => turbofish_args.push(cp.ident.to_token_stream()), + GenericParam::Lifetime(_) => {} + } + } + let mut def_wheres: Vec = context .into_iter() .flat_map(Context::where_predicates) .cloned() .collect(); + if let Some(context) = context { + def_wheres.extend(crate::model_where_predicates( + &context.sig, + context.outer.as_ref(), + )); + } // `Self` in a method context: rewrite it to a synthetic generic, then pass // the real `Self` via turbofish (legal in expression position). - let uses_self = - context.is_some_and(Context::is_method) && tokens_contain_self(&closure.to_token_stream()); - if uses_self { + if crate::tokens_contain_ident(&closure.to_token_stream(), "Self") { let synth: syn::Ident = format_ident!("__ThrustSelf"); for param in &mut fn_params { SelfRewriter { synth: &synth }.visit_fn_arg_mut(param); @@ -153,25 +160,6 @@ fn expand_invariant(closure: &syn::ExprClosure, context: Option<&Context>) -> sy turbofish_args.push(quote!(Self)); } - // Re-declare the in-scope generics on the formula function and pass them - // through the turbofish. - for param in context.into_iter().flat_map(Context::params) { - def_params.push(param.to_token_stream()); - match param { - GenericParam::Type(tp) => turbofish_args.push(tp.ident.to_token_stream()), - GenericParam::Const(cp) => turbofish_args.push(cp.ident.to_token_stream()), - GenericParam::Lifetime(_) => {} - } - } - - // `Model` bounds for those generics (closure-typed params are skipped). - if let Some(context) = context { - def_wheres.extend(crate::model_where_predicates( - &context.sig, - context.outer.as_ref(), - )); - } - let model_ty_params = crate::fn_params_with_model_ty(&fn_params); let body = &closure.body; @@ -206,9 +194,6 @@ fn expand_invariant(closure: &syn::ExprClosure, context: Option<&Context>) -> sy }) } -/// Rewrites a bare `Self` type path to a synthetic type parameter, so the type -/// can be named inside a nested formula function. Qualified paths such as -/// `Self::Assoc` are left untouched (and are not supported in invariants). struct SelfRewriter<'a> { synth: &'a syn::Ident, } @@ -224,11 +209,3 @@ impl VisitMut for SelfRewriter<'_> { } } } - -fn tokens_contain_self(tokens: &TokenStream2) -> bool { - tokens.clone().into_iter().any(|tt| match tt { - TokenTree::Ident(ident) => ident == "Self", - TokenTree::Group(group) => tokens_contain_self(&group.stream()), - _ => false, - }) -} diff --git a/thrust-macros/src/invariant_context.rs b/thrust-macros/src/invariant_context.rs index 9efd03e..a4e0dab 100644 --- a/thrust-macros/src/invariant_context.rs +++ b/thrust-macros/src/invariant_context.rs @@ -5,39 +5,19 @@ //! an invariant may refer to generic- and `Self`-typed variables that the //! standalone `invariant!` macro cannot see. //! -//! It is applied to a single function. A method recovers its enclosing -//! `impl`/`trait` header from the `#[thrust::_outer_context(..)]` attribute -//! stamped by `#[thrust_macros::context]`, the same mechanism `requires`/ -//! `ensures` use: -//! -//! ```ignore -//! #[thrust_macros::context] -//! impl Foo { -//! #[thrust_macros::invariant_context] -//! fn f(&self) { .. } -//! } -//! ``` -//! -//! This macro does **not** expand the invariants itself: it only injects the -//! context, by rewriting each `invariant!(CLOSURE)` call into a -//! `thrust_macros::_invariant_with_context!(#outer_attr #signature; CLOSURE)` -//! call, pasting the host function's signature (and, for methods, a -//! `#[thrust::_outer_context(..)]` attribute carrying the enclosing -//! `impl`/`trait` header) so the expansion can recover the in-scope context. -//! //! It also extends the host function's where clause with the `Model` predicates //! (see [`crate::model_where_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 proc_macro2::{TokenStream as TokenStream2, TokenTree}; +use proc_macro2::TokenStream as TokenStream2; use quote::{quote, ToTokens}; -use syn::{visit_mut::VisitMut, Generics, Signature}; +use syn::{visit_mut::VisitMut, Signature}; use crate::FnOuterItem; -pub(super) fn expand(item: TokenStream) -> TokenStream { +pub fn expand(item: TokenStream) -> TokenStream { let mut item_fn = syn::parse_macro_input!(item as syn::ItemFn); let outer = match crate::extract_outer_context(&item_fn.attrs) { @@ -46,35 +26,36 @@ pub(super) fn expand(item: TokenStream) -> TokenStream { }; let sig = item_fn.sig.clone(); - let mut injector = ContextInjector::new(&sig, outer.as_ref()); + let mut injector = ContextInjector { + sig: &sig, + outer: outer.as_ref(), + self_used: false, + }; injector.visit_block_mut(&mut item_fn.block); - injector.inject_model_bounds(&mut item_fn.sig.generics); + + let mut predicates = crate::model_where_predicates(&sig, outer.as_ref()); + if injector.self_used { + predicates.extend(crate::model_predicates("e!(Self))); + } + if !predicates.is_empty() { + item_fn + .sig + .generics + .make_where_clause() + .predicates + .extend(predicates); + } item_fn.into_token_stream().into() } struct ContextInjector<'a> { - /// The host function's signature, pasted verbatim into each rewritten call. sig: &'a Signature, - /// The enclosing `impl`/`trait` header, for a method. outer: Option<&'a FnOuterItem>, - /// Whether any rewritten invariant references `Self`. self_used: bool, } impl<'a> ContextInjector<'a> { - fn new(sig: &'a Signature, outer: Option<&'a FnOuterItem>) -> Self { - Self { - sig, - outer, - self_used: false, - } - } - - /// Builds the context-carrying argument for `_invariant_with_context!` from - /// the closure inside an `invariant!(CLOSURE)` call: the host signature - /// pasted as-is, tagged (for methods) with the enclosing `impl`/`trait` - /// header, followed by `; CLOSURE`. fn inject_context(&self, closure: &TokenStream2) -> TokenStream2 { let sig = self.sig; let outer_attr = self @@ -87,20 +68,6 @@ impl<'a> ContextInjector<'a> { #closure } } - - /// Extends the host where clause with the `Model` predicates for every - /// in-scope type parameter (and for `Self` when an invariant names it), so - /// the host can instantiate the `Model`-bounded formula functions. - fn inject_model_bounds(&self, generics: &mut Generics) { - let mut predicates = crate::model_where_predicates(self.sig, self.outer); - if self.self_used { - predicates.extend(crate::model_predicates("e!(Self))); - } - if predicates.is_empty() { - return; - } - generics.make_where_clause().predicates.extend(predicates); - } } impl VisitMut for ContextInjector<'_> { @@ -108,7 +75,7 @@ impl VisitMut for ContextInjector<'_> { if !is_invariant_macro(&mac.path) { return; } - if self.outer.is_some() && tokens_contain_self(&mac.tokens) { + if crate::tokens_contain_ident(&mac.tokens, "Self") { self.self_used = true; } mac.tokens = self.inject_context(&mac.tokens); @@ -119,11 +86,3 @@ impl VisitMut for ContextInjector<'_> { fn is_invariant_macro(path: &syn::Path) -> bool { path.segments.last().is_some_and(|s| s.ident == "invariant") } - -fn tokens_contain_self(tokens: &TokenStream2) -> bool { - tokens.clone().into_iter().any(|tt| match tt { - TokenTree::Ident(ident) => ident == "Self", - TokenTree::Group(group) => tokens_contain_self(&group.stream()), - _ => false, - }) -} From bd80c7cf559bca615c30233f9af8a8a8f7e85db4 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 4 Jun 2026 02:39:06 +0000 Subject: [PATCH 11/13] thrust-macros: reconcile rebase with main's fn_outer_item module MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Resolution of the rebase onto latest main: main (#101) extracted FnOuterItem into its own fn_outer_item module and dropped it from the crate root. Adapt the loop-invariant work accordingly — reference crate::fn_outer_item::FnOuterItem and keep the shared helpers (extract_outer_context, has_fn_bound, model_predicates, model_where_predicates, tokens_contain_ident) in the crate root. --- thrust-macros/src/invariant.rs | 2 +- thrust-macros/src/invariant_context.rs | 2 +- thrust-macros/src/lib.rs | 164 +++++++++++++++++++++++-- thrust-macros/src/spec.rs | 8 +- 4 files changed, 160 insertions(+), 16 deletions(-) diff --git a/thrust-macros/src/invariant.rs b/thrust-macros/src/invariant.rs index baba89d..3dace70 100644 --- a/thrust-macros/src/invariant.rs +++ b/thrust-macros/src/invariant.rs @@ -39,7 +39,7 @@ use syn::{ FnArg, GenericParam, Signature, WherePredicate, }; -use crate::FnOuterItem; +use crate::fn_outer_item::FnOuterItem; static COUNTER: AtomicUsize = AtomicUsize::new(0); diff --git a/thrust-macros/src/invariant_context.rs b/thrust-macros/src/invariant_context.rs index a4e0dab..f0dc102 100644 --- a/thrust-macros/src/invariant_context.rs +++ b/thrust-macros/src/invariant_context.rs @@ -15,7 +15,7 @@ use proc_macro2::TokenStream as TokenStream2; use quote::{quote, ToTokens}; use syn::{visit_mut::VisitMut, Signature}; -use crate::FnOuterItem; +use crate::fn_outer_item::FnOuterItem; pub fn expand(item: TokenStream) -> TokenStream { let mut item_fn = syn::parse_macro_input!(item as syn::ItemFn); diff --git a/thrust-macros/src/lib.rs b/thrust-macros/src/lib.rs index c23b043..f1f51c6 100644 --- a/thrust-macros/src/lib.rs +++ b/thrust-macros/src/lib.rs @@ -1,5 +1,7 @@ use proc_macro::TokenStream; -use proc_macro2::TokenStream as TokenStream2; +use proc_macro2::{TokenStream as TokenStream2, TokenTree as TokenTree2}; +use quote::quote; +use syn::{FnArg, GenericParam, TypeParamBound}; mod context; mod fn_outer_item; @@ -7,6 +9,8 @@ mod invariant; mod invariant_context; mod spec; +use fn_outer_item::FnOuterItem; + #[proc_macro_attribute] pub fn context(_attr: TokenStream, item: TokenStream) -> TokenStream { context::expand(item) @@ -45,10 +49,6 @@ pub fn _invariant_with_context(input: TokenStream) -> TokenStream { /// invariant may refer to generic- and `Self`-typed variables that the /// standalone `invariant!` macro cannot see. Each such call is rewritten into /// `thrust_macros::_invariant_with_context!`. -/// -/// Applied to a single function; a method recovers its enclosing -/// `impl`/`trait` context from `#[thrust_macros::context]` on the surrounding -/// block, as `requires`/`ensures` do. #[proc_macro_attribute] pub fn invariant_context(_attr: TokenStream, item: TokenStream) -> TokenStream { invariant_context::expand(item) @@ -74,24 +74,168 @@ pub fn _requires_ensures(attr: TokenStream, item: TokenStream) -> TokenStream { spec::expand_requires_ensures(attr, item) } +/// Reads the `#[thrust::_outer_context(..)]` attribute stamped onto methods by +/// `#[thrust_macros::context]` (and threaded by `invariant_context`), returning +/// the enclosing `impl`/`trait` header it carries, or `None` if absent. +fn extract_outer_context(attrs: &[syn::Attribute]) -> syn::Result> { + let outer_context_path: syn::Path = syn::parse_quote!(thrust::_outer_context); + let mut outer_context = None; + for attr in attrs { + if attr.path() != &outer_context_path { + continue; + } + if outer_context.is_some() { + return Err(syn::Error::new_spanned( + attr, + "multiple _outer_context attributes found; expected at most one", + )); + } + outer_context = Some(attr.parse_args()?); + } + Ok(outer_context) +} + +fn has_fn_bound<'a>(bounds: impl IntoIterator) -> bool { + bounds.into_iter().any(|b| { + let TypeParamBound::Trait(tb) = b else { + return false; + }; + tb.path + .segments + .last() + .is_some_and(|s| matches!(s.ident.to_string().as_str(), "Fn" | "FnOnce" | "FnMut")) + }) +} + +fn model_predicates(ty: &impl quote::ToTokens) -> [syn::WherePredicate; 2] { + [ + syn::parse_quote!(#ty: thrust_models::Model), + syn::parse_quote!(<#ty as thrust_models::Model>::Ty: PartialEq), + ] +} + +/// `T: Model` / `::Ty: PartialEq` predicates for every type param +/// in scope for `sig` (its own, plus the outer `impl`/`trait`'s and, for a +/// trait, `Self`) that does not carry an `Fn`/`FnOnce`/`FnMut` bound, plus the +/// same for any generic associated-type projection appearing in `sig`. +fn model_where_predicates( + sig: &syn::Signature, + outer_context: Option<&FnOuterItem>, +) -> Vec { + struct GenericTypeParam { + ident: syn::Ident, + bounds: Vec, + } + + impl From for GenericTypeParam { + fn from(tp: syn::TypeParam) -> Self { + Self { + ident: tp.ident, + bounds: tp.bounds.into_iter().collect(), + } + } + } + + let mut generic_type_params: Vec = Vec::new(); + for param in &sig.generics.params { + let GenericParam::Type(tp) = param else { + continue; + }; + generic_type_params.push(tp.clone().into()); + } + if let Some(outer_item) = outer_context { + for param in &outer_item.generics().params { + let GenericParam::Type(tp) = param else { + continue; + }; + generic_type_params.push(tp.clone().into()); + } + if let FnOuterItem::ItemTrait(outer_item) = &outer_item { + generic_type_params.push(GenericTypeParam { + ident: quote::format_ident!("Self"), + bounds: outer_item.supertraits.iter().cloned().collect(), + }); + } + } + generic_type_params.retain(|p| !has_fn_bound(&p.bounds)); + + let mut predicates: Vec = Vec::new(); + for param in &generic_type_params { + predicates.extend(model_predicates(¶m.ident)); + } + + struct Visitor { + generic_type_params: Vec, + generic_paths: Vec, + } + + impl syn::visit::Visit<'_> for Visitor { + fn visit_type_path(&mut self, tp: &syn::TypePath) { + for param in &self.generic_type_params { + if let Some(qself) = &tp.qself { + let param = ¶m.ident; + let param_ty: syn::Type = syn::parse_quote!(#param); + if *qself.ty == param_ty { + self.generic_paths.push(tp.clone()); + } + } + if tp.path.segments.len() > 1 + && tp.path.segments.first().unwrap().ident == param.ident + && tp.qself.is_none() + { + self.generic_paths.push(tp.clone()); + } + } + syn::visit::visit_type_path(self, tp); + } + } + + let mut visitor = Visitor { + generic_type_params, + generic_paths: Vec::new(), + }; + use syn::visit::Visit as _; + for arg in &sig.inputs { + visitor.visit_fn_arg(arg); + } + visitor.visit_return_type(&sig.output); + for tp in visitor.generic_paths { + predicates.extend(model_predicates(&tp)); + } + + predicates +} + /// Maps each function parameter `x: T` to `x: ::Ty`. fn fn_params_with_model_ty<'ast, I>(args: I) -> TokenStream2 where - I: IntoIterator, + I: IntoIterator, { - let mut model_inputs: Vec = Vec::new(); + let mut model_inputs: Vec = Vec::new(); for arg in args { match arg { - syn::FnArg::Receiver(receiver) => { + FnArg::Receiver(receiver) => { let ty = &receiver.ty; model_inputs.push(syn::parse_quote!(self_: <#ty as thrust_models::Model>::Ty)); } - syn::FnArg::Typed(pt) => { + FnArg::Typed(pt) => { let pat = &pt.pat; let ty = &pt.ty; model_inputs.push(syn::parse_quote!(#pat: <#ty as thrust_models::Model>::Ty)); } } } - quote::quote!(#(#model_inputs),*) + quote!(#(#model_inputs),*) +} + +fn tokens_contain_ident(tokens: &TokenStream2, target: T) -> bool +where + T: AsRef, +{ + let target = target.as_ref(); + tokens.clone().into_iter().any(|tt| match tt { + TokenTree2::Ident(ident) => ident == target, + TokenTree2::Group(group) => tokens_contain_ident(&group.stream(), target), + _ => false, + }) } diff --git a/thrust-macros/src/spec.rs b/thrust-macros/src/spec.rs index 1faf6a3..3d86505 100644 --- a/thrust-macros/src/spec.rs +++ b/thrust-macros/src/spec.rs @@ -15,7 +15,7 @@ use syn::{ use crate::{fn_outer_item::FnOuterItem, fn_params_with_model_ty}; -pub fn expand_predicate(item: TokenStream) -> TokenStream { +pub(super) fn expand_predicate(item: TokenStream) -> TokenStream { let func = parse_macro_input!(item as FnItemWithSignature); let outer_context = match extract_outer_context(&func) { Ok(ctx) => ctx, @@ -45,7 +45,7 @@ pub fn expand_predicate(item: TokenStream) -> TokenStream { } } -pub fn expand_requires(attr: TokenStream, item: TokenStream) -> TokenStream { +pub(super) fn expand_requires(attr: TokenStream, item: TokenStream) -> TokenStream { let expr = TokenStream2::from(attr); let mut func = parse_macro_input!(item as FnItemWithSignature); @@ -60,7 +60,7 @@ pub fn expand_requires(attr: TokenStream, item: TokenStream) -> TokenStream { func.into_token_stream().into() } -pub fn expand_ensures(attr: TokenStream, item: TokenStream) -> TokenStream { +pub(super) fn expand_ensures(attr: TokenStream, item: TokenStream) -> TokenStream { let expr = TokenStream2::from(attr); let mut func = parse_macro_input!(item as FnItemWithSignature); @@ -75,7 +75,7 @@ pub fn expand_ensures(attr: TokenStream, item: TokenStream) -> TokenStream { func.into_token_stream().into() } -pub fn expand_requires_ensures(attr: TokenStream, item: TokenStream) -> TokenStream { +pub(super) fn expand_requires_ensures(attr: TokenStream, item: TokenStream) -> TokenStream { use syn::parse::Parser as _; let parser = Punctuated::::parse_separated_nonempty; let mut exprs = match parser.parse(attr.clone()) { From df7c40c3075a8d9979ad3d2a6248b5dcbfb97f51 Mon Sep 17 00:00:00 2001 From: coord_e Date: Thu, 4 Jun 2026 13:52:05 +0900 Subject: [PATCH 12/13] style fix --- thrust-macros/src/invariant_context.rs | 1 + thrust-macros/src/lib.rs | 22 ++++++++++------------ thrust-macros/src/spec.rs | 8 ++++---- 3 files changed, 15 insertions(+), 16 deletions(-) diff --git a/thrust-macros/src/invariant_context.rs b/thrust-macros/src/invariant_context.rs index f0dc102..af3a30a 100644 --- a/thrust-macros/src/invariant_context.rs +++ b/thrust-macros/src/invariant_context.rs @@ -84,5 +84,6 @@ impl VisitMut for ContextInjector<'_> { } fn is_invariant_macro(path: &syn::Path) -> bool { + // TODO: identify the macro precisely path.segments.last().is_some_and(|s| s.ident == "invariant") } diff --git a/thrust-macros/src/lib.rs b/thrust-macros/src/lib.rs index f1f51c6..a0eb9bd 100644 --- a/thrust-macros/src/lib.rs +++ b/thrust-macros/src/lib.rs @@ -1,7 +1,5 @@ use proc_macro::TokenStream; use proc_macro2::{TokenStream as TokenStream2, TokenTree as TokenTree2}; -use quote::quote; -use syn::{FnArg, GenericParam, TypeParamBound}; mod context; mod fn_outer_item; @@ -95,9 +93,9 @@ fn extract_outer_context(attrs: &[syn::Attribute]) -> syn::Result(bounds: impl IntoIterator) -> bool { +fn has_fn_bound<'a>(bounds: impl IntoIterator) -> bool { bounds.into_iter().any(|b| { - let TypeParamBound::Trait(tb) = b else { + let syn::TypeParamBound::Trait(tb) = b else { return false; }; tb.path @@ -124,7 +122,7 @@ fn model_where_predicates( ) -> Vec { struct GenericTypeParam { ident: syn::Ident, - bounds: Vec, + bounds: Vec, } impl From for GenericTypeParam { @@ -138,14 +136,14 @@ fn model_where_predicates( let mut generic_type_params: Vec = Vec::new(); for param in &sig.generics.params { - let GenericParam::Type(tp) = param else { + let syn::GenericParam::Type(tp) = param else { continue; }; generic_type_params.push(tp.clone().into()); } if let Some(outer_item) = outer_context { for param in &outer_item.generics().params { - let GenericParam::Type(tp) = param else { + let syn::GenericParam::Type(tp) = param else { continue; }; generic_type_params.push(tp.clone().into()); @@ -209,23 +207,23 @@ fn model_where_predicates( /// Maps each function parameter `x: T` to `x: ::Ty`. fn fn_params_with_model_ty<'ast, I>(args: I) -> TokenStream2 where - I: IntoIterator, + I: IntoIterator, { - let mut model_inputs: Vec = Vec::new(); + let mut model_inputs: Vec = Vec::new(); for arg in args { match arg { - FnArg::Receiver(receiver) => { + syn::FnArg::Receiver(receiver) => { let ty = &receiver.ty; model_inputs.push(syn::parse_quote!(self_: <#ty as thrust_models::Model>::Ty)); } - FnArg::Typed(pt) => { + syn::FnArg::Typed(pt) => { let pat = &pt.pat; let ty = &pt.ty; model_inputs.push(syn::parse_quote!(#pat: <#ty as thrust_models::Model>::Ty)); } } } - quote!(#(#model_inputs),*) + quote::quote!(#(#model_inputs),*) } fn tokens_contain_ident(tokens: &TokenStream2, target: T) -> bool diff --git a/thrust-macros/src/spec.rs b/thrust-macros/src/spec.rs index 3d86505..1faf6a3 100644 --- a/thrust-macros/src/spec.rs +++ b/thrust-macros/src/spec.rs @@ -15,7 +15,7 @@ use syn::{ use crate::{fn_outer_item::FnOuterItem, fn_params_with_model_ty}; -pub(super) fn expand_predicate(item: TokenStream) -> TokenStream { +pub fn expand_predicate(item: TokenStream) -> TokenStream { let func = parse_macro_input!(item as FnItemWithSignature); let outer_context = match extract_outer_context(&func) { Ok(ctx) => ctx, @@ -45,7 +45,7 @@ pub(super) fn expand_predicate(item: TokenStream) -> TokenStream { } } -pub(super) fn expand_requires(attr: TokenStream, item: TokenStream) -> TokenStream { +pub fn expand_requires(attr: TokenStream, item: TokenStream) -> TokenStream { let expr = TokenStream2::from(attr); let mut func = parse_macro_input!(item as FnItemWithSignature); @@ -60,7 +60,7 @@ pub(super) fn expand_requires(attr: TokenStream, item: TokenStream) -> TokenStre func.into_token_stream().into() } -pub(super) fn expand_ensures(attr: TokenStream, item: TokenStream) -> TokenStream { +pub fn expand_ensures(attr: TokenStream, item: TokenStream) -> TokenStream { let expr = TokenStream2::from(attr); let mut func = parse_macro_input!(item as FnItemWithSignature); @@ -75,7 +75,7 @@ pub(super) fn expand_ensures(attr: TokenStream, item: TokenStream) -> TokenStrea func.into_token_stream().into() } -pub(super) fn expand_requires_ensures(attr: TokenStream, item: TokenStream) -> TokenStream { +pub fn expand_requires_ensures(attr: TokenStream, item: TokenStream) -> TokenStream { use syn::parse::Parser as _; let parser = Punctuated::::parse_separated_nonempty; let mut exprs = match parser.parse(attr.clone()) { From 6ac606420279eba413bc59cd5a85eeb5ef5140e2 Mon Sep 17 00:00:00 2001 From: coord_e Date: Thu, 4 Jun 2026 20:19:52 +0900 Subject: [PATCH 13/13] style fix --- thrust-macros/src/invariant.rs | 26 ++++++++++++++++---------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/thrust-macros/src/invariant.rs b/thrust-macros/src/invariant.rs index 3dace70..afdd447 100644 --- a/thrust-macros/src/invariant.rs +++ b/thrust-macros/src/invariant.rs @@ -47,7 +47,10 @@ static COUNTER: AtomicUsize = AtomicUsize::new(0); /// context. pub fn expand(input: TokenStream) -> TokenStream { let closure = parse_macro_input!(input as syn::ExprClosure); - expand_invariant(&closure, None).into_token_stream().into() + match expand_invariant(&closure, None) { + Ok(expr) => expr.into_token_stream().into(), + Err(e) => e.to_compile_error().into(), + } } /// Expands `_invariant_with_context!(#outer_attr #sig; CLOSURE)`, the form @@ -73,9 +76,10 @@ pub fn expand_with_context(input: TokenStream) -> TokenStream { } let WithContext { closure, context } = parse_macro_input!(input as WithContext); - expand_invariant(&closure, Some(&context)) - .into_token_stream() - .into() + match expand_invariant(&closure, Some(&context)) { + Ok(expr) => expr.into_token_stream().into(), + Err(e) => e.to_compile_error().into(), + } } /// The enclosing context threaded into an invariant by @@ -110,15 +114,17 @@ impl Context { /// Expands a predicate closure into a `#[thrust::formula_fn]` plus a marker /// call. With `context`, the in-scope generics (and, in methods, `Self`) are /// re-declared on the formula function and instantiated via turbofish. -fn expand_invariant(closure: &syn::ExprClosure, context: Option<&Context>) -> syn::Expr { +fn expand_invariant( + closure: &syn::ExprClosure, + context: Option<&Context>, +) -> syn::Result { let mut fn_params: Vec = Vec::new(); for param in &closure.inputs { let syn::Pat::Type(pt) = param else { - let err = syn::Error::new_spanned( + return Err(syn::Error::new_spanned( param, "invariant closure parameters must have explicit types, e.g. `|x: i64| ...`", - ); - return syn::Expr::Verbatim(err.to_compile_error()); + )); }; let pat = &pt.pat; let ty = &pt.ty; @@ -182,7 +188,7 @@ fn expand_invariant(closure: &syn::ExprClosure, context: Option<&Context>) -> sy quote!(::<#(#turbofish_args),*>) }; - syn::parse_quote!({ + Ok(syn::parse_quote!({ #[allow(unused_variables)] #[allow(non_snake_case)] #[thrust::formula_fn] @@ -191,7 +197,7 @@ fn expand_invariant(closure: &syn::ExprClosure, context: Option<&Context>) -> sy } thrust_models::__invariant_marker(#name #turbofish) - }) + })) } struct SelfRewriter<'a> {