Equate a parameter with its argument only when the body never assigns it

src/analyze/basic_block.rs

@@ -3,6 +3,7 @@ use std::collections::HashMap;
 
 use num_bigint::BigInt;
 use rustc_hir::def::DefKind;
+use rustc_index::bit_set::DenseBitSet;
 use rustc_index::IndexVec;
 use rustc_middle::mir::{
     self, BasicBlock, Body, Local, Operand, Rvalue, StatementKind, TerminatorKind,
@@ -51,6 +52,48 @@ pub fn needs_own_precondition(body: &Body<'_>, bb: BasicBlock) -> bool {
     pred_term.successors().filter(|s| *s == bb).count() > 1
 }
 
+/// For each block, the parameters that may have been assigned since the function's entry when
+/// control reaches the block.
+pub fn params_assigned_on_entry(body: &Body<'_>) -> IndexVec<BasicBlock, DenseBitSet<Local>> {
+    struct Visitor {
+        arg_count: usize,
+        sites: Vec<(Local, BasicBlock)>,
+    }
+    impl<'tcx> mir::visit::Visitor<'tcx> for Visitor {
+        fn visit_local(
+            &mut self,
+            local: Local,
+            ctxt: mir::visit::PlaceContext,
+            location: mir::Location,
+        ) {
+            let is_param = (1..=self.arg_count).contains(&local.index());
+            if is_param && ctxt.is_place_assignment() {
+                self.sites.push((local, location.block));
+            }
+        }
+    }
+    let mut visitor = Visitor {
+        arg_count: body.arg_count,
+        sites: Vec::new(),
+    };
+    use mir::visit::Visitor as _;
+    visitor.visit_body(body);
+
+    let mut assigned = IndexVec::from_elem_n(
+        DenseBitSet::new_empty(body.local_decls.len()),
+        body.basic_blocks.len(),
+    );
+    for (local, site) in visitor.sites {
+        let mut worklist: Vec<_> = body.basic_blocks[site].terminator().successors().collect();
+        while let Some(bb) = worklist.pop() {
+            if assigned[bb].insert(local) {
+                worklist.extend(body.basic_blocks[bb].terminator().successors());
+            }
+        }
+    }
+    assigned
+}
+
 /// Whether every value of the integer type `inner` is also a value of the integer type `outer`.
 fn int_ty_includes<'tcx>(
     tcx: TyCtxt<'tcx>,
@@ -190,6 +233,7 @@ pub struct Analyzer<'tcx, 'ctx> {
 
     local_def_id: LocalDefId,
     drop_points: DropPoints,
+    assigned_params: DenseBitSet<Local>,
     basic_block: BasicBlock,
     body: Cow<'tcx, Body<'tcx>>,
 
@@ -1678,9 +1722,14 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
                     // 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.
+                    // equate them when the local is non-flow and not assigned on any path to this
+                    // block, since it then still holds the argument. A non-flow local can be
+                    // assigned again after being moved out.
                     let local = analyze::local_of_function_param(outer_idx);
-                    if !param_ty.to_sort().is_singleton() && self.env.is_non_flow_local(local) {
+                    if !param_ty.to_sort().is_singleton()
+                        && self.env.is_non_flow_local(local)
+                        && !self.assigned_params.contains(local)
+                    {
                         assumption.body.push_conj(
                             chc::Term::var(PlaceTypeVar::Var(var))
                                 .equal_to(self.env.local_type(local).term),
@@ -1718,6 +1767,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         let tcx = ctx.tcx;
         let drop_points = DropPoints::default();
         let body = Cow::Borrowed(tcx.optimized_mir(local_def_id.to_def_id()));
+        let assigned_params = DenseBitSet::new_empty(body.local_decls.len());
         let env = ctx.new_env();
         let local_decls = body.local_decls.clone();
         let prophecy_vars = Default::default();
@@ -1727,6 +1777,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
             tcx,
             local_def_id,
             drop_points,
+            assigned_params,
             basic_block,
             body,
             type_builder,
@@ -1741,6 +1792,13 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         self
     }
 
+    /// Sets the parameters that may have been assigned since the function's entry when control
+    /// reaches this block.
+    pub fn assigned_params(&mut self, assigned_params: DenseBitSet<Local>) -> &mut Self {
+        self.assigned_params = assigned_params;
+        self
+    }
+
     pub fn body(&mut self, body: Body<'tcx>) -> &mut Self {
         self.local_decls = body.local_decls.clone();
         self.body = Cow::Owned(body);

src/analyze/local_def.rs

@@ -990,6 +990,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
 
     fn analyze_basic_blocks(&mut self, expected_fn_ty: &rty::RefinedType) {
         let expected_fn_ty = expected_fn_ty.ty.as_function().unwrap();
+        let assigned_params = analyze::basic_block::params_assigned_on_entry(&self.body);
         // Reverse postorder guarantees each block that inherits its precondition
         // is visited after the predecessor that lazily materialized its type.
         for (bb, data) in mir::traversal::reverse_postorder(&self.body) {
@@ -1005,6 +1006,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
                 .basic_block_analyzer(self.local_def_id, bb)
                 .body(self.body.clone())
                 .drop_points(drop_points)
+                .assigned_params(assigned_params[bb].clone())
                 .run(&rty, expected_fn_ty);
         }
     }

tests/ui/fail/fn_param_reinit_after_loop.rs

@@ -0,0 +1,37 @@
+//@error-in-other-file: Unsat
+//@compile-flags: -C debug-assertions=off
+
+struct S {
+    x: i64,
+}
+
+impl thrust_models::Model for S {
+    type Ty = Self;
+}
+
+#[thrust_macros::requires(true)]
+#[thrust_macros::ensures(true)]
+#[thrust::trusted]
+fn rand() -> bool {
+    unimplemented!()
+}
+
+fn id(x: i64) -> i64 {
+    x
+}
+
+#[thrust_macros::ensures(result == a.x + 1)]
+#[thrust_macros::context]
+fn keep_argument(mut a: S) -> i64 {
+    let mut v = a.x;
+    while rand() {
+        thrust_macros::invariant!(|a: S, v: i64| v == a.x);
+        v = a.x;
+    }
+    let t = a;
+    a = S { x: id(t.x) };
+    let _ = a.x;
+    v
+}
+
+fn main() {}

tests/ui/fail/fn_param_reinit_after_move.rs

@@ -0,0 +1,25 @@
+//@error-in-other-file: Unsat
+//@compile-flags: -C debug-assertions=off
+
+struct S {
+    a: i64,
+}
+
+impl thrust_models::Model for S {
+    type Ty = Self;
+}
+
+fn id(x: i64) -> i64 {
+    x
+}
+
+fn f(mut s: S) -> i64 {
+    let t = s;
+    s = S { a: id(t.a + 1) };
+    if id(0) == 0 {}
+    s.a
+}
+
+fn main() {
+    assert!(f(S { a: 0 }) == 0);
+}

tests/ui/pass/fn_param_reinit_after_loop.rs

@@ -0,0 +1,37 @@
+//@check-pass
+//@compile-flags: -C debug-assertions=off
+
+struct S {
+    x: i64,
+}
+
+impl thrust_models::Model for S {
+    type Ty = Self;
+}
+
+#[thrust_macros::requires(true)]
+#[thrust_macros::ensures(true)]
+#[thrust::trusted]
+fn rand() -> bool {
+    unimplemented!()
+}
+
+fn id(x: i64) -> i64 {
+    x
+}
+
+#[thrust_macros::ensures(result == a.x)]
+#[thrust_macros::context]
+fn keep_argument(mut a: S) -> i64 {
+    let mut v = a.x;
+    while rand() {
+        thrust_macros::invariant!(|a: S, v: i64| v == a.x);
+        v = a.x;
+    }
+    let t = a;
+    a = S { x: id(t.x) };
+    let _ = a.x;
+    v
+}
+
+fn main() {}

tests/ui/pass/fn_param_reinit_after_move.rs

@@ -0,0 +1,25 @@
+//@check-pass
+//@compile-flags: -C debug-assertions=off
+
+struct S {
+    a: i64,
+}
+
+impl thrust_models::Model for S {
+    type Ty = Self;
+}
+
+fn id(x: i64) -> i64 {
+    x
+}
+
+fn f(mut s: S) -> i64 {
+    let t = s;
+    s = S { a: id(t.a + 1) };
+    if id(0) == 0 {}
+    s.a
+}
+
+fn main() {
+    assert!(f(S { a: 0 }) == 1);
+}