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_env_path() -> [Symbol; 2] {
+    [Symbol::intern("thrust"), Symbol::intern("closure_env")]
+}
+
 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_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<Inner>`; 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<rty::FunctionParamIdx>,
+        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<rty::FunctionParamIdx>,
+        capture: &mir_ty::CapturedPlace<'tcx>,
+        upvar_ty: mir_ty::Ty<'tcx>,
+        subpat: &'tcx rustc_hir::Pat<'tcx>,
+    ) -> chc::Term<rty::FunctionParamIdx> {
+        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<rty::Closed>,
     ) -> Option<chc::Term<rty::FunctionParamIdx>> {

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

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

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

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

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