diff --git a/src/analyze/annot_fn.rs b/src/analyze/annot_fn.rs index 7d0200a8..92a8b2c7 100644 --- a/src/analyze/annot_fn.rs +++ b/src/analyze/annot_fn.rs @@ -19,6 +19,20 @@ pub struct FormulaFn<'tcx> { formula: chc::Formula, } +/// The source name a parameter of a formula function lifted out of a function body +/// (`invariant!`, `ghost!`) refers to. +/// +/// The lifted function is free, where `self` is not a legal parameter name, so a formula +/// naming the receiver gets a synthetic parameter instead. It stands for the value that +/// debug info records as `self`. +pub fn lifted_param_source_name(ident: rustc_span::symbol::Ident) -> rustc_span::Symbol { + if ident.name.as_str() == "__thrust_self" { + rustc_span::Symbol::intern("self") + } else { + ident.name + } +} + impl<'a, D> Pretty<'a, D, termcolor::ColorSpec> for &FormulaFn<'_> where D: pretty::DocAllocator<'a, termcolor::ColorSpec>, diff --git a/src/analyze/basic_block.rs b/src/analyze/basic_block.rs index 97ae3ec0..3e885f7e 100644 --- a/src/analyze/basic_block.rs +++ b/src/analyze/basic_block.rs @@ -1054,9 +1054,10 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> { .skip(1) .map(|ident| { let ident = ident.expect("ghost term parameters must be named"); - let operand = self.operand_of_name(ident.name).unwrap_or_else(|| { + let name = analyze::annot_fn::lifted_param_source_name(ident); + let operand = self.operand_of_name(name).unwrap_or_else(|| { self.tcx.dcx().fatal(format!( - "ghost term refers to `{ident}`, which is not a live variable here" + "ghost term refers to `{name}`, which is not a live variable here" )) }); self.operand_refined_type(operand) diff --git a/src/analyze/local_def.rs b/src/analyze/local_def.rs index a98eb1ce..e4954ee8 100644 --- a/src/analyze/local_def.rs +++ b/src/analyze/local_def.rs @@ -853,13 +853,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> { .unwrap_or(*input_ty) }; - // The synthetic `__thrust_self` parameter (emitted when an invariant refers to the receiver - // `self`) maps to the loop-carried receiver, which appears as `self` in debug info. - let name = if ident.name.as_str() == "__thrust_self" { - rustc_span::Symbol::intern("self") - } else { - ident.name - }; + let name = analyze::annot_fn::lifted_param_source_name(ident); if input_ty .ty_adt_def() diff --git a/tests/ui/fail/ghost_generic.rs b/tests/ui/fail/ghost_generic.rs new file mode 100644 index 00000000..5dc8c4d9 --- /dev/null +++ b/tests/ui/fail/ghost_generic.rs @@ -0,0 +1,19 @@ +//@error-in-other-file: Unsat +//@compile-flags: -C debug-assertions=off -A unused-variables + +use thrust_models::Ghost; + +#[thrust_macros::requires(g == v)] +fn expect_same(g: Ghost, v: T) { + let _ = g; +} + +#[thrust_macros::context] +fn record(a: T, b: T) { + let g = thrust_macros::ghost!(|b: T| -> T { b }); + expect_same(g, a); +} + +fn main() { + record(3_i64, 5_i64); +} diff --git a/tests/ui/fail/ghost_self.rs b/tests/ui/fail/ghost_self.rs new file mode 100644 index 00000000..d0ab9794 --- /dev/null +++ b/tests/ui/fail/ghost_self.rs @@ -0,0 +1,26 @@ +//@error-in-other-file: Unsat +//@compile-flags: -C debug-assertions=off -A unused-variables + +use thrust_models::model::{Int, Seq}; +use thrust_models::Ghost; + +struct Counter { + count: i64, + seen: Ghost>, +} + +impl thrust_models::Model for Counter { + type Ty = (Int, Seq); +} + +#[thrust_macros::context] +impl Counter { + #[thrust_macros::requires((*self).1.len() == (*self).0)] + #[thrust_macros::ensures((!self).1.len() == (!self).0)] + fn record(&mut self, x: i64) { + self.count += 1; + self.seen = thrust_macros::ghost!(|self: &mut Self, x: i64| -> Seq { (*self).1 }); + } +} + +fn main() {} diff --git a/tests/ui/pass/ghost_generic.rs b/tests/ui/pass/ghost_generic.rs new file mode 100644 index 00000000..9962c703 --- /dev/null +++ b/tests/ui/pass/ghost_generic.rs @@ -0,0 +1,19 @@ +//@check-pass +//@compile-flags: -C debug-assertions=off -A unused-variables + +use thrust_models::Ghost; + +#[thrust_macros::requires(g == v)] +fn expect_same(g: Ghost, v: T) { + let _ = g; +} + +#[thrust_macros::context] +fn record(a: T, b: T) { + let g = thrust_macros::ghost!(|a: T| -> T { a }); + expect_same(g, a); +} + +fn main() { + record(3_i64, 5_i64); +} diff --git a/tests/ui/pass/ghost_self.rs b/tests/ui/pass/ghost_self.rs new file mode 100644 index 00000000..8e7aec8c --- /dev/null +++ b/tests/ui/pass/ghost_self.rs @@ -0,0 +1,27 @@ +//@check-pass +//@compile-flags: -C debug-assertions=off -A unused-variables + +use thrust_models::model::{Int, Seq}; +use thrust_models::Ghost; + +struct Counter { + count: i64, + seen: Ghost>, +} + +impl thrust_models::Model for Counter { + type Ty = (Int, Seq); +} + +#[thrust_macros::context] +impl Counter { + #[thrust_macros::requires((*self).1.len() == (*self).0)] + #[thrust_macros::ensures((!self).1.len() == (!self).0)] + fn record(&mut self, x: i64) { + self.count += 1; + self.seen = + thrust_macros::ghost!(|self: &mut Self, x: i64| -> Seq { (*self).1.push(x) }); + } +} + +fn main() {} diff --git a/thrust-macros/src/context.rs b/thrust-macros/src/context.rs index 0ba272a5..a358c798 100644 --- a/thrust-macros/src/context.rs +++ b/thrust-macros/src/context.rs @@ -2,10 +2,11 @@ //! //! Makes the enclosing context available to the specifications written inside an item: //! -//! - On a function, every `thrust_macros::invariant!(...)` in the body is rewritten into -//! its context-carrying counterpart, carrying the host signature (and, in a method, the -//! enclosing `impl`/`trait` header), so an invariant may refer to generic- and -//! `Self`-typed variables that the standalone macro cannot see. +//! - On a function, every `thrust_macros::invariant!(...)` and +//! `thrust_macros::ghost!(...)` in the body is rewritten into its context-carrying +//! counterpart, carrying the host signature (and, in a method, the enclosing +//! `impl`/`trait` header), so a formula may refer to generic- and `Self`-typed +//! variables that the standalone macros cannot see. //! - On an `impl`/`trait`, each method is stamped with the enclosing header so //! method-level `requires`/`ensures` can recover the outer generics, and its body is //! threaded as above. @@ -164,6 +165,7 @@ fn context_carrying_form(path: &syn::Path) -> Option { // TODO: identify the macro precisely match path.segments.last()?.ident.to_string().as_str() { "invariant" => Some(syn::parse_quote!(::thrust_macros::_invariant_with_context)), + "ghost" => Some(syn::parse_quote!(::thrust_macros::_ghost_with_context)), _ => None, } } diff --git a/thrust-macros/src/formula_fn_lifting.rs b/thrust-macros/src/formula_fn_lifting.rs new file mode 100644 index 00000000..12b16d43 --- /dev/null +++ b/thrust-macros/src/formula_fn_lifting.rs @@ -0,0 +1,375 @@ +//! Lifting a formula written inside a function body into a standalone +//! `#[thrust::formula_fn]` item that a marker call refers to. +//! +//! The item is a free function, so it inherits neither the enclosing function's generics +//! nor `Self`: a formula naming a generic- or `Self`-typed variable only type-checks once +//! those are re-declared on it and instantiated where it is referred to. +//! [`EnclosingContext`] carries what to re-declare, threaded in by +//! `#[thrust_macros::context]`; without it a formula only sees concrete types. + +use proc_macro2::TokenStream as TokenStream2; +use quote::{format_ident, quote, ToTokens}; +use syn::{ + parse::{Parse, ParseStream}, + visit_mut::VisitMut, + FnArg, GenericParam, Signature, WherePredicate, +}; + +use crate::{fn_outer_item::FnOuterItem, FormulaFnTypeLowering}; + +/// The context a formula is written in: the host function's signature and, for a method, +/// its `impl`/`trait` header. +pub struct EnclosingContext { + sig: Signature, + outer: Option, +} + +impl Parse for EnclosingContext { + 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::()?; + Ok(Self { sig, outer }) + } +} + +impl EnclosingContext { + /// The generic params in scope: the host signature's own, plus the outer + /// `impl`/`trait`'s for a method. + fn generic_params(&self) -> impl Iterator { + self.sig + .generics + .params + .iter() + .chain(self.outer.iter().flat_map(|o| o.generics().params.iter())) + } + + /// 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()))) + } + + fn type_lowering(&self) -> FormulaFnTypeLowering<'_> { + match &self.outer { + Some(outer) => FormulaFnTypeLowering::with_outer_context(&self.sig, outer), + None => FormulaFnTypeLowering::new(&self.sig), + } + } +} + +/// A spec macro's closure together with the context it was written in, the form +/// `#[thrust_macros::context]` rewrites the macro's tokens into. +pub struct ClosureWithContext { + pub context: EnclosingContext, + pub closure: syn::ExprClosure, +} + +impl Parse for ClosureWithContext { + fn parse(input: ParseStream) -> syn::Result { + let context = input.parse()?; + let closure = input.parse()?; + Ok(Self { context, closure }) + } +} + +/// A formula lifted into a `#[thrust::formula_fn]` item, with the expression naming that +/// item — generic arguments included — for the marker call to take. +pub struct LiftedFormulaFn { + pub item: syn::ItemFn, + pub reference: syn::Expr, +} + +/// Lifts `body`, a formula over `params`, into a `#[thrust::formula_fn]` called `name`. +/// +/// The parameters are lowered to their model types. The receiver `self` is renamed to +/// `__thrust_self`, which the analyzer binds back to the receiver value. Under a +/// `context`, the generics in scope there are re-declared on the item and instantiated at +/// the reference; in a method, `Self` becomes the concrete self type of the `impl`, or a +/// synthetic type parameter instantiated with the real `Self` in a trait. +pub fn lift( + name: &syn::Ident, + params: &[FnArg], + body: &syn::Expr, + context: Option<&EnclosingContext>, +) -> syn::Result { + let mut params = params.to_vec(); + let mut body = body.clone(); + + let mut def_params: Vec = Vec::new(); + let mut turbofish_args: Vec = Vec::new(); + for param in context + .into_iter() + .flat_map(EnclosingContext::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(EnclosingContext::where_predicates) + .cloned() + .collect(); + + let dummy_sig = syn::parse_quote!(fn f()); + let type_lowering = match context { + Some(context) => context.type_lowering(), + None => FormulaFnTypeLowering::new(&dummy_sig), + }; + + def_wheres.extend(type_lowering.model_where_predicates()); + + // A formula may refer to the receiver value `self`; the lifted formula function is free, so + // rewrite `self` to a `__thrust_self` parameter. The analyzer binds it back to the receiver. + let mut rewriter = SelfValueRewriter { + to: format_ident!("__thrust_self"), + }; + for param in &mut params { + rewriter.visit_fn_arg_mut(param); + } + rewriter.visit_expr_mut(&mut body); + + let self_used = params + .iter() + .any(|param| crate::tokens_contain_ident(¶m.to_token_stream(), "Self")) + || crate::tokens_contain_ident(&body.to_token_stream(), "Self") + || def_wheres + .iter() + .any(|pred| crate::tokens_contain_ident(&pred.to_token_stream(), "Self")); + if self_used { + let Some(outer) = context.and_then(|context| context.outer.as_ref()) else { + return Err(syn::Error::new_spanned( + body, + "formula cannot refer to `Self` without an enclosing impl/trait context", + )); + }; + + match outer { + FnOuterItem::ItemImpl(item_impl) => { + // `Self` in an impl method context: rewrite it to the concrete self type everywhere + // TODO: Support generic/trait impl + let self_ty = &item_impl.self_ty; + let mut rewriter = SelfTypeRewriter { + to: *self_ty.clone(), + }; + for param in &mut params { + rewriter.visit_fn_arg_mut(param); + } + rewriter.visit_expr_mut(&mut body); + for pred in &mut def_wheres { + rewriter.visit_where_predicate_mut(pred); + } + } + FnOuterItem::ItemTrait(item_trait) => { + // `Self` in a trait method context: rewrite it to a synthetic generic everywhere + // it reaches the formula function — parameters, body, and the propagated + // where-clause predicates — then pass the real `Self` via turbofish (legal + // in expression position). + let synth: syn::Ident = format_ident!("__ThrustSelf"); + def_wheres.push(syn::parse_quote!(#synth: ?Sized)); + + let mut rewriter = SelfTypeRewriter { + to: syn::parse_quote!(#synth), + }; + for param in &mut params { + rewriter.visit_fn_arg_mut(param); + } + rewriter.visit_expr_mut(&mut body); + for pred in &mut def_wheres { + rewriter.visit_where_predicate_mut(pred); + } + def_params.push(quote!(#synth)); + def_wheres.extend(type_lowering.model_where_predicates_for(&synth)); + + // Mirror the host's implicit `Self: Trait` bound onto the synthetic + // generic so trait associated types (`Self::Item`) and predicates + // (`Self::step`) remain resolvable on it. + let trait_ident = &item_trait.ident; + let (_, ty_generics, _) = item_trait.generics.split_for_impl(); + def_wheres.push(syn::parse_quote!(#synth: #trait_ident #ty_generics)); + + turbofish_args.push(quote!(Self)); + + // Rewriting `Self` to the synthetic generic can yield predicates that + // duplicate the synthetic generic's own `Model` bounds; drop the dups. + let mut seen = std::collections::HashSet::new(); + def_wheres.retain(|pred| seen.insert(pred.to_token_stream().to_string())); + } + } + } + + let model_ty_params = type_lowering.lower_params(¶ms); + + 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),*>) + }; + + Ok(LiftedFormulaFn { + item: syn::parse_quote!( + #[allow(unused_variables)] + #[allow(non_snake_case)] + #[thrust::formula_fn] + fn #name #def_generics(#model_ty_params) -> bool #where_clause { + #body + } + ), + reference: syn::parse_quote!(#name #turbofish), + }) +} + +struct SelfValueRewriter { + to: syn::Ident, +} + +impl VisitMut for SelfValueRewriter { + fn visit_pat_ident_mut(&mut self, pat: &mut syn::PatIdent) { + if pat.ident == "self" { + pat.ident = self.to.clone(); + } + syn::visit_mut::visit_pat_ident_mut(self, pat); + } + + fn visit_fn_arg_mut(&mut self, arg: &mut syn::FnArg) { + match arg { + syn::FnArg::Receiver(receiver) => { + let to = &self.to; + let ty = crate::receiver_type(receiver); + *arg = syn::parse_quote!(#to: #ty); + } + syn::FnArg::Typed(_) => { /* handled by visit_pat_ident_mut */ } + } + + syn::visit_mut::visit_fn_arg_mut(self, arg); + } + + fn visit_expr_path_mut(&mut self, expr_path: &mut syn::ExprPath) { + if expr_path.qself.is_some() { + syn::visit_mut::visit_expr_path_mut(self, expr_path); + return; + } + + if expr_path.path.leading_colon.is_some() || expr_path.path.segments.len() != 1 { + syn::visit_mut::visit_expr_path_mut(self, expr_path); + return; + } + + if expr_path.path.segments[0].ident == "self" { + expr_path.path.segments[0].ident = self.to.clone(); + return; + } + + syn::visit_mut::visit_expr_path_mut(self, expr_path); + } + + fn visit_macro_mut(&mut self, mac: &mut syn::Macro) { + if !is_formula_macro(&mac.path) { + syn::visit_mut::visit_macro_mut(self, mac); + return; + } + + let expanded = crate::formula::expand(mac.tokens.clone()); + let Ok(mut expr) = syn::parse2::(expanded) else { + return; + }; + self.visit_expr_mut(&mut expr); + mac.tokens = expr.into_token_stream(); + } +} + +struct SelfTypeRewriter { + to: syn::Type, +} + +impl VisitMut for SelfTypeRewriter { + fn visit_type_mut(&mut self, ty: &mut syn::Type) { + syn::visit_mut::visit_type_mut(self, ty); + + let syn::Type::Path(type_path) = ty else { + return; + }; + + if type_path.qself.is_some() || type_path.path.leading_colon.is_some() { + return; + } + + let mut segments = type_path.path.segments.iter(); + + if segments.next().is_none_or(|first| first.ident != "Self") { + return; + } + + let tail: syn::punctuated::Punctuated<_, syn::Token![::]> = segments.cloned().collect(); + + if tail.is_empty() { + *ty = self.to.clone(); + } else { + let to = &self.to; + *ty = syn::parse_quote!(<#to>::#tail) + }; + } + + fn visit_expr_path_mut(&mut self, expr_path: &mut syn::ExprPath) { + syn::visit_mut::visit_expr_path_mut(self, expr_path); + + if expr_path.qself.is_some() || expr_path.path.leading_colon.is_some() { + return; + } + + let mut segments = expr_path.path.segments.iter(); + + if segments.next().is_none_or(|first| first.ident != "Self") { + return; + } + + let tail: syn::punctuated::Punctuated<_, syn::Token![::]> = segments.cloned().collect(); + + if tail.is_empty() { + return; + } + + let to = &self.to; + *expr_path = syn::parse_quote!(<#to>::#tail); + } + + fn visit_macro_mut(&mut self, mac: &mut syn::Macro) { + if !is_formula_macro(&mac.path) { + syn::visit_mut::visit_macro_mut(self, mac); + return; + } + + let expanded = crate::formula::expand(mac.tokens.clone()); + let Ok(mut expr) = syn::parse2::(expanded) else { + return; + }; + self.visit_expr_mut(&mut expr); + mac.tokens = expr.into_token_stream(); + } +} + +fn is_formula_macro(path: &syn::Path) -> bool { + // TODO: identify the macro precisely + path.segments + .last() + .is_some_and(|seg| seg.ident == "formula") +} diff --git a/thrust-macros/src/ghost.rs b/thrust-macros/src/ghost.rs index 794b173d..163d76ac 100644 --- a/thrust-macros/src/ghost.rs +++ b/thrust-macros/src/ghost.rs @@ -1,9 +1,16 @@ -//! Expansion of `thrust_macros::ghost!` into a `#[thrust::formula_fn]` relating the +//! Expansion of `thrust_macros::ghost!` and its context-carrying sibling +//! `thrust_macros::_ghost_with_context!` into a `#[thrust::formula_fn]` relating the //! introduced value to the ghost term, plus a marker call the analyzer intercepts. //! //! The value is parameter `0`, as in an `ensures` formula function: that is the layout //! the analyzer reads the term back with. Unlike `ensures`, it is bound to a synthetic //! name, leaving `result` free for a term to name a live variable with. +//! +//! `ghost!(|x: i64| -> Seq { .. })` only sees concrete types. +//! `_ghost_with_context!(..)` additionally carries the enclosing generic context (see +//! [`mod@crate::formula_fn_lifting`]), so a term may name generic- and `Self`-typed +//! variables; `#[thrust_macros::context]` rewrites each `ghost!` it finds into +//! that form. use std::sync::atomic::{AtomicUsize, Ordering}; @@ -11,23 +18,41 @@ use proc_macro::TokenStream; use quote::{format_ident, ToTokens}; use syn::FnArg; -use crate::FormulaFnTypeLowering; +use crate::formula_fn_lifting::{self, ClosureWithContext, EnclosingContext, LiftedFormulaFn}; static COUNTER: AtomicUsize = AtomicUsize::new(0); +/// Expands `ghost!(CLOSURE)`: a bare ghost term with no threaded context. pub fn expand(input: TokenStream) -> TokenStream { let input = crate::formula::wrap_closure_body(input.into()); let closure = match syn::parse2::(input) { Ok(closure) => closure, Err(e) => return e.to_compile_error().into(), }; - match expand_ghost(&closure) { + match expand_ghost(&closure, None) { + Ok(expr) => expr.into_token_stream().into(), + Err(e) => e.to_compile_error().into(), + } +} + +/// Expands `_ghost_with_context!(#outer_attr #sig; CLOSURE)`, the form +/// `#[thrust_macros::context]` rewrites each `ghost!` into. +pub fn expand_with_context(input: TokenStream) -> TokenStream { + let input = crate::formula::wrap_closure_body(input.into()); + let ClosureWithContext { closure, context } = match syn::parse2::(input) { + Ok(parsed) => parsed, + Err(e) => return e.to_compile_error().into(), + }; + match expand_ghost(&closure, Some(&context)) { Ok(expr) => expr.into_token_stream().into(), Err(e) => e.to_compile_error().into(), } } -fn expand_ghost(closure: &syn::ExprClosure) -> syn::Result { +fn expand_ghost( + closure: &syn::ExprClosure, + context: Option<&EnclosingContext>, +) -> syn::Result { let syn::ReturnType::Type(_, value_ty) = &closure.output else { return Err(syn::Error::new_spanned( closure, @@ -37,7 +62,7 @@ fn expand_ghost(closure: &syn::ExprClosure) -> syn::Result { // Synthetic, so that a term may name a live variable called `result`. let value = format_ident!("__thrust_ghost_value"); - let mut fn_params: Vec = vec![syn::parse_quote!(#value: #value_ty)]; + let mut params: Vec = vec![syn::parse_quote!(#value: #value_ty)]; for param in &closure.inputs { let syn::Pat::Type(pt) = param else { return Err(syn::Error::new_spanned( @@ -47,22 +72,20 @@ fn expand_ghost(closure: &syn::ExprClosure) -> syn::Result { }; let pat = &pt.pat; let ty = &pt.ty; - fn_params.push(syn::parse_quote!(#pat: #ty)); + params.push(syn::parse_quote!(#pat: #ty)); } - let dummy_sig = syn::parse_quote!(fn f()); - let model_ty_params = FormulaFnTypeLowering::new(&dummy_sig).lower_params(&fn_params); + let term = &closure.body; + let body = syn::parse_quote!(#value == (#term)); - let body = &closure.body; let id = COUNTER.fetch_add(1, Ordering::Relaxed); let name = format_ident!("_thrust_ghost_{}", id); + let LiftedFormulaFn { item, reference } = + formula_fn_lifting::lift(&name, ¶ms, &body, context)?; Ok(syn::parse_quote!({ - #[thrust::formula_fn] - fn #name(#model_ty_params) -> bool { - #value == (#body) - } + #item - thrust_models::__ghost_marker::<_, #value_ty>(#name) + thrust_models::__ghost_marker::<_, #value_ty>(#reference) })) } diff --git a/thrust-macros/src/invariant.rs b/thrust-macros/src/invariant.rs index 9c519aa8..79a38d97 100644 --- a/thrust-macros/src/invariant.rs +++ b/thrust-macros/src/invariant.rs @@ -8,11 +8,11 @@ //! - `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::context]` -//! rewrites each `invariant!` it finds into this form, pasting the host -//! function's signature (and, in methods, a `#[thrust::_outer_context(..)]` -//! attribute carrying the enclosing `impl`/`trait` header) ahead of the -//! closure: +//! context (see [`mod@crate::formula_fn_lifting`]). It is never written by hand: +//! `#[thrust_macros::context]` rewrites each `invariant!` it finds into this form, +//! pasting the host function's signature (and, in methods, a +//! `#[thrust::_outer_context(..)]` attribute carrying the enclosing `impl`/`trait` +//! header) ahead of the closure: //! //! ```ignore //! _invariant_with_context!( @@ -21,24 +21,14 @@ //! |x: T, v: T| x == v //! ) //! ``` -//! -//! 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). use std::sync::atomic::{AtomicUsize, Ordering}; use proc_macro::TokenStream; -use proc_macro2::TokenStream as TokenStream2; -use quote::{format_ident, quote, ToTokens}; -use syn::{ - parse::{Parse, ParseStream}, - visit_mut::VisitMut, - FnArg, GenericParam, Signature, WherePredicate, -}; +use quote::{format_ident, ToTokens}; +use syn::FnArg; -use crate::{fn_outer_item::FnOuterItem, FormulaFnTypeLowering}; +use crate::formula_fn_lifting::{self, ClosureWithContext, EnclosingContext, LiftedFormulaFn}; static COUNTER: AtomicUsize = AtomicUsize::new(0); @@ -59,27 +49,8 @@ pub fn expand(input: TokenStream) -> TokenStream { /// Expands `_invariant_with_context!(#outer_attr #sig; CLOSURE)`, the form /// `#[thrust_macros::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 input = crate::formula::wrap_closure_body(input.into()); - let WithContext { closure, context } = match syn::parse2::(input) { + let ClosureWithContext { closure, context } = match syn::parse2::(input) { Ok(parsed) => parsed, Err(e) => return e.to_compile_error().into(), }; @@ -89,51 +60,12 @@ pub fn expand_with_context(input: TokenStream) -> TokenStream { } } -/// The enclosing context threaded into an invariant by -/// `#[thrust_macros::context]`: the host function signature and, for a -/// method, its `impl`/`trait` header. A standalone `invariant!` has none. -struct Context { - sig: Signature, - outer: Option, -} - -impl Context { - /// The generic params in scope: the host signature's own, plus the outer - /// `impl`/`trait`'s for a method. - fn generic_params(&self) -> impl Iterator { - self.sig - .generics - .params - .iter() - .chain(self.outer.iter().flat_map(|o| o.generics().params.iter())) - } - - /// 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()))) - } - - fn type_lowering(&self) -> FormulaFnTypeLowering<'_> { - if let Some(outer) = &self.outer { - FormulaFnTypeLowering::with_outer_context(&self.sig, outer) - } else { - FormulaFnTypeLowering::new(&self.sig) - } - } -} - -/// 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. +/// Expands a predicate closure into a `#[thrust::formula_fn]` plus a marker call. fn expand_invariant( closure: &syn::ExprClosure, - context: Option<&Context>, + context: Option<&EnclosingContext>, ) -> syn::Result { - let mut fn_params: Vec = Vec::new(); + let mut params: Vec = Vec::new(); for param in &closure.inputs { let syn::Pat::Type(pt) = param else { return Err(syn::Error::new_spanned( @@ -143,279 +75,17 @@ fn expand_invariant( }; 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 turbofish_args: Vec = Vec::new(); - 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(); - - let dummy_sig = syn::parse_quote!(fn f()); - let type_lowering = if let Some(context) = context { - context.type_lowering() - } else { - FormulaFnTypeLowering::new(&dummy_sig) - }; - - def_wheres.extend(type_lowering.model_where_predicates()); - - let mut body = closure.body.clone(); - - // An invariant may refer to the receiver value `self`; the lifted formula function is free, so - // rewrite `self` to a `__thrust_self` parameter. The analyzer binds it back to the loop-carried receiver. - let mut rewriter = SelfValueRewriter { - to: format_ident!("__thrust_self"), - }; - for param in &mut fn_params { - rewriter.visit_fn_arg_mut(param); + params.push(syn::parse_quote!(#pat: #ty)); } - rewriter.visit_expr_mut(&mut body); - - let self_used = crate::tokens_contain_ident(&closure.to_token_stream(), "Self") - || def_wheres - .iter() - .any(|pred| crate::tokens_contain_ident(&pred.to_token_stream(), "Self")); - if self_used { - let Some(outer) = context.and_then(|context| context.outer.as_ref()) else { - return Err(syn::Error::new_spanned( - closure, - "invariant closure cannot refer to `Self` without an enclosing impl/trait context", - )); - }; - - match outer { - FnOuterItem::ItemImpl(item_impl) => { - // `Self` in an impl method context: rewrite it to the concrete self type everywhere - // TODO: Support generic/trait impl - let self_ty = &item_impl.self_ty; - let mut rewriter = SelfTypeRewriter { - to: *self_ty.clone(), - }; - for param in &mut fn_params { - rewriter.visit_fn_arg_mut(param); - } - rewriter.visit_expr_mut(&mut body); - for pred in &mut def_wheres { - rewriter.visit_where_predicate_mut(pred); - } - } - FnOuterItem::ItemTrait(item_trait) => { - // `Self` in a trait method context: rewrite it to a synthetic generic everywhere - // it reaches the formula function — parameters, body, and the propagated - // where-clause predicates — then pass the real `Self` via turbofish (legal - // in expression position). - let synth: syn::Ident = format_ident!("__ThrustSelf"); - def_wheres.push(syn::parse_quote!(#synth: ?Sized)); - - let mut rewriter = SelfTypeRewriter { - to: syn::parse_quote!(#synth), - }; - for param in &mut fn_params { - rewriter.visit_fn_arg_mut(param); - } - rewriter.visit_expr_mut(&mut body); - for pred in &mut def_wheres { - rewriter.visit_where_predicate_mut(pred); - } - def_params.push(quote!(#synth)); - def_wheres.extend(type_lowering.model_where_predicates_for(&synth)); - - // Mirror the host's implicit `Self: Trait` bound onto the synthetic - // generic so trait associated types (`Self::Item`) and predicates - // (`Self::step`) remain resolvable on it. - let trait_ident = &item_trait.ident; - let (_, ty_generics, _) = item_trait.generics.split_for_impl(); - def_wheres.push(syn::parse_quote!(#synth: #trait_ident #ty_generics)); - - turbofish_args.push(quote!(Self)); - - // Rewriting `Self` to the synthetic generic can yield predicates that - // duplicate the synthetic generic's own `Model` bounds; drop the dups. - let mut seen = std::collections::HashSet::new(); - def_wheres.retain(|pred| seen.insert(pred.to_token_stream().to_string())); - } - } - } - - let model_ty_params = type_lowering.lower_params(&fn_params); - let body = &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),*>) - }; + let LiftedFormulaFn { item, reference } = + formula_fn_lifting::lift(&name, ¶ms, &closure.body, context)?; Ok(syn::parse_quote!({ - #[allow(unused_variables)] - #[allow(non_snake_case)] - #[thrust::formula_fn] - fn #name #def_generics(#model_ty_params) -> bool #where_clause { - #body - } + #item - thrust_models::__invariant_marker(#name #turbofish) + thrust_models::__invariant_marker(#reference) })) } - -struct SelfValueRewriter { - to: syn::Ident, -} - -impl VisitMut for SelfValueRewriter { - fn visit_pat_ident_mut(&mut self, pat: &mut syn::PatIdent) { - if pat.ident == "self" { - pat.ident = self.to.clone(); - } - syn::visit_mut::visit_pat_ident_mut(self, pat); - } - - fn visit_fn_arg_mut(&mut self, arg: &mut syn::FnArg) { - match arg { - syn::FnArg::Receiver(receiver) => { - let to = &self.to; - let ty = crate::receiver_type(receiver); - *arg = syn::parse_quote!(#to: #ty); - } - syn::FnArg::Typed(_) => { /* handled by visit_pat_ident_mut */ } - } - - syn::visit_mut::visit_fn_arg_mut(self, arg); - } - - fn visit_expr_path_mut(&mut self, expr_path: &mut syn::ExprPath) { - if expr_path.qself.is_some() { - syn::visit_mut::visit_expr_path_mut(self, expr_path); - return; - } - - if expr_path.path.leading_colon.is_some() || expr_path.path.segments.len() != 1 { - syn::visit_mut::visit_expr_path_mut(self, expr_path); - return; - } - - if expr_path.path.segments[0].ident == "self" { - expr_path.path.segments[0].ident = self.to.clone(); - return; - } - - syn::visit_mut::visit_expr_path_mut(self, expr_path); - } - - fn visit_macro_mut(&mut self, mac: &mut syn::Macro) { - if !is_formula_macro(&mac.path) { - syn::visit_mut::visit_macro_mut(self, mac); - return; - } - - let expanded = crate::formula::expand(mac.tokens.clone()); - let Ok(mut expr) = syn::parse2::(expanded) else { - return; - }; - self.visit_expr_mut(&mut expr); - mac.tokens = expr.into_token_stream(); - } -} - -struct SelfTypeRewriter { - to: syn::Type, -} - -impl VisitMut for SelfTypeRewriter { - fn visit_type_mut(&mut self, ty: &mut syn::Type) { - syn::visit_mut::visit_type_mut(self, ty); - - let syn::Type::Path(type_path) = ty else { - return; - }; - - if type_path.qself.is_some() || type_path.path.leading_colon.is_some() { - return; - } - - let mut segments = type_path.path.segments.iter(); - - if segments.next().is_none_or(|first| first.ident != "Self") { - return; - } - - let tail: syn::punctuated::Punctuated<_, syn::Token![::]> = segments.cloned().collect(); - - if tail.is_empty() { - *ty = self.to.clone(); - } else { - let to = &self.to; - *ty = syn::parse_quote!(<#to>::#tail) - }; - } - - fn visit_expr_path_mut(&mut self, expr_path: &mut syn::ExprPath) { - syn::visit_mut::visit_expr_path_mut(self, expr_path); - - if expr_path.qself.is_some() || expr_path.path.leading_colon.is_some() { - return; - } - - let mut segments = expr_path.path.segments.iter(); - - if segments.next().is_none_or(|first| first.ident != "Self") { - return; - } - - let tail: syn::punctuated::Punctuated<_, syn::Token![::]> = segments.cloned().collect(); - - if tail.is_empty() { - return; - } - - let to = &self.to; - *expr_path = syn::parse_quote!(<#to>::#tail); - } - - fn visit_macro_mut(&mut self, mac: &mut syn::Macro) { - if !is_formula_macro(&mac.path) { - syn::visit_mut::visit_macro_mut(self, mac); - return; - } - - let expanded = crate::formula::expand(mac.tokens.clone()); - let Ok(mut expr) = syn::parse2::(expanded) else { - return; - }; - self.visit_expr_mut(&mut expr); - mac.tokens = expr.into_token_stream(); - } -} - -fn is_formula_macro(path: &syn::Path) -> bool { - // TODO: identify the macro precisely - path.segments - .last() - .is_some_and(|seg| seg.ident == "formula") -} diff --git a/thrust-macros/src/lib.rs b/thrust-macros/src/lib.rs index 2d81ec0d..aa4aecea 100644 --- a/thrust-macros/src/lib.rs +++ b/thrust-macros/src/lib.rs @@ -5,6 +5,7 @@ mod closure; mod context; mod fn_outer_item; mod formula; +mod formula_fn_lifting; mod formula_fn_type_lowering; mod ghost; mod invariant; @@ -52,11 +53,21 @@ pub fn ghost(input: TokenStream) -> TokenStream { ghost::expand(input) } +/// Context-carrying counterpart of `ghost!`, emitted by +/// `#[thrust_macros::context]`. Not intended to be written by hand: +/// it takes a `fn` header carrying the threaded generics/where clause whose +/// body is the ghost term closure (see [`ghost`]). +#[proc_macro] +pub fn _ghost_with_context(input: TokenStream) -> TokenStream { + ghost::expand_with_context(input) +} + /// Makes the enclosing context available to the specifications written inside an /// item. On an `impl`/`trait`, each method recovers the outer generics (and `Self`) /// in its `requires`/`ensures`; on a function — including a method reached that way — -/// every `thrust_macros::invariant!(...)` in the body may refer to generic- and -/// `Self`-typed variables that the standalone macro cannot see. See [`mod@context`]. +/// every `thrust_macros::invariant!(...)` and `thrust_macros::ghost!(...)` in the body +/// may refer to generic- and `Self`-typed variables that the standalone macros cannot +/// see. See [`mod@context`]. #[proc_macro_attribute] pub fn context(_attr: TokenStream, item: TokenStream) -> TokenStream { context::expand(item)