Bridge the rust-call receiver in pre!/post!

src/analyze/annot_fn.rs

@@ -416,6 +416,41 @@ impl<'a, 'tcx> AnnotFnTranslator<'a, 'tcx> {
         }
     }
 
+    /// The receiver term as the closure's own body takes it.
+    ///
+    /// That body takes the upvars by `&`, by `&mut`, or by value, following the kind inferred
+    /// for the closure, while `pre!`/`post!` reach the closure through the parameter the
+    /// annotated function declares. A call bridges the two by borrowing the closure into the
+    /// receiver the body takes; the same borrow is taken here on the term.
+    fn closure_receiver_term(
+        &self,
+        receiver: &'tcx rustc_hir::Expr<'tcx>,
+        fn_ty: &rty::FunctionType,
+    ) -> chc::Term<rty::FunctionParamIdx> {
+        let held_as = match self.expr_ty(receiver).kind() {
+            mir_ty::TyKind::Adt(adt, _) if Some(adt.did()) == self.def_ids.mut_model() => {
+                Some(rty::RefKind::Mut)
+            }
+            mir_ty::TyKind::Ref(_, _, mir_ty::Mutability::Not) => Some(rty::RefKind::Immut),
+            _ => None,
+        };
+        let upvars_ty = &fn_ty.params[rty::FunctionParamIdx::from(0usize)].ty;
+        let received_as = match upvars_ty.as_pointer().map(|ty| ty.kind) {
+            Some(rty::PointerKind::Ref(kind)) => Some(kind),
+            _ => None,
+        };
+
+        let term = self.to_term(receiver);
+        match (held_as, received_as) {
+            (None, Some(rty::RefKind::Immut)) => chc::Term::box_(term),
+            (None, Some(rty::RefKind::Mut)) => chc::Term::mut_(term.clone(), term),
+            (Some(rty::RefKind::Mut), Some(rty::RefKind::Immut)) => {
+                chc::Term::box_(term.mut_current())
+            }
+            _ => term,
+        }
+    }
+
     /// Resolves the [`rty::FunctionType`] of the closure contract referred to by the receiver.
     ///
     /// The receiver type is instantiated to the actual closure type in the formula function; its
@@ -467,7 +502,7 @@ impl<'a, 'tcx> AnnotFnTranslator<'a, 'tcx> {
             "closure precondition arity mismatch: closure takes {} argument(s)",
             fn_ty.params.len() - 1
         );
-        let param_args: Vec<_> = std::iter::once(self.to_term(receiver))
+        let param_args: Vec<_> = std::iter::once(self.closure_receiver_term(receiver, &fn_ty))
             .chain(logical_args)
             .collect();
         FormulaOrTerm::Formula(fn_ty.precondition_formula(&param_args))
@@ -497,7 +532,7 @@ impl<'a, 'tcx> AnnotFnTranslator<'a, 'tcx> {
             "closure postcondition arity mismatch: closure takes {} argument(s)",
             fn_ty.params.len() - 1
         );
-        let param_args: Vec<_> = std::iter::once(self.to_term(receiver))
+        let param_args: Vec<_> = std::iter::once(self.closure_receiver_term(receiver, &fn_ty))
             .chain(logical_args)
             .collect();
         let result = self.to_term(result);

tests/ui/fail/closure_mut_capture_pre_post.rs

@@ -0,0 +1,19 @@
+//@error-in-other-file: Unsat
+//@compile-flags: -C debug-assertions=off
+
+#[thrust_macros::requires(thrust_macros::pre!(f()))]
+#[thrust_macros::ensures(thrust_macros::post!(f(), result))]
+fn call<F: FnMut() -> i64>(mut f: F) -> i64 {
+    f()
+}
+
+fn main() {
+    let mut cnt: i64 = 0;
+    let f = || -> i64 {
+        cnt += 1;
+        cnt
+    };
+    let r = call(f);
+    // `f` increments `cnt` once, so `r == 1`
+    assert!(r == 2);
+}

tests/ui/fail/closure_receiver_mut_model_byval.rs

