Let a closure specification name the closure's captured variables

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"),

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<Inner>`; 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<rty::FunctionParamIdx>,
+ 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<rty::Closed>,
  ) -> Option<chc::Term<rty::FunctionParamIdx>> {
@@ -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,

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<F: FnOnce(i32) -> 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);
+}

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);
+}

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<F: FnOnce(i32) -> 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);
+}

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<F: FnOnce(i32) -> 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);
+}

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);
+}

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<F: FnOnce(i32) -> 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);
+}