diff --git a/src/analyze/annot.rs b/src/analyze/annot.rs index 869b9bf9..c6b7a288 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_env_path() -> [Symbol; 2] { + [Symbol::intern("thrust"), Symbol::intern("closure_env")] +} + 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..d81de20f 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_env_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,118 @@ impl<'a, 'tcx> AnnotFnTranslator<'a, 'tcx> { } } + /// Whether the parameter stands in for the closure environment, as marked by + /// `#[thrust::closure_env]` on a `closure!` specification. + fn is_closure_env_param(&self, param: &rustc_hir::Param<'tcx>) -> bool { + let attr_path = analyze::annot::closure_env_path(); + self.tcx + .hir_attrs(param.hir_id) + .iter() + .any(|attr| attr.path_matches(&attr_path)) + } + + /// Binds the names of a closure specification's environment pattern to the + /// closure's captured variables. + /// + /// The pattern lists the captures a clause names, in the order it wrote them, + /// while the environment holds every capture in the order rustc chose. The two + /// are therefore matched up by name. + fn build_env_from_captures( + &mut self, + env: chc::Term, + pat: &'tcx rustc_hir::Pat<'tcx>, + ) { + let rustc_hir::PatKind::Tuple(subpats, _) = pat.kind else { + panic!( + "closure environment is 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 upvar_tys = closure_args.as_closure().upvar_tys(); + // A closure called through `&mut self` receives its environment behind a `Mut`, + // a shape the analyzer does not yet represent consistently across a closure's + // definition and its call sites. + let env_ty = self.analyzer.fn_sig(closure_def_id.to_def_id()).inputs()[0]; + if !subpats.is_empty() + && matches!( + env_ty.kind(), + mir_ty::TyKind::Ref(_, _, mir_ty::Mutability::Mut) + ) + { + self.tcx.dcx().span_fatal( + pat.span, + "this closure is called through `&mut`, so a specification cannot name its captures yet", + ); + } + 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) + else { + self.tcx.dcx().span_fatal( + subpat.span, + format!("`{}` is not captured by this closure", ident), + ); + }; + let term = self.capture_term( + env.clone().tuple_proj(idx), + captures[idx], + upvar_tys[idx], + subpat, + ); + self.env.insert(hir_id, term); + } + } + + /// The value a capture name stands for, reporting against `subpat` when the + /// restated type does not describe what the closure captured. + /// + /// A shared borrow is read through, so that adding or removing `move` does not + /// change how a clause names the variable. A mutable borrow is left as the `Mut` + /// it is, since a clause has to say whether it means the value on entry or the + /// one on exit. + fn capture_term( + &self, + term: chc::Term, + capture: &mir_ty::CapturedPlace<'tcx>, + upvar_ty: mir_ty::Ty<'tcx>, + subpat: &'tcx rustc_hir::Pat<'tcx>, + ) -> chc::Term { + if !capture.place.projections.is_empty() { + self.tcx.dcx().span_fatal( + subpat.span, + format!( + "`{}` is captured field by field, which a closure specification cannot name", + capture.var_ident + ), + ); + } + let (named_ty, term) = match upvar_ty.kind() { + mir_ty::TyKind::Ref(_, referent_ty, mir_ty::Mutability::Not) => { + (*referent_ty, term.box_current()) + } + _ => (upvar_ty, term), + }; + if self.type_builder.build(self.pat_ty(subpat)).to_sort() + != self.type_builder.build(named_ty).to_sort() + { + self.tcx.dcx().span_fatal( + subpat.span, + format!("`{}` is captured as `{}`", capture.var_ident, named_ty), + ); + } + term + } + fn singleton_term_for_ty( ty: &rty::Type, ) -> Option> { diff --git a/tests/ui/fail/closure_captures.rs b/tests/ui/fail/closure_captures.rs new file mode 100644 index 00000000..d18c2d69 --- /dev/null +++ b/tests/ui/fail/closure_captures.rs @@ -0,0 +1,20 @@ +//@error-in-other-file: Unsat +//@compile-flags: -C debug-assertions=off + +// The declared postcondition carries the captured `n`, which is 5, so `r` is 8. +#[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), + ensures(result == x + n), + |x: i32| -> i32 { x + n }, + ); + let r = apply(3, f); + assert!(r == 9); +} diff --git a/tests/ui/fail/closure_captures_mut.rs b/tests/ui/fail/closure_captures_mut.rs new file mode 100644 index 00000000..046e867e --- /dev/null +++ b/tests/ui/fail/closure_captures_mut.rs @@ -0,0 +1,22 @@ +//@error-in-other-file: Unsat +//@compile-flags: -C debug-assertions=off + +// The declared postcondition pins `result` to `x + 1`, which is 4. +#[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/fail/closure_captures_order.rs b/tests/ui/fail/closure_captures_order.rs new file mode 100644 index 00000000..6670bdf2 --- /dev/null +++ b/tests/ui/fail/closure_captures_order.rs @@ -0,0 +1,22 @@ +//@error-in-other-file: Unsat +//@compile-flags: -C debug-assertions=off + +// `captures` lists `n` first while the closure captures `b` first; `n` still carries +// its own value, 5, so `r` is 8. +#[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 b = true; + let f = thrust_macros::closure!( + captures(n: i32, b: bool), + ensures(result == x + n), + move |x: i32| -> i32 { if b { x + n } else { x } }, + ); + let r = apply(3, f); + assert!(r == 9); +} diff --git a/tests/ui/pass/closure_captures.rs b/tests/ui/pass/closure_captures.rs new file mode 100644 index 00000000..51708fcf --- /dev/null +++ b/tests/ui/pass/closure_captures.rs @@ -0,0 +1,21 @@ +//@check-pass +//@compile-flags: -C debug-assertions=off + +// A closure specification naming a captured variable. `n` is captured by reference, +// which the specification reads through: the clause names the variable, not the borrow. +#[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), + ensures(result == x + n), + |x: i32| -> i32 { x + n }, + ); + let r = apply(3, f); + assert!(r == 8); +} diff --git a/tests/ui/pass/closure_captures_mut.rs b/tests/ui/pass/closure_captures_mut.rs new file mode 100644 index 00000000..0d658305 --- /dev/null +++ b/tests/ui/pass/closure_captures_mut.rs @@ -0,0 +1,26 @@ +//@check-pass +//@compile-flags: -C debug-assertions=off + +// A capture taken by mutable borrow is named as the `&mut` it is, so that the clause +// can say both what it was on entry (`*acc`) and what it becomes (`!acc`). +// +// The closure is passed straight to `apply`: binding it to a `let` first would have it +// called through `&mut`, which a specification cannot name its captures through yet. +#[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/tests/ui/pass/closure_captures_order.rs b/tests/ui/pass/closure_captures_order.rs new file mode 100644 index 00000000..8c234815 --- /dev/null +++ b/tests/ui/pass/closure_captures_order.rs @@ -0,0 +1,23 @@ +//@check-pass +//@compile-flags: -C debug-assertions=off + +// The closure captures `b` before `n`, since that is the order its body first uses +// them, while `captures` lists `n` first. Matching the two up by name is what makes +// `n` resolve to the second captured value rather than the first. +#[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 b = true; + let f = thrust_macros::closure!( + captures(n: i32, b: bool), + ensures(result == x + n), + move |x: i32| -> i32 { if b { x + n } else { x } }, + ); + let r = apply(3, f); + assert!(r == 8); +} diff --git a/thrust-macros/src/closure.rs b/thrust-macros/src/closure.rs index 9c154544..8a6217ea 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 have outside 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 env = env_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 `[env, arg1, .., argN]`, the environment holding its + // captures. The companions take the environment in that same leading position, so + // their parameters line up with the closure's. + let env_model = type_lowering.lower_params([&env]); + 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_env] #env_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_env] #env_model, + #arg_models + ) -> bool { #body } @@ -164,6 +190,24 @@ fn expand_closure(spec: ClosureSpec) -> syn::Result { Ok(closure) } +/// The companion parameter holding the closure environment: a tuple of the captures a +/// clause names, which the plugin matches up with the real environment by name. +fn env_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()