Carry what a type states across a basic block boundary

src/analyze.rs

@@ -585,6 +585,28 @@ impl<'tcx> Analyzer<'tcx> {
         self.basic_blocks.entry(def_id).or_default().insert(bb, def);
     }
 
+    /// Installs the types of a basic block's parameters.
+    ///
+    /// A block whose parameters are typed from MIR types alone carries an unrefined
+    /// specification for every function type they contain. This overwrites those
+    /// parameters with types whose specifications were recovered from elsewhere.
+    pub fn register_basic_block_param_tys(
+        &mut self,
+        def_id: LocalDefId,
+        bb: BasicBlock,
+        tys: impl IntoIterator<Item = (rty::FunctionParamIdx, rty::Type<rty::FunctionParamIdx>)>,
+    ) {
+        let bb_def = self
+            .basic_blocks
+            .get_mut(&def_id)
+            .unwrap()
+            .get_mut(&bb)
+            .unwrap();
+        for (idx, ty) in tys {
+            bb_def.ty.set_param_ty(idx, ty);
+        }
+    }
+
     pub fn register_basic_block_precondition(
         &mut self,
         def_id: LocalDefId,

src/analyze/basic_block.rs

@@ -816,11 +816,45 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         }
         capture.push_env_state(&self.env);
         let precondition = capture.finish(&self.env);
+        let param_tys = self.inherited_param_tys(bty);
 
+        self.ctx
+            .register_basic_block_param_tys(self.local_def_id, bb, param_tys);
         self.ctx
             .register_basic_block_precondition(self.local_def_id, bb, precondition);
     }
 
+    /// Takes the types the env holds for a goto target's params.
+    ///
+    /// The captured precondition carries what a parameter's *value* satisfies, and the
+    /// target's params are otherwise built from their MIR types alone. Everything a type
+    /// states by itself is therefore missing from the target: the refinements nested in
+    /// it, and the specification a function type spells out. Those are handed over here.
+    ///
+    /// A type in the env is closed — a refinement nested in one constrains the value at
+    /// its own position and names nothing from the env — so it transfers as it stands.
+    fn inherited_param_tys(
+        &self,
+        bty: &BasicBlockType,
+    ) -> Vec<(rty::FunctionParamIdx, rty::Type<rty::FunctionParamIdx>)> {
+        let mut tys = Vec::new();
+        for (param_idx, param_rty) in bty.as_ref().params.iter_enumerated() {
+            // An `OuterFnParam` copy of an argument names that argument's entry value, and
+            // takes its type from the outer function's signature rather than from the env.
+            let BasicBlockTypeParamKind::Local(local, _) = bty.param_kind(param_idx) else {
+                continue;
+            };
+            let ty = self.env.local_type(local).ty.assert_closed().vacuous();
+            assert_eq!(
+                ty.to_sort(),
+                param_rty.ty.to_sort(),
+                "env holds {local:?} at a different sort than the goto target's parameter"
+            );
+            tys.push((param_idx, ty));
+        }
+        tys
+    }
+
     fn with_assumptions<F, T>(&mut self, assumptions: Vec<impl Into<Assumption>>, callback: F) -> T
     where
         F: FnOnce(&mut Self) -> T,

src/refine/basic_block.rs

@@ -137,6 +137,15 @@ impl BasicBlockType {
         self.ty.clone()
     }
 