@@ -0,0 +1,28 @@
+//@error-in-other-file: Unsat
+//@compile-flags: -C debug-assertions=off
+//@rustc-env: THRUST_SOLVER=tests/thrust-pcsat-wrapper
+
+use thrust_models::{
+    exists,
+    model::{Int, Mut},
+};
+
+#[thrust_macros::ensures(exists(|g, h, i: Int|
+  thrust_macros::post!(Mut::new(f, g)(), i)
+  && thrust_macros::post!(Mut::new(g, h)(), result)
+))]
+fn call_twice<F: FnMut() -> i64>(mut f: F) -> i64 {
+    f();
+    f()
+}
+
+fn main() {
+    let mut cnt: i64 = 0;
+    let f = move || -> i64 {
+        cnt += 1;
+        cnt
+    };
+    let r = call_twice(f);
+    // `f` increments `cnt` on each of the two calls, so `r == 2`
+    assert!(r == 3);
+}

tests/ui/fail/closure_ref_mut_pre_post.rs

@@ -0,0 +1,16 @@
+//@error-in-other-file: Unsat
+//@compile-flags: -C debug-assertions=off
+
+#[thrust_macros::requires(thrust_macros::pre!(f()))]
+#[thrust_macros::ensures(thrust_macros::post!(f(), result))]
+fn call<F: Fn() -> i64>(f: &mut F) -> i64 {
+    f()
+}
+
+fn main() {
+    let k: i64 = 1;
+    let mut f = || -> i64 { k };
+    let r = call(&mut f);
+    // `f` returns `k`, which is 1
+    assert!(r == 2);
+}

tests/ui/pass/closure_captures_fn_once.rs

@@ -2,8 +2,7 @@
 //@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`.
+// makes it `FnMut`, which holds its upvars behind another `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 {

tests/ui/pass/closure_mut_capture_pre_post.rs

@@ -0,0 +1,20 @@
+//@check-pass
+//@compile-flags: -C debug-assertions=off
+
+// A closure that mutates a capture receives its upvars behind a `Mut`, while the
+// higher-order function names the closure by value in `pre!`/`post!`.
+#[thrust_macros::requires(thrust_macros::pre!(f()))]
+#[thrust_macros::ensures(thrust_macros::post!(f(), result))]
+fn call<F: FnMut() -> i64>(mut f: F) -> i64 {
+    f()
+}
+
+fn main() {
+    let mut cnt: i64 = 0;
+    let f = || -> i64 {
+        cnt += 1;
+        cnt
+    };
+    let r = call(f);
+    assert!(r == 1);
+}

tests/ui/pass/closure_receiver_mut_model_byval.rs

@@ -0,0 +1,30 @@
+//@check-pass
+//@compile-flags: -C debug-assertions=off
+//@rustc-env: THRUST_SOLVER=tests/thrust-pcsat-wrapper
+
+use thrust_models::{
+    exists,
+    model::{Int, Mut},
+};
+
+// Naming the closure by value leaves its upvars as the call found them, which cannot carry
+// the upvars from one call to the next. `Mut::new` builds the receiver instead, naming the
+// upvars between the two calls.
+#[thrust_macros::ensures(exists(|g, h, i: Int|
+  thrust_macros::post!(Mut::new(f, g)(), i)
+  && thrust_macros::post!(Mut::new(g, h)(), result)
+))]
+fn call_twice<F: FnMut() -> i64>(mut f: F) -> i64 {
+    f();
+    f()
+}
+
+fn main() {
+    let mut cnt: i64 = 0;
+    let f = move || -> i64 {
+        cnt += 1;
+        cnt
+    };
+    let r = call_twice(f);
+    assert!(r == 2);
+}

tests/ui/pass/closure_ref_mut_pre_post.rs

@@ -0,0 +1,17 @@
+//@check-pass
+//@compile-flags: -C debug-assertions=off
+
+// The higher-order function names the closure through a `&mut` in `pre!`/`post!`, while a
+// closure that only reads its captures receives its upvars as they are.
+#[thrust_macros::requires(thrust_macros::pre!(f()))]
+#[thrust_macros::ensures(thrust_macros::post!(f(), result))]
+fn call<F: Fn() -> i64>(f: &mut F) -> i64 {
+    f()
+}
+
+fn main() {
+    let k: i64 = 1;
+    let mut f = || -> i64 { k };
+    let r = call(&mut f);
+    assert!(r == 1);
+}