[codex] Preserve outer parameter links in loop invariants

src/analyze/basic_block.rs

@@ -1305,6 +1305,19 @@ 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/refine/env.rs

@@ -922,6 +922,10 @@ 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),

tests/ui/fail/loop_invariant_outer_param.rs

@@ -0,0 +1,24 @@
+//@error-in-other-file: Unsat
+//@compile-flags: -C debug-assertions=off
+
+#[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| true);
+        v = a;
+    }
+
+    v
+}
+
+fn main() {}

tests/ui/pass/loop_invariant_outer_param.rs

@@ -0,0 +1,24 @@
+//@check-pass
+//@compile-flags: -C debug-assertions=off
+
+#[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() {}