+    /// Replaces the type of the parameter at `idx`, keeping its refinement.
+    pub fn set_param_ty(
+        &mut self,
+        idx: rty::FunctionParamIdx,
+        ty: rty::Type<rty::FunctionParamIdx>,
+    ) {
+        self.ty.params[idx].ty = ty;
+    }
+
     pub fn set_precondition(&mut self, refinement: rty::Refinement<rty::FunctionParamIdx>) {
         let last_param_idx = self.ty.params.last_index().unwrap();
         self.ty.params.raw.last_mut().unwrap().refinement = refinement.map_var(|v| {

tests/ui/fail/fn_ptr_call_in_branch.rs

@@ -0,0 +1,18 @@
+//@error-in-other-file: Unsat
+//@compile-flags: -C debug-assertions=off
+
+fn add1(x: i64) -> i64 {
+    x + 1
+}
+
+// `add1(0)` is 1 rather than 0.
+#[thrust::callable]
+fn check(c: bool) {
+    let f: fn(i64) -> i64 = add1;
+    if c {
+        let a = f(0);
+        assert!(a == 0);
+    }
+}
+
+fn main() {}

tests/ui/fail/fn_ptr_call_twice.rs

@@ -0,0 +1,15 @@
+//@error-in-other-file: Unsat
+//@compile-flags: -C debug-assertions=off
+
+fn incr(m: &mut i64) {
+    *m += 1;
+}
+
+// `x` is incremented twice, so it is 2 rather than 1 here.
+fn main() {
+    let f: fn(&mut i64) = incr;
+    let mut x = 0;
+    f(&mut x);
+    f(&mut x);
+    assert!(x == 1);
+}

tests/ui/fail/fn_ptr_in_tuple_call_twice.rs

@@ -0,0 +1,14 @@
+//@error-in-other-file: Unsat
+//@compile-flags: -C debug-assertions=off
+
+fn add1(x: i64) -> i64 {
+    x + 1
+}
+
+// `add1` is applied twice, so the result is 2 rather than 1.
+fn main() {
+    let p: (fn(i64) -> i64,) = (add1,);
+    let a = (p.0)(0);
+    let b = (p.0)(a);
+    assert!(b == 1);
+}

tests/ui/fail/fn_ptr_param_call_twice.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() -> i64 { unimplemented!() }
+
+fn incr(m: &mut i64) {
+    *m += 1;
+}
+
+fn app(f: fn(&mut i64), mut x: i64) -> i64 {
+    f(&mut x);
+    f(&mut x);
+    x
+}
+
+// `x` is incremented twice, so it is `i + 2` rather than `i + 1` here.
+fn main() {
+    let i = rand();
+    let x = app(incr, i);
+    assert!(x == i + 1);
+}

tests/ui/pass/fn_ptr_call_in_branch.rs

@@ -0,0 +1,20 @@
+//@check-pass
+//@compile-flags: -C debug-assertions=off
+
+fn add1(x: i64) -> i64 {
+    x + 1
+}
+
+// The cast that produces `f` and the call of `f` sit in different basic blocks.
+// The callee's specification must survive that boundary; without it the call's
+// result is unconstrained.
+#[thrust::callable]
+fn check(c: bool) {
+    let f: fn(i64) -> i64 = add1;
+    if c {
+        let a = f(0);
+        assert!(a == 1);
+    }
+}
+
+fn main() {}

tests/ui/pass/fn_ptr_call_twice.rs

@@ -0,0 +1,17 @@
+//@check-pass
+//@compile-flags: -C debug-assertions=off
+
+fn incr(m: &mut i64) {
+    *m += 1;
+}
+
+// A call ends its basic block, so the second call sees `f` re-entering the block
+// it lives in. The callee's specification must survive that boundary; without it
+// the second call's effect on `x` is unconstrained.
+fn main() {
+    let f: fn(&mut i64) = incr;
+    let mut x = 0;
+    f(&mut x);
+    f(&mut x);
+    assert!(x == 2);
+}

tests/ui/pass/fn_ptr_in_tuple_call_twice.rs

@@ -0,0 +1,15 @@
+//@check-pass
+//@compile-flags: -C debug-assertions=off
+
+fn add1(x: i64) -> i64 {
+    x + 1
+}
+
+// The specification has to reach a function type nested inside another type, not
+// just one a local holds directly.
+fn main() {
+    let p: (fn(i64) -> i64,) = (add1,);
+    let a = (p.0)(0);
+    let b = (p.0)(a);
+    assert!(b == 2);
+}

tests/ui/pass/fn_ptr_param_call_twice.rs

@@ -0,0 +1,26 @@
+//@check-pass
+//@compile-flags: -C debug-assertions=off
+
+#[thrust_macros::requires(true)]
+#[thrust_macros::ensures(true)]
+#[thrust::trusted]
+fn rand() -> i64 { unimplemented!() }
+
+fn incr(m: &mut i64) {
+    *m += 1;
+}
+
+// A call ends its basic block, so the second call sees `f` re-entering the block
+// it lives in. The specification the caller supplied for `f` must survive that
+// boundary; without it the second call's effect on `x` is unconstrained.
+fn app(f: fn(&mut i64), mut x: i64) -> i64 {
+    f(&mut x);
+    f(&mut x);
+    x
+}
+
+fn main() {
+    let i = rand();
+    let x = app(incr, i);
+    assert!(x == i + 2);
+}