From 97421eb307de2817a2db53f3db0843dc898dfec9 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 23 Jul 2026 11:29:17 +0000 Subject: [PATCH] Allow user-specified closure pre/post via closure! macro MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closure pre-/post-conditions were always inferred as predicate-variable templates. Add `thrust_macros::closure!(requires(..), ensures(..), |x: T| -> R { .. })` so a user can pin them explicitly, mirroring Prusti's `closure!`. The macro rewrites the closure body to carry `#[thrust::formula_fn]` companions plus `#[thrust::requires_path]` / `#[thrust::ensures_path]` path statements — the same markers the plugin already reads for named-fn specs. Each clause is optional; an omitted side stays inferred (a pvar). Plugin changes: - `expected_ty` shifts a closure formula's argument indices by one to account for the environment parameter that leads a closure's `FunctionType` params. - `refine_local_defs` registers all formula fns in a pre-pass so a closure's spec companions (nested in its own body) resolve regardless of `mir_keys` order. Adds passing/failing UI test pairs for both-clause and ensures-only closures. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_019mT8dkzKVXeCA4DhFTn9wj --- src/analyze/crate_.rs | 18 ++ src/analyze/local_def.rs | 46 ++++- tests/ui/fail/closure_ensures_only.rs | 17 ++ tests/ui/fail/closure_requires_ensures.rs | 20 ++ tests/ui/pass/closure_ensures_only.rs | 19 ++ tests/ui/pass/closure_requires_ensures.rs | 20 ++ thrust-macros/src/closure.rs | 217 ++++++++++++++++++++++ thrust-macros/src/lib.rs | 10 + 8 files changed, 360 insertions(+), 7 deletions(-) create mode 100644 tests/ui/fail/closure_ensures_only.rs create mode 100644 tests/ui/fail/closure_requires_ensures.rs create mode 100644 tests/ui/pass/closure_ensures_only.rs create mode 100644 tests/ui/pass/closure_requires_ensures.rs create mode 100644 thrust-macros/src/closure.rs diff --git a/src/analyze/crate_.rs b/src/analyze/crate_.rs index 74198f18..5edd8c8d 100644 --- a/src/analyze/crate_.rs +++ b/src/analyze/crate_.rs @@ -60,6 +60,24 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> { } fn refine_local_defs(&mut self) { + // Pre-pass: register every formula fn before any def's type is computed, so + // that a def's `expected_ty` can resolve the formula fns its + // `requires`/`ensures` refer to regardless of the order `mir_keys` yields + // them. A closure's spec formula fns are nested inside the closure body and + // would otherwise be visited only after the closure itself. + for local_def_id in self.tcx.mir_keys(()) { + if !self.tcx.def_kind(*local_def_id).is_fn_like() { + continue; + } + if self + .ctx + .local_def_analyzer(*local_def_id) + .is_annotated_as_formula_fn() + { + self.ctx.register_formula_fn(*local_def_id); + self.skip_analysis.insert(*local_def_id); + } + } for local_def_id in self.tcx.mir_keys(()) { if self.tcx.def_kind(*local_def_id).is_fn_like() { self.refine_fn_def(*local_def_id); diff --git a/src/analyze/local_def.rs b/src/analyze/local_def.rs index 68b0d646..6d952d45 100644 --- a/src/analyze/local_def.rs +++ b/src/analyze/local_def.rs @@ -281,18 +281,50 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> { let trait_item_ty = self.trait_item_ty(); let is_fully_annotated = self.is_fully_annotated(); + // A closure's `FunctionType` carries its environment as the first parameter + // (`[env, arg1, .., argN]`), but a user-written `requires`/`ensures` formula + // function is expressed over the logical arguments only (`[arg1, .., argN]`). + // For closures we therefore shift the formula's argument indices by one to + // skip the environment param. + let is_closure = matches!( + self.tcx.def_kind(self.local_def_id), + rustc_hir::def::DefKind::Closure + ); let mut builder = self.type_builder.for_function_template(&mut self.ctx, sig); if let Some(require) = require_annot { - let formula = require.map_var(|idx| { - if idx.index() == sig.inputs().len() - 1 { - rty::RefinedTypeVar::Value - } else { - rty::RefinedTypeVar::Free(idx) - } - }); + let formula = if is_closure { + let num_args = sig.inputs().len() - 1; + require.map_var(|idx| { + if num_args > 0 && idx.index() == num_args - 1 { + rty::RefinedTypeVar::Value + } else { + rty::RefinedTypeVar::Free(rty::FunctionParamIdx::from_usize( + idx.index() + 1, + )) + } + }) + } else { + require.map_var(|idx| { + if idx.index() == sig.inputs().len() - 1 { + rty::RefinedTypeVar::Value + } else { + rty::RefinedTypeVar::Free(idx) + } + }) + }; builder.param_refinement(formula.into()); } if let Some(ensure) = ensure_annot { + let ensure = if is_closure { + ensure.map_var(|v| match v { + rty::RefinedTypeVar::Free(idx) => rty::RefinedTypeVar::Free( + rty::FunctionParamIdx::from_usize(idx.index() + 1), + ), + other => other, + }) + } else { + ensure + }; builder.ret_refinement(ensure.into()); } for (position, refinement) in refinement_annots { diff --git a/tests/ui/fail/closure_ensures_only.rs b/tests/ui/fail/closure_ensures_only.rs new file mode 100644 index 00000000..cd8294d1 --- /dev/null +++ b/tests/ui/fail/closure_ensures_only.rs @@ -0,0 +1,17 @@ +//@error-in-other-file: Unsat +//@compile-flags: -C debug-assertions=off + +#[thrust_macros::requires(thrust_macros::pre!(f(x)))] +#[thrust_macros::ensures(thrust_macros::post!(f(x), result))] +fn apply i32>(x: i32, f: F) -> i32 { + f(x) +} + +fn main() { + let f = thrust_macros::closure!( + ensures(result == x + 1), + |x: i32| -> i32 { x + 1 }, + ); + let r = apply(3, f); + assert!(r == 5); +} diff --git a/tests/ui/fail/closure_requires_ensures.rs b/tests/ui/fail/closure_requires_ensures.rs new file mode 100644 index 00000000..615c5782 --- /dev/null +++ b/tests/ui/fail/closure_requires_ensures.rs @@ -0,0 +1,20 @@ +//@error-in-other-file: Unsat +//@compile-flags: -C debug-assertions=off + +// The fixed postcondition `result == x + 1` gives `r == 4`, so asserting `r == 5` +// must fail verification. +#[thrust_macros::requires(thrust_macros::pre!(f(x)))] +#[thrust_macros::ensures(thrust_macros::post!(f(x), result))] +fn apply i32>(x: i32, f: F) -> i32 { + f(x) +} + +fn main() { + let f = thrust_macros::closure!( + requires(x > 0), + ensures(result == x + 1), + |x: i32| -> i32 { x + 1 }, + ); + let r = apply(3, f); + assert!(r == 5); +} diff --git a/tests/ui/pass/closure_ensures_only.rs b/tests/ui/pass/closure_ensures_only.rs new file mode 100644 index 00000000..b5462d7c --- /dev/null +++ b/tests/ui/pass/closure_ensures_only.rs @@ -0,0 +1,19 @@ +//@check-pass +//@compile-flags: -C debug-assertions=off + +// A closure that specifies only `ensures`; its precondition stays inferred as a +// predicate variable. +#[thrust_macros::requires(thrust_macros::pre!(f(x)))] +#[thrust_macros::ensures(thrust_macros::post!(f(x), result))] +fn apply i32>(x: i32, f: F) -> i32 { + f(x) +} + +fn main() { + let f = thrust_macros::closure!( + ensures(result == x + 1), + |x: i32| -> i32 { x + 1 }, + ); + let r = apply(3, f); + assert!(r == 4); +} diff --git a/tests/ui/pass/closure_requires_ensures.rs b/tests/ui/pass/closure_requires_ensures.rs new file mode 100644 index 00000000..b437abd4 --- /dev/null +++ b/tests/ui/pass/closure_requires_ensures.rs @@ -0,0 +1,20 @@ +//@check-pass +//@compile-flags: -C debug-assertions=off + +// A closure whose pre-/post-condition is given explicitly via `closure!` rather +// than inferred as predicate variables. The caller relies on the fixed spec. +#[thrust_macros::requires(thrust_macros::pre!(f(x)))] +#[thrust_macros::ensures(thrust_macros::post!(f(x), result))] +fn apply i32>(x: i32, f: F) -> i32 { + f(x) +} + +fn main() { + let f = thrust_macros::closure!( + requires(x > 0), + ensures(result == x + 1), + |x: i32| -> i32 { x + 1 }, + ); + let r = apply(3, f); + assert!(r == 4); +} diff --git a/thrust-macros/src/closure.rs b/thrust-macros/src/closure.rs new file mode 100644 index 00000000..f7ad803d --- /dev/null +++ b/thrust-macros/src/closure.rs @@ -0,0 +1,217 @@ +//! Expansion of `thrust_macros::closure!`, which attaches an explicit +//! `requires`/`ensures` specification to a closure expression. +//! +//! Rust attributes cannot sit on a closure *expression*, so the spec is written +//! as leading `requires(..)` / `ensures(..)` clauses in a wrapper macro: +//! +//! ```ignore +//! let f = thrust_macros::closure!( +//! requires(x > 0), +//! ensures(result == x + 1), +//! |x: i32| -> i32 { x + 1 }, +//! ); +//! ``` +//! +//! The expansion rewrites the closure so its body block gains, as leading +//! statements, `#[thrust::formula_fn]` companions plus `#[thrust::requires_path]` +//! / `#[thrust::ensures_path]` path statements — the exact same markers the +//! plugin already reads for named `fn` specs (see `spec.rs`). The plugin then +//! installs those formulas onto the closure's `FunctionType` in place of the +//! inferred predicate-variable templates. Each clause is optional: omit one and +//! that side stays inferred (a pvar). The closure header supplies the parameter +//! and return types, so nothing is restated. +//! +//! Generic/`Self` context threading (the `invariant_context` machinery) is not +//! supported here: closures in generic contexts that refer to generic- or +//! `Self`-typed values are out of scope for now. + +use std::sync::atomic::{AtomicUsize, Ordering}; + +use proc_macro::TokenStream; +use proc_macro2::TokenStream as TokenStream2; +use quote::{format_ident, quote, ToTokens}; +use syn::{ + parenthesized, + parse::{Parse, ParseStream}, + FnArg, +}; + +use crate::FormulaFnTypeLowering; + +mod kw { + syn::custom_keyword!(requires); + syn::custom_keyword!(ensures); +} + +static COUNTER: AtomicUsize = AtomicUsize::new(0); + +struct ClosureSpec { + /// Preprocessed `requires` predicates (already run through `formula::expand`). + requires: Vec, + /// Preprocessed `ensures` predicates. + ensures: Vec, + closure: syn::ExprClosure, +} + +impl Parse for ClosureSpec { + fn parse(input: ParseStream) -> syn::Result { + let mut requires = Vec::new(); + let mut ensures = Vec::new(); + + loop { + if input.peek(kw::requires) { + input.parse::()?; + let content; + parenthesized!(content in input); + let raw: TokenStream2 = content.parse()?; + requires.push(crate::formula::expand(raw)); + } else if input.peek(kw::ensures) { + input.parse::()?; + let content; + parenthesized!(content in input); + let raw: TokenStream2 = content.parse()?; + ensures.push(crate::formula::expand(raw)); + } else { + break; + } + input.parse::>()?; + } + + let closure: syn::ExprClosure = input.parse()?; + input.parse::>()?; + + Ok(Self { + requires, + ensures, + closure, + }) + } +} + +pub fn expand(input: TokenStream) -> TokenStream { + let spec = match syn::parse::(input) { + Ok(spec) => spec, + Err(e) => return e.to_compile_error().into(), + }; + match expand_closure(spec) { + Ok(expr) => expr.into_token_stream().into(), + Err(e) => e.to_compile_error().into(), + } +} + +fn expand_closure(spec: ClosureSpec) -> syn::Result { + let ClosureSpec { + requires, + ensures, + mut closure, + } = spec; + + // The closure parameters, restated as `fn` arguments so their types can be + // lowered to model types. + let mut fn_params: Vec = Vec::new(); + for param in &closure.inputs { + let syn::Pat::Type(pt) = param else { + return Err(syn::Error::new_spanned( + param, + "closure!'s closure parameters must have explicit types, e.g. `|x: i32| ...`", + )); + }; + let pat = &pt.pat; + let ty = &pt.ty; + fn_params.push(syn::parse_quote!(#pat: #ty)); + } + + if !ensures.is_empty() && matches!(closure.output, syn::ReturnType::Default) { + return Err(syn::Error::new_spanned( + &closure, + "closure! with `ensures` must declare an explicit return type, e.g. `|x: i32| -> i32 { .. }`", + )); + } + + let dummy_sig: syn::Signature = syn::parse_quote!(fn __thrust_closure_spec()); + let type_lowering = FormulaFnTypeLowering::new(&dummy_sig); + let model_params = type_lowering.lower_params(&fn_params); + + let id = COUNTER.fetch_add(1, Ordering::Relaxed); + + let mut prelude: Vec = Vec::new(); + + if !requires.is_empty() { + let name = format_ident!("_thrust_closure_requires_{}", id); + let body = conjoin(&requires); + prelude.push(quote! { + #[allow(unused_variables, non_snake_case)] + #[thrust::formula_fn] + fn #name(#model_params) -> bool { + #body + } + }); + prelude.push(quote! { + #[thrust::requires_path] + #name; + }); + } + + if !ensures.is_empty() { + let name = format_ident!("_thrust_closure_ensures_{}", id); + let body = conjoin(&ensures); + let ret_model = type_lowering.lower_return_type(&closure.output); + prelude.push(quote! { + #[allow(unused_variables, non_snake_case)] + #[thrust::formula_fn] + fn #name(result: #ret_model, #model_params) -> bool { + #body + } + }); + prelude.push(quote! { + #[thrust::ensures_path] + #name; + }); + } + + // Splice the prelude into the closure body. When the body is already a block, + // prepend into it (rather than nesting a second block) so the tail expression + // stays bare and no `unused_braces` warning is emitted. + let prelude_block: syn::Block = syn::parse_quote!({ #(#prelude)* }); + let new_body: syn::Expr = match *closure.body { + syn::Expr::Block(ref expr_block) + if expr_block.attrs.is_empty() && expr_block.label.is_none() => + { + let mut block = expr_block.block.clone(); + let mut stmts = prelude_block.stmts; + stmts.append(&mut block.stmts); + block.stmts = stmts; + syn::Expr::Block(syn::ExprBlock { + attrs: Vec::new(), + label: None, + block, + }) + } + ref orig_body => { + let mut block = prelude_block; + block.stmts.push(syn::Stmt::Expr(orig_body.clone(), None)); + syn::Expr::Block(syn::ExprBlock { + attrs: Vec::new(), + label: None, + block, + }) + } + }; + closure.body = Box::new(new_body); + + Ok(closure) +} + +/// Conjoins predicate expressions with `&&`, defaulting to `true` when empty. +fn conjoin(preds: &[TokenStream2]) -> TokenStream2 { + if preds.is_empty() { + return quote!(true); + } + let mut iter = preds.iter(); + let first = iter.next().unwrap(); + let mut acc = quote!((#first)); + for pred in iter { + acc = quote!(#acc && (#pred)); + } + acc +} diff --git a/thrust-macros/src/lib.rs b/thrust-macros/src/lib.rs index 2c89f573..6c37118a 100644 --- a/thrust-macros/src/lib.rs +++ b/thrust-macros/src/lib.rs @@ -1,6 +1,7 @@ use proc_macro::TokenStream; use proc_macro2::{TokenStream as TokenStream2, TokenTree as TokenTree2}; +mod closure; mod context; mod fn_outer_item; mod formula; @@ -28,6 +29,15 @@ pub fn post(input: TokenStream) -> TokenStream { pre_post::expand_post(input) } +/// `closure!(requires(..), ensures(..), |x: T| -> R { .. })` attaches an +/// explicit pre-/post-condition to a closure expression. Each of `requires` and +/// `ensures` is optional (omitting one leaves that side inferred) and may be +/// repeated (conjoined). See [`mod@closure`]. +#[proc_macro] +pub fn closure(input: TokenStream) -> TokenStream { + closure::expand(input) +} + #[proc_macro_attribute] pub fn context(_attr: TokenStream, item: TokenStream) -> TokenStream { context::expand(item)