diff --git a/src/analyze/annot.rs b/src/analyze/annot.rs index 869b9bf9..a4df8ae0 100644 --- a/src/analyze/annot.rs +++ b/src/analyze/annot.rs @@ -46,6 +46,10 @@ pub fn refinement_path_path() -> [Symbol; 2] { [Symbol::intern("thrust"), Symbol::intern("refinement_path")] } +pub fn closure_upvars_path() -> [Symbol; 2] { + [Symbol::intern("thrust"), Symbol::intern("closure_upvars")] +} + pub fn model_ty_path() -> [Symbol; 3] { [ Symbol::intern("thrust"), diff --git a/src/analyze/annot_fn.rs b/src/analyze/annot_fn.rs index 9aba2f3c..1eea3402 100644 --- a/src/analyze/annot_fn.rs +++ b/src/analyze/annot_fn.rs @@ -211,6 +211,10 @@ impl<'a, 'tcx> AnnotFnTranslator<'a, 'tcx> { fn build_env_from_params(&mut self) { for (idx, param) in self.body.params.iter().enumerate() { let param_idx = rty::FunctionParamIdx::from(idx); + if self.is_closure_upvars_param(param) { + self.build_env_from_captures(chc::Term::var(param_idx), param.pat); + continue; + } let mir_ty = self.pat_ty(param.pat); // `at_entry()` yields the `Inner` of a `FnParam`; classify by it so // a singleton wrapped argument collapses like any other singleton below. @@ -227,6 +231,67 @@ impl<'a, 'tcx> AnnotFnTranslator<'a, 'tcx> { } } + /// Whether the parameter holds the closure's upvars, as marked by + /// `#[thrust::closure_upvars]` on a `closure!` specification. + fn is_closure_upvars_param(&self, param: &rustc_hir::Param<'tcx>) -> bool { + let attr_path = analyze::annot::closure_upvars_path(); + self.tcx + .hir_attrs(param.hir_id) + .iter() + .any(|attr| attr.path_matches(&attr_path)) + } + + /// Binds the names of a closure specification's upvars pattern to the closure's + /// captured variables. + /// + /// The pattern lists the captures a clause names, in the order it wrote them, + /// while the upvars hold every capture in the order rustc chose. The two are + /// therefore matched up by name. + fn build_env_from_captures( + &mut self, + upvars_term: chc::Term, + pat: &'tcx rustc_hir::Pat<'tcx>, + ) { + let rustc_hir::PatKind::Tuple(subpats, _) = pat.kind else { + panic!( + "closure upvars are expected to be a tuple pattern: {:?}", + pat + ); + }; + let closure_def_id = self.tcx.local_parent(self.local_def_id); + let captures = self.tcx.closure_captures(closure_def_id); + let closure_ty = self.tcx.type_of(closure_def_id).instantiate_identity(); + let mir_ty::TyKind::Closure(_, closure_args) = closure_ty.kind() else { + panic!("closure specification is expected to sit inside a closure"); + }; + let mut upvar_terms = Vec::new(); + for i in 0..captures.len() { + let upvar_term = match closure_args.as_closure().kind() { + mir_ty::ClosureKind::Fn => upvars_term.clone().box_current().tuple_proj(i), + mir_ty::ClosureKind::FnMut => chc::Term::mut_( + upvars_term.clone().mut_current().tuple_proj(i), + upvars_term.clone().mut_final().tuple_proj(i), + ), + mir_ty::ClosureKind::FnOnce => upvars_term.clone().tuple_proj(i), + }; + upvar_terms.push(upvar_term); + } + for subpat in subpats { + let rustc_hir::PatKind::Binding(_, hir_id, ident, None) = subpat.kind else { + panic!("closure capture is expected to be a binding: {:?}", subpat); + }; + let Some(idx) = captures.iter().position(|capture| { + capture.var_ident.name == ident.name && capture.place.projections.is_empty() + }) else { + self.tcx.dcx().span_fatal( + subpat.span, + format!("`{}` is not captured by this closure", ident), + ); + }; + self.env.insert(hir_id, upvar_terms[idx].clone()); + } + } + fn singleton_term_for_ty( ty: &rty::Type, ) -> Option> { @@ -378,7 +443,7 @@ impl<'a, 'tcx> AnnotFnTranslator<'a, 'tcx> { } /// The values of a closure's parameters: the closure's first (RustCall) parameter is its - /// environment, which is the closure value itself, followed by the logical arguments. + /// upvars, which are the closure value itself, followed by the logical arguments. fn translate_closure_precondition( &self, receiver: &'tcx rustc_hir::Expr<'tcx>, @@ -394,8 +459,8 @@ impl<'a, 'tcx> AnnotFnTranslator<'a, 'tcx> { ) }); let logical_args = self.closure_spec_args(args); - // `fn_ty` is a closure (RustCall ABI), so its parameters are `[env, args..]`, where the - // environment is the closure value itself. + // `fn_ty` is a closure (RustCall ABI), so its parameters are `[upvars, args..]`, where + // the upvars are the closure value itself. assert_eq!( logical_args.len(), fn_ty.params.len() - 1, @@ -424,8 +489,8 @@ impl<'a, 'tcx> AnnotFnTranslator<'a, 'tcx> { ) }); let logical_args = self.closure_spec_args(args); - // `fn_ty` is a closure (RustCall ABI), so its parameters are `[env, args..]`, where the - // environment is the closure value itself. + // `fn_ty` is a closure (RustCall ABI), so its parameters are `[upvars, args..]`, where + // the upvars are the closure value itself. assert_eq!( logical_args.len(), fn_ty.params.len() - 1, diff --git a/tests/ui/fail/closure_captures.rs b/tests/ui/fail/closure_captures.rs new file mode 100644 index 00000000..0da513e9 --- /dev/null +++ b/tests/ui/fail/closure_captures.rs @@ -0,0 +1,20 @@ +//@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 n = 5; + let f = thrust_macros::closure!( + captures(n: i32), + requires(x > n), + ensures(result == x + n), + |x: i32| -> i32 { x + n }, + ); + let r = apply(3, f); + assert!(r == 8); +} diff --git a/tests/ui/fail/closure_captures_fn_mut.rs b/tests/ui/fail/closure_captures_fn_mut.rs new file mode 100644 index 00000000..91aaae16 --- /dev/null +++ b/tests/ui/fail/closure_captures_fn_mut.rs @@ -0,0 +1,16 @@ +//@error-in-other-file: Unsat +//@compile-flags: -C debug-assertions=off + +fn main() { + let mut acc = 0; + let mut f = thrust_macros::closure!( + captures(acc: &mut &mut i32), + ensures(result == x + 1 && *(!acc) == *(*acc) + 1), + |x: i32| -> i32 { + acc += 1; + x + acc + }, + ); + let r = f(3); + assert!(r == 5); +} diff --git a/tests/ui/fail/closure_captures_fn_once.rs b/tests/ui/fail/closure_captures_fn_once.rs new file mode 100644 index 00000000..bd2efe00 --- /dev/null +++ b/tests/ui/fail/closure_captures_fn_once.rs @@ -0,0 +1,21 @@ +//@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 mut acc = 0; + let r = apply( + 3, + thrust_macros::closure!( + captures(acc: &mut i32), + ensures(result == x + 1 && !acc == *acc + 1), + |x: i32| -> i32 { acc += 1; x + acc }, + ), + ); + assert!(r == 5); +} diff --git a/tests/ui/pass/closure_captures.rs b/tests/ui/pass/closure_captures.rs new file mode 100644 index 00000000..560f80da --- /dev/null +++ b/tests/ui/pass/closure_captures.rs @@ -0,0 +1,20 @@ +//@check-pass +//@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 n = 5; + let f = thrust_macros::closure!( + captures(n: i32), + requires(x > n), + ensures(result == x + n), + |x: i32| -> i32 { x + n }, + ); + let r = apply(7, f); + assert!(r == 12); +} diff --git a/tests/ui/pass/closure_captures_fn_mut.rs b/tests/ui/pass/closure_captures_fn_mut.rs new file mode 100644 index 00000000..d5023dbe --- /dev/null +++ b/tests/ui/pass/closure_captures_fn_mut.rs @@ -0,0 +1,20 @@ +//@check-pass +//@compile-flags: -C debug-assertions=off + +// A `FnMut` closure holds its upvars behind a `Mut`, so a mutable-borrow capture has +// two levels: the outer holds the slot on entry (`*acc`) and on exit (`!acc`), the +// inner is the borrow, whose current value is the counter. Hence `*(!acc)`, not +// `!(*acc)`. +fn main() { + let mut acc = 0; + let mut f = thrust_macros::closure!( + captures(acc: &mut &mut i32), + ensures(result == x + 1 && *(!acc) == *(*acc) + 1), + |x: i32| -> i32 { + acc += 1; + x + acc + }, + ); + let r = f(3); + assert!(r == 4); +} diff --git a/tests/ui/pass/closure_captures_fn_once.rs b/tests/ui/pass/closure_captures_fn_once.rs new file mode 100644 index 00000000..a9115dbe --- /dev/null +++ b/tests/ui/pass/closure_captures_fn_once.rs @@ -0,0 +1,24 @@ +//@check-pass +//@compile-flags: -C debug-assertions=off + +// Passed straight to `apply` to keep the closure `FnOnce`: binding it to a `let` first +// makes it `FnMut`, and `pre!`/`post!` hand a `FnMut` closure upvars stripped of their +// `Mut`. +#[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 mut acc = 0; + let r = apply( + 3, + thrust_macros::closure!( + captures(acc: &mut i32), + ensures(result == x + 1 && !acc == *acc + 1), + |x: i32| -> i32 { acc += 1; x + acc }, + ), + ); + assert!(r == 4); +} diff --git a/thrust-macros/src/closure.rs b/thrust-macros/src/closure.rs index 9c154544..33845a9f 100644 --- a/thrust-macros/src/closure.rs +++ b/thrust-macros/src/closure.rs @@ -3,9 +3,10 @@ //! //! ```ignore //! let f = thrust_macros::closure!( +//! captures(n: i32), //! requires(x > 0), -//! ensures(result > x), -//! |x: i32| -> i32 { x + 1 }, +//! ensures(result > x + n), +//! |x: i32| -> i32 { x + n + 1 }, //! ); //! ``` //! @@ -16,6 +17,10 @@ //! `spec.rs`). Each clause is optional (an omitted one leaves that side inferred) //! and may be repeated, in which case its predicates are conjoined. //! +//! `captures` restates the captured variables a clause wants to name, with the types +//! they are captured by the closure. Only the ones a clause names need restating, and +//! in any order: the plugin matches them against the closure's real captures by name. +//! //! A clause sees no threaded generic or `Self` context, so a closure in a generic //! context cannot refer to generic- or `Self`-typed values. @@ -31,11 +36,13 @@ use syn::{ use crate::FormulaFnTypeLowering; mod kw { + syn::custom_keyword!(captures); syn::custom_keyword!(requires); syn::custom_keyword!(ensures); } struct ClosureSpec { + captures: Vec, requires: Vec, ensures: Vec, closure: syn::ExprClosure, @@ -43,22 +50,30 @@ struct ClosureSpec { impl Parse for ClosureSpec { fn parse(input: ParseStream) -> syn::Result { + let mut captures = Vec::new(); let mut requires = Vec::new(); let mut ensures = Vec::new(); loop { - let clause = if input.peek(kw::requires) { - input.parse::()?; - &mut requires - } else if input.peek(kw::ensures) { - input.parse::()?; - &mut ensures + if input.peek(kw::captures) { + input.parse::()?; + let content; + parenthesized!(content in input); + captures.extend(content.parse_terminated(FnArg::parse, syn::Token![,])?); } else { - break; - }; - let content; - parenthesized!(content in input); - clause.push(content.parse()?); + let clause = if input.peek(kw::requires) { + input.parse::()?; + &mut requires + } else if input.peek(kw::ensures) { + input.parse::()?; + &mut ensures + } else { + break; + }; + let content; + parenthesized!(content in input); + clause.push(content.parse()?); + } input.parse::>()?; } @@ -66,6 +81,7 @@ impl Parse for ClosureSpec { input.parse::>()?; Ok(Self { + captures, requires, ensures, closure, @@ -86,15 +102,14 @@ pub fn expand(input: TokenStream) -> TokenStream { fn expand_closure(spec: ClosureSpec) -> syn::Result { let ClosureSpec { + captures, requires, ensures, mut closure, } = spec; - // A closure's parameters are `[env, arg1, .., argN]`, the environment being the - // closure value itself. A clause names only the arguments, so the companions take - // a dummy parameter in the environment's place to keep the positions aligned. - let mut fn_params: Vec = vec![syn::parse_quote!(_thrust_closure_env: ())]; + let upvars = upvars_param(&captures)?; + let mut arg_params: Vec = Vec::new(); for param in &closure.inputs { let syn::Pat::Type(pt) = param else { return Err(syn::Error::new_spanned( @@ -104,7 +119,7 @@ fn expand_closure(spec: ClosureSpec) -> syn::Result { }; let pat = &pt.pat; let ty = &pt.ty; - fn_params.push(syn::parse_quote!(#pat: #ty)); + arg_params.push(syn::parse_quote!(#pat: #ty)); } if !ensures.is_empty() && matches!(closure.output, syn::ReturnType::Default) { @@ -118,14 +133,21 @@ fn expand_closure(spec: ClosureSpec) -> syn::Result { // clause has none of its own. let spec_sig: syn::Signature = syn::parse_quote!(fn closure_spec()); let type_lowering = FormulaFnTypeLowering::new(&spec_sig); - let model_params = type_lowering.lower_params(&fn_params); + // A closure's parameters are `[upvars, arg1, .., argN]`. The companions take the + // upvars in that same leading position, so their parameters line up with the + // closure's. + let upvars_model = type_lowering.lower_params([&upvars]); + let arg_models = type_lowering.lower_params(&arg_params); let mut prelude: Vec = Vec::new(); if let Some(body) = conjoin(requires) { prelude.push(quote! { #[allow(unused_variables, non_snake_case)] #[thrust::formula_fn] - fn _thrust_closure_requires(#model_params) -> bool { + fn _thrust_closure_requires( + #[thrust::closure_upvars] #upvars_model, + #arg_models + ) -> bool { #body } @@ -138,7 +160,11 @@ fn expand_closure(spec: ClosureSpec) -> syn::Result { prelude.push(quote! { #[allow(unused_variables, non_snake_case)] #[thrust::formula_fn] - fn _thrust_closure_ensures(result: #ret_model, #model_params) -> bool { + fn _thrust_closure_ensures( + result: #ret_model, + #[thrust::closure_upvars] #upvars_model, + #arg_models + ) -> bool { #body } @@ -164,6 +190,24 @@ fn expand_closure(spec: ClosureSpec) -> syn::Result { Ok(closure) } +/// The companion parameter holding the closure's upvars: a tuple of the captures a +/// clause names, which the plugin matches up with the real upvars by name. +fn upvars_param(captures: &[FnArg]) -> syn::Result { + let mut names = Vec::new(); + let mut tys = Vec::new(); + for capture in captures { + let FnArg::Typed(capture) = capture else { + return Err(syn::Error::new_spanned( + capture, + "closure! captures are written as `name: Type`", + )); + }; + names.push(&capture.pat); + tys.push(&capture.ty); + } + Ok(syn::parse_quote!((#(#names,)*): (#(#tys,)*))) +} + fn conjoin(preds: Vec) -> Option { preds .into_iter()