Relate OuterFnParam to locals via FnParam::at_here in loop invariants

src/analyze/annot.rs

@@ -233,6 +233,22 @@ pub fn fn_param_at_entry_path() -> [Symbol; 3] {
     ]
 }
 
+pub fn fn_param_at_here_path() -> [Symbol; 3] {
+    [
+        Symbol::intern("thrust"),
+        Symbol::intern("def"),
+        Symbol::intern("fn_param_at_here"),
+    ]
+}
+
+pub fn fn_param_is_not_changed_path() -> [Symbol; 3] {
+    [
+        Symbol::intern("thrust"),
+        Symbol::intern("def"),
+        Symbol::intern("fn_param_is_not_changed"),
+    ]
+}
+
 pub fn closure_precondition_path() -> [Symbol; 3] {
     [
         Symbol::intern("thrust"),

src/analyze/annot_fn.rs

@@ -264,6 +264,43 @@ impl<'a, 'tcx> AnnotFnTranslator<'a, 'tcx> {
             .is_some_and(|def| Some(def.did()) == self.def_ids.fn_param_wrapper())
     }
 
+    /// A `FnParam<T>` invariant parameter denotes the function argument's value at function entry
+    /// (`OuterFnParam`); `at_entry()` is therefore the identity on the parameter.
+    fn fn_param_at_entry_term(
+        &self,
+        receiver: &'tcx rustc_hir::Expr<'tcx>,
+    ) -> chc::Term<rty::FunctionParamIdx> {
+        self.to_term(receiver)
+    }
+
+    /// `at_here()` denotes the argument's *current* local at the loop header, which is a distinct
+    /// basic-block parameter from the entry value. The two are wired to separate parameters in
+    /// [`crate::analyze::local_def::Analyzer::build_invariant_precondition`]: each invariant
+    /// parameter `i` has its entry value at variable `i` and its "here" value at variable
+    /// `params.len() + i`.
+    fn fn_param_at_here_term(
+        &self,
+        receiver: &'tcx rustc_hir::Expr<'tcx>,
+    ) -> chc::Term<rty::FunctionParamIdx> {
+        let chc::Term::Var(idx) = self.to_term(receiver) else {
+            panic!("FnParam::at_here receiver must be a function-parameter binding");
+        };
+        chc::Term::var(rty::FunctionParamIdx::from(
+            self.body.params.len() + idx.index(),
+        ))
+    }
+
+    fn fn_param_is_not_changed(
+        &self,
+        receiver: &'tcx rustc_hir::Expr<'tcx>,
+    ) -> FormulaOrTerm<rty::FunctionParamIdx> {
+        FormulaOrTerm::BinOp(
+            self.fn_param_at_entry_term(receiver),
+            AmbiguousBinOp::Eq,
+            self.fn_param_at_here_term(receiver),
+        )
+    }
+
     fn build_env_from_pat(
         &mut self,
         param: chc::Term<rty::FunctionParamIdx>,
@@ -669,8 +706,21 @@ impl<'a, 'tcx> AnnotFnTranslator<'a, 'tcx> {
                             args.is_empty(),
                             "FnParam::at_entry does not take any arguments"
                         );
-                        let t = self.to_term(receiver);
-                        return FormulaOrTerm::Term(t);
+                        return FormulaOrTerm::Term(self.fn_param_at_entry_term(receiver));
+                    }
+                    if Some(def_id) == self.def_ids.fn_param_at_here() {
+                        assert!(
+                            args.is_empty(),
+                            "FnParam::at_here does not take any arguments"
+                        );
+                        return FormulaOrTerm::Term(self.fn_param_at_here_term(receiver));
+                    }
+                    if Some(def_id) == self.def_ids.fn_param_is_not_changed() {
+                        assert!(
+                            args.is_empty(),
+                            "FnParam::is_not_changed does not take any arguments"
+                        );
+                        return self.fn_param_is_not_changed(receiver);
                     }
                     if Some(def_id) == self.def_ids.seq_len() {
                         assert!(args.is_empty(), "Seq::len does not take any arguments");
@@ -759,8 +809,19 @@ impl<'a, 'tcx> AnnotFnTranslator<'a, 'tcx> {
                         }
                         if Some(def_id) == self.def_ids.fn_param_at_entry() {
                             assert_eq!(args.len(), 1, "FnParam::at_entry takes exactly 1 argument");
-                            let t = self.to_term(&args[0]);
-                            return FormulaOrTerm::Term(t);
+                            return FormulaOrTerm::Term(self.fn_param_at_entry_term(&args[0]));
+                        }
+                        if Some(def_id) == self.def_ids.fn_param_at_here() {
+                            assert_eq!(args.len(), 1, "FnParam::at_here takes exactly 1 argument");
+                            return FormulaOrTerm::Term(self.fn_param_at_here_term(&args[0]));
+                        }
+                        if Some(def_id) == self.def_ids.fn_param_is_not_changed() {
+                            assert_eq!(
+                                args.len(),
+                                1,
+                                "FnParam::is_not_changed takes exactly 1 argument"
+                            );
+                            return self.fn_param_is_not_changed(&args[0]);
                         }
                         if Some(def_id) == self.def_ids.implies() {
                             let [lhs, rhs] = args else {

src/analyze/basic_block.rs

@@ -1305,19 +1305,6 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
                     let var = self.env.immut_bind_tmp(param_unrefined_rty);
                     param_terms.insert(param_idx, chc::Term::var(var.into()));
                     outer_fn_param_vars.insert(outer_idx, var);
-
-                    // XXX: A user-supplied loop invariant replaces the inferred precondition, so
-                    // it no longer carries the relationship between a function's entry argument
-                    // (`OuterFnParam`) and its current local. The invariant cannot express both
-                    // views of the same argument with the current syntax. As an ad-hoc workaround,
-                    // equate them when the local is non-flow, since its representation cannot change.
-                    let local = analyze::local_of_function_param(outer_idx);
-                    if !param_ty.to_sort().is_singleton() && self.env.is_non_flow_local(local) {
-                        assumption.body.push_conj(
-                            chc::Term::var(PlaceTypeVar::Var(var))
-                                .equal_to(self.env.local_type(local).term),
-                        );
-                    }
                 }
                 BasicBlockTypeParamKind::Synthetic => {}
             }

src/analyze/did_cache.rs

@@ -38,6 +38,8 @@ struct DefIds {
 
     fn_param_wrapper: OnceCell<Option<DefId>>,
     fn_param_at_entry: OnceCell<Option<DefId>>,
+    fn_param_at_here: OnceCell<Option<DefId>>,
+    fn_param_is_not_changed: OnceCell<Option<DefId>>,
 
     closure_precondition: OnceCell<Option<DefId>>,
     closure_postcondition: OnceCell<Option<DefId>>,
@@ -270,6 +272,19 @@ impl<'tcx> DefIdCache<'tcx> {
             .get_or_init(|| self.annotated_def(&crate::analyze::annot::fn_param_at_entry_path()))
     }
 
+    pub fn fn_param_at_here(&self) -> Option<DefId> {
+        *self
+            .def_ids
+            .fn_param_at_here
+            .get_or_init(|| self.annotated_def(&crate::analyze::annot::fn_param_at_here_path()))
+    }
+
+    pub fn fn_param_is_not_changed(&self) -> Option<DefId> {
+        *self.def_ids.fn_param_is_not_changed.get_or_init(|| {
+            self.annotated_def(&crate::analyze::annot::fn_param_is_not_changed_path())
+        })
+    }
+
     pub fn closure_precondition(&self) -> Option<DefId> {
         *self
             .def_ids

src/analyze/local_def.rs

@@ -1009,7 +1009,18 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
             .fn_sig(formula_def_id.to_def_id())
             .instantiate(self.tcx, generic_args);
 
-        let mut mapping: Vec<rty::FunctionParamIdx> = Vec::with_capacity(idents.len());
+        // Each invariant parameter `i` is substituted by a term over `bty`'s parameters. A
+        // `FnParam<T>` parameter denotes the argument's value at function entry (`OuterFnParam`);
+        // its current local at the loop header (`at_here()`) is a *separate* basic-block parameter.
+        // To expose both from a single invariant parameter without conflating them, variable `i`
+        // carries the entry value and variable `params.len() + i` carries the "here" value (see
+        // `AnnotFnTranslator::fn_param_at_here_term`). Ordinary parameters map to their loop-header
+        // local and have no distinct "here" value.
+        let free = |idx| chc::Term::var(rty::RefinedTypeVar::Free(idx));
+        let mut entry_mapping: Vec<chc::Term<rty::RefinedTypeVar<rty::FunctionParamIdx>>> =
+            Vec::with_capacity(idents.len());
+        let mut here_mapping: Vec<chc::Term<rty::RefinedTypeVar<rty::FunctionParamIdx>>> =
+            Vec::with_capacity(idents.len());
         for (ident_opt, input_ty) in idents.iter().zip(sig.skip_binder().inputs()) {
             let name = ident_opt.expect("invariant parameters must be named").name;
             let input_ty = {
@@ -1038,21 +1049,44 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
                     ))
                 });
                 let param_idx = crate::analyze::function_param_of_local(local);
-                mapping.push(bty.param_of_outer_fn_param(param_idx).unwrap());
+                let at_entry = free(bty.param_of_outer_fn_param(param_idx).unwrap());
+                // `at_here()` is the argument's current local at the loop header. It may have been
+                // moved or dropped (not live), in which case there is no distinct "here" value and
+                // we fall back to the entry value.
+                let at_here = bty
+                    .param_of_local(local)
+                    .map(free)
+                    .unwrap_or_else(|| at_entry.clone());
+                entry_mapping.push(at_entry);
+                here_mapping.push(at_here);
             } else {
+                // Referring to a function argument's local without `FnParam` is ambiguous about
+                // which view (entry vs. current) is meant, so it is rejected. The receiver `self`
+                // is exempt: it is loop-carried and always denotes its current value.
+                if name.as_str() != "self" && self.function_param_local_of_name(name).is_some() {
+                    self.tcx.dcx().fatal(format!(
+                        "loop invariant refers to function parameter `{name}` directly; \
+                         use `{name}: FnParam<..>` with `.at_entry()`/`.at_here()` to choose \
+                         between its value at function entry and its current value"
+                    ))
+                }
                 let local = self.local_of_name_in_bb(name, bty).unwrap_or_else(|| {
                     self.tcx.dcx().fatal(format!(
                         "loop invariant refers to `{name}`, which is not a live variable at the loop header"
                     ))
                 });
-                mapping.push(bty.param_of_local(local).unwrap());
+                let term = free(bty.param_of_local(local).unwrap());
+                // Ordinary parameters have no separate "here" value; the slot is never referenced.
+                entry_mapping.push(term.clone());
+                here_mapping.push(term);
             }
         }
 
+        let mapping: Vec<_> = entry_mapping.into_iter().chain(here_mapping).collect();
         formula_fn
             .formula()
             .clone()
-            .subst_var(|idx| chc::Term::var(rty::RefinedTypeVar::Free(mapping[idx.index()])))
+            .subst_var(|idx| mapping[idx.index()].clone())
             .into()
     }
 

src/refine/env.rs

@@ -922,10 +922,6 @@ where
         self.locals.contains_key(&local) || self.flow_locals.contains_key(&local)
     }
 
-    pub fn is_non_flow_local(&self, local: Local) -> bool {
-        self.locals.contains_key(&local)
-    }
-
     fn flow_binding(&self, var: Var) -> Option<&FlowBinding> {
         match var {
             Var::Local(local) => self.flow_locals.get(&local),

std.rs

@@ -415,6 +415,20 @@ mod thrust_models {
         pub fn at_entry(self) -> T {
             unimplemented!()
         }
+
+        #[allow(dead_code)]
+        #[thrust::def::fn_param_at_here]
+        #[thrust::ignored]
+        pub fn at_here(self) -> T {
+            unimplemented!()
+        }
+
+        #[allow(dead_code)]
+        #[thrust::def::fn_param_is_not_changed]
+        #[thrust::ignored]
+        pub fn is_not_changed(&self) -> bool {
+            unimplemented!()
+        }
     }
 }
 

tests/ui/fail/loop_invariant_fn_param_at_here.rs

@@ -0,0 +1,26 @@
+//@error-in-other-file: Unsat
+//@compile-flags: -C debug-assertions=off
+
+// `a` is mutated in the loop, so its current value (`at_here`) diverges from its entry value.
+// Claiming `is_not_changed()` (i.e. `at_entry() == at_here()`) is therefore not inductive: it holds
+// on entry but is broken by `a += 1`, so verification fails.
+
+#[thrust_macros::requires(true)]
+#[thrust_macros::ensures(true)]
+#[thrust::trusted]
+fn rand() -> bool {
+    unimplemented!()
+}
+
+#[thrust_macros::ensures(result == a)]
+#[thrust_macros::invariant_context]
+fn count_up(mut a: i64) -> i64 {
+    while rand() {
+        thrust_macros::invariant!(|a: thrust_models::FnParam<i64>| a.is_not_changed());
+        a += 1;
+    }
+
+    a
+}
+
+fn main() {}

tests/ui/fail/loop_invariant_fn_param_at_here_mut_borrow.rs

@@ -0,0 +1,27 @@
+//@error-in-other-file: Unsat
+//@compile-flags: -C debug-assertions=off
+
+// `a` is mutably borrowed and incremented through the borrow, so its current value (`at_here`)
+// diverges from its entry value. Claiming `is_not_changed()` (i.e. `at_entry() == at_here()`) is
+// therefore not inductive: it holds on entry but is broken by `*r += 1`, so verification fails.
+
+#[thrust_macros::requires(true)]
+#[thrust_macros::ensures(true)]
+#[thrust::trusted]
+fn rand() -> bool {
+    unimplemented!()
+}
+
+#[thrust_macros::ensures(result == a)]
+#[thrust_macros::invariant_context]
+fn count_up(mut a: i64) -> i64 {
+    while rand() {
+        thrust_macros::invariant!(|a: thrust_models::FnParam<i64>| a.is_not_changed());
+        let r = &mut a;
+        *r += 1;
+    }
+
+    a
+}
+
+fn main() {}

tests/ui/fail/loop_invariant_fn_param_direct.rs

@@ -0,0 +1,28 @@
+//@error-in-other-file: refers to function parameter `a` directly
+//@compile-flags: -C debug-assertions=off
+
+// Naively referring to a function argument's local directly (without `FnParam`) is rejected: the
+// choice between its entry value and its current value would be implicit. Compare with the paired
+// pass test, which makes that choice explicit via `FnParam` and `at_entry()`/`at_here()`.
+
+#[thrust_macros::requires(true)]
+#[thrust_macros::ensures(true)]
+#[thrust::trusted]
+fn rand() -> bool {
+    unimplemented!()
+}
+
+#[thrust_macros::ensures(result == a)]
+#[thrust_macros::invariant_context]
+fn keep_argument(a: i64) -> i64 {
+    let mut v = a;
+
+    while rand() {
+        thrust_macros::invariant!(|a: i64, v: i64| v == a);
+        v = a;
+    }
+
+    v
+}
+
+fn main() {}

tests/ui/fail/loop_invariant_generic.rs

@@ -10,7 +10,7 @@ fn rand() -> i64 { unimplemented!() }
 fn keep<T: Copy + PartialEq>(v: T) {
   let mut x = v;
   while rand() == 0 {
-    thrust_macros::invariant!(|v: T| v == v);
+    thrust_macros::invariant!(|v: thrust_models::FnParam<T>| v.is_not_changed());
     x = v;
   }
   assert!(x == v);

tests/ui/fail/loop_invariant_outer_param.rs

@@ -1,6 +1,10 @@
 //@error-in-other-file: Unsat
 //@compile-flags: -C debug-assertions=off
 
+// Without `is_not_changed()` the invariant only relates `v` to the argument's *current* local
+// (`at_here`); nothing ties that back to the entry value, so the `result == a` postcondition (whose
+// `a` is the entry value) cannot be discharged.
+
 #[thrust_macros::requires(true)]
 #[thrust_macros::ensures(true)]
 #[thrust::trusted]
@@ -14,7 +18,7 @@ fn keep_argument(a: i64) -> i64 {
     let mut v = a;
 
     while rand() {
-        thrust_macros::invariant!(|a: i64, v: i64| true);
+        thrust_macros::invariant!(|a: thrust_models::FnParam<i64>, v: i64| v == a.at_here());
         v = a;
     }