Take the precondition of the entry block from the signature

src/analyze/local_def.rs

@@ -9,7 +9,6 @@ use rustc_span::def_id::{DefId, LocalDefId};
 
 use crate::analyze;
 use crate::chc;
-use crate::pretty::PrettyDisplayExt as _;
 use crate::refine::{self, BasicBlockType, TypeBuilder};
 use crate::rty;
 
@@ -889,7 +888,73 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
             .into()
     }
 
-    fn refine_basic_blocks(&mut self) {
+    /// The precondition of the entry block, which is the precondition of the function itself.
+    ///
+    /// The entry block takes each parameter of the function twice: as the local that holds it and
+    /// as the value it has on entry ([`refine::BasicBlockTypeParamKind::OuterFnParam`]), which the
+    /// postcondition of the function is stated in terms of. Both denote the argument, so the
+    /// precondition equates them along with carrying the refinement of every parameter.
+    fn entry_precondition(
+        &self,
+        expected: &rty::FunctionType,
+        bty: &BasicBlockType,
+    ) -> rty::Refinement<rty::FunctionParamIdx> {
+        // The value of a parameter of the function among the parameters of the entry block. One of
+        // a singleton sort is denoted by its only value, as it is no variable of the constraints
+        // (see `refine::Env::var_type`), and so is the synthetic parameter that a function without
+        // one carries its precondition on, which has no local in the entry block.
+        let param_term = |idx: rty::FunctionParamIdx| {
+            let sort = expected.params[idx].ty.to_sort();
+            match bty.param_of_local(analyze::local_of_function_param(idx)) {
+                Some(param_idx) if !sort.is_singleton() => {
+                    chc::Term::var(rty::RefinedTypeVar::Free(param_idx))
+                }
+                _ => chc::Term::default_for(&sort),
+            }
+        };
+
+        let mut precondition = rty::Refinement::top();
+        for (idx, param) in expected.params.iter_enumerated() {
+            precondition.push_conj(param.refinement.clone().subst_var(|v| match v {
+                rty::RefinedTypeVar::Value => param_term(idx),
+                rty::RefinedTypeVar::Free(free_idx) => param_term(free_idx),
+                rty::RefinedTypeVar::Existential(ev) => {
+                    chc::Term::var(rty::RefinedTypeVar::Existential(ev))
+                }
+            }));
+
+            if let Some(outer_param_idx) = bty.param_of_outer_fn_param(idx) {
+                precondition.push_conj(
+                    chc::Term::var(rty::RefinedTypeVar::Free(outer_param_idx))
+                        .equal_to(param_term(idx))
+                        .into(),
+                );
+            }
+        }
+        precondition
+    }
+
+    /// The type of the entry block, which takes the arguments of the call under the precondition
+    /// of the function, leaving nothing about its state to be inferred.
+    fn entry_block_ty(
+        &self,
+        expected: &rty::RefinedType,
+        live_locals: Vec<(Local, TypeAndMut<'tcx>)>,
+        ret_ty: mir_ty::Ty<'tcx>,
+    ) -> BasicBlockType {
+        let mut expected_fn = expected.ty.as_function().cloned().unwrap();
+        self.elaborate_mut_params(&mut expected_fn);
+
+        let mut bty = self
+            .type_builder
+            .build_basic_block(&self.body, live_locals, ret_ty);
+        bty.install_signature_types(&expected_fn.params);
+        let precondition = self.entry_precondition(&expected_fn, &bty);
+        bty.set_precondition(precondition);
+        bty
+    }
+
+    fn refine_basic_blocks(&mut self, expected: &rty::RefinedType) {
         use rustc_mir_dataflow::Analysis as _;
         let loop_invariants = self.collect_loop_invariant_annotations();
         let mut results = rustc_mir_dataflow::impls::MaybeLiveLocals
@@ -952,6 +1017,10 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
                 bty.set_precondition(inv);
                 self.ctx
                     .register_basic_block_ty_with_precondition(self.local_def_id, bb, bty);
+            } else if bb == mir::START_BLOCK {
+                let bty = self.entry_block_ty(expected, live_locals, ret_ty);
+                self.ctx
+                    .register_basic_block_ty_with_precondition(self.local_def_id, bb, bty);
             } else if analyze::basic_block::needs_own_precondition(&self.body, bb) {
                 let bty = self
                     .type_builder
@@ -1034,79 +1103,6 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
             }
         });
     }
-
-    /// Drop excessive parameters from the BB-side entry function type that do not
-    /// correspond to any function argument. These are introduced by ZST locals whose
-    /// liveness analysis treats them as live without an explicit def.
-    fn drop_bb_zst_params(&self, bb_ty: &BasicBlockType) -> rty::FunctionType {
-        let mut fn_ty = bb_ty.to_function_ty();
-        let arg_locals: HashSet<_> = self.body.args_iter().collect();
-
-        for idx in bb_ty.local_params().rev() {
-            let local = bb_ty.local_of_param(idx).unwrap();
-            if !arg_locals.contains(&local) {
-                fn_ty.remove_param(idx);
-            }
-        }
-
-        // A function type must keep at least one parameter to host the precondition
-        // predicate. When the function has no real argument, both the expected type and
-        // the BB type carry a synthetic unit parameter (see
-        // `crate::refine::TypeBuilder::build_basic_block`). That synthetic has no
-        // backing local, so it survives the drop loop untouched. If instead the entry
-        // block exposed only ZST-local parameters (e.g. `RETURN_PLACE`), dropping them
-        // empties the type, and we re-introduce the synthetic unit parameter carrying
-        // the precondition refinement of the last dropped parameter.
-        if self.body.arg_count == 0 && fn_ty.params.is_empty() {
-            let refinement = bb_ty.as_ref().last_param().unwrap().refinement.clone();
-            fn_ty
-                .params
-                .push(rty::RefinedType::new(rty::Type::unit(), refinement));
-        }
-
-        fn_ty
-    }
-
-    /// Drop function parameters from `expected_ty` whose corresponding local is unused
-    /// (and thus not represented) in the BB-side entry function type.
-    fn drop_unused_expected_params(
-        &self,
-        expected_ty: &mut rty::FunctionType,
-        bb_ty: &BasicBlockType,
-    ) {
-        if self.body.arg_count == 0 {
-            return;
-        }
-        let arg_locals: HashSet<_> = self.body.args_iter().collect();
-        let present_arg_locals: HashSet<_> = bb_ty
-            .locals()
-            .filter(|local| arg_locals.contains(local))
-            .collect();
-        for idx in expected_ty.params.indices().rev() {
-            let arg_local = analyze::local_of_function_param(idx);
-            if !present_arg_locals.contains(&arg_local) {
-                expected_ty.remove_param(idx);
-            }
-        }
-    }
-
-    fn assert_entry(&mut self, expected: &rty::RefinedType) {
-        let mut entry_ty = self
-            .ctx
-            .basic_block_ty_with_precondition(self.local_def_id, mir::START_BLOCK)
-            .clone();
-        tracing::debug!(expected = %expected.display(), entry = %entry_ty.display(), "assert_entry before");
-        let mut expected = expected.ty.as_function().cloned().unwrap();
-        self.elaborate_mut_params(&mut expected);
-
-        entry_ty.truncate_outer_fn_params();
-        self.drop_unused_expected_params(&mut expected, &entry_ty);
-        let entry_ty = self.drop_bb_zst_params(&entry_ty);
-
-        tracing::debug!(expected = %expected.display(), entry = %entry_ty.display(), "assert_entry after");
-        let clauses = rty::relate_sub_param_types(&entry_ty.params, &expected.params);
-        self.ctx.extend_clauses(clauses);
-    }
 }
 
 impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
@@ -1145,8 +1141,8 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         self.unelaborate_derefs();
         analyze::reconstruct_slice_indexing::reconstruct(self.tcx, &mut self.body);
         self.reassign_local_mutabilities();
-        self.refine_basic_blocks();
+
+        self.refine_basic_blocks(expected);
         self.analyze_basic_blocks(expected);
-        self.assert_entry(expected);
     }
 }

src/refine/basic_block.rs

@@ -1,12 +1,11 @@
 //! The refinement type for a basic block.
 
-use std::collections::HashMap;
-
 use pretty::{termcolor, Pretty};
 use rustc_index::IndexVec;
-use rustc_middle::mir::Local;
+use rustc_middle::mir::{self, Local};
 use rustc_middle::ty as mir_ty;
 
+use crate::chc;
 use crate::rty;
 
 #[derive(Debug, Clone)]
@@ -105,17 +104,6 @@ impl BasicBlockType {
         }
     }
 
-    pub fn local_params(&self) -> impl DoubleEndedIterator<Item = rty::FunctionParamIdx> + '_ {
-        self.locals.indices()
-    }
-
-    pub fn locals(&self) -> impl Iterator<Item = Local> + '_ {
-        self.ty
-            .params
-            .iter_enumerated()
-            .filter_map(|(idx, _)| self.local_of_param(idx))
-    }
-
     pub fn param_of_local(&self, local: Local) -> Option<rty::FunctionParamIdx> {
         self.locals
             .iter_enumerated()
@@ -133,8 +121,48 @@ impl BasicBlockType {
         }
     }
 
-    pub fn to_function_ty(&self) -> rty::FunctionType {
-        self.ty.clone()
+    /// Replaces the type of every parameter that holds a parameter of the function with the type
+    /// the signature of the function gives it.
+    ///
+    /// A signature can refine a type where the MIR type it is built from has nothing to say, as in
+    /// `Vec<{ v: i32 | v > 0 }>` or in the pre- and postcondition of a function-typed parameter.
+    /// The entry block is entered with the arguments of the call, so those are the types it takes.
+    pub fn install_signature_types(
+        &mut self,
+        params: &IndexVec<rty::FunctionParamIdx, rty::RefinedType<rty::FunctionParamIdx>>,
+    ) {
+        let param_of_fn_param = |idx| {
+            self.param_of_local(crate::analyze::local_of_function_param(idx))
+                .expect("the entry block takes every parameter of the function")
+        };
+        let signature_types: Vec<_> = self
+            .ty
+            .params
+            .indices()
+            .filter_map(|idx| {
+                let fn_param_idx = self.fn_param_of_param(idx)?;
+                let ty = params[fn_param_idx]
+                    .ty
+                    .clone()
+                    .subst_var(|idx| chc::Term::var(param_of_fn_param(idx)));
+                Some((idx, ty))
+            })
+            .collect();
+        for (idx, ty) in signature_types {
+            self.ty.params[idx].ty = ty;
+        }
+    }
+
+    /// The parameter of the function held by the parameter `idx`, if it holds one.
+    fn fn_param_of_param(&self, idx: rty::FunctionParamIdx) -> Option<rty::FunctionParamIdx> {
+        match self.param_kind(idx) {
+            BasicBlockTypeParamKind::Local(local, _) if local != mir::RETURN_PLACE => {
+                let fn_param_idx = crate::analyze::function_param_of_local(local);
+                (fn_param_idx.index() < self.outer_fn_param_count).then_some(fn_param_idx)
+            }
+            BasicBlockTypeParamKind::OuterFnParam(fn_param_idx) => Some(fn_param_idx),
+            _ => None,
+        }
     }
 
     pub fn set_precondition(&mut self, refinement: rty::Refinement<rty::FunctionParamIdx>) {
@@ -147,59 +175,4 @@ impl BasicBlockType {
             }
         });
     }
-
-    /// Inner function type of BasicBlockType contains extra parameters that carry original
-    /// function parameter values. `truncate_outer_fn_params` removes these extra parameters
-    /// to subtype output of [`BasicBlockType::to_function_ty`] against the function type.
-    ///
-    /// before: (_1: int, _2: int, int, { int |  p4 ν $0 $1 $2 }) → { int | p5 ν $0 $1 $2 $3 }
-    /// after:  (_1: int, _2: { int | p4 v $0 $1 $0 }) → { int | p5 ν $0 $1 _1 _2 }
-    ///
-    /// FIXME: this should be (&self) -> FunctionType
-    pub fn truncate_outer_fn_params(&mut self) {
-        let last_param_idx = self.ty.params.last_index().unwrap();
-        let last_param_ty = self.ty.params.raw.last().unwrap();
-
-        let mut mapping = HashMap::new();
-        for (idx, param_ty) in self.ty.params.iter_enumerated() {
-            let mapped_idx = if let Some(outer_idx) = self.param_kind(idx).outer_fn_param_idx() {
-                let corresponding_local = crate::analyze::local_of_function_param(outer_idx);
-                self.param_of_local(corresponding_local).unwrap()
-            } else {
-                idx
-            };
-            mapping.insert(idx, mapped_idx);
-
-            // to be sure
-            if idx != last_param_idx {
-                assert!(param_ty.refinement.is_top());
-            }
-        }
-
-        let last_param_refinement = last_param_ty.refinement.clone().map_var(|v| {
-            let idx = match v {
-                rty::RefinedTypeVar::Free(idx) => idx,
-                rty::RefinedTypeVar::Value => last_param_idx,
-                v => return v,
-            };
-            let mapped_idx = mapping[&idx];
-            if Some(mapped_idx) == self.locals.last_index() {
-                rty::RefinedTypeVar::Value
-            } else {
-                rty::RefinedTypeVar::Free(mapped_idx)
-            }
-        });
-
-        if !self.locals.is_empty() {
-            self.ty.params.truncate(self.locals.len());
-        }
-
-        self.ty.params.raw.last_mut().unwrap().refinement = last_param_refinement;
-        self.ty.ret.refinement = self
-            .ty
-            .ret
-            .refinement
-            .clone()
-            .map_free_var(|idx| mapping[&idx]);
-    }
 }

src/refine/template.rs

@@ -263,7 +263,7 @@ impl<'tcx> TypeBuilder<'tcx> {
     }
 
     pub fn build_basic_block<I>(
-        &mut self,
+        &self,
         body: &rustc_middle::mir::Body<'tcx>,
         live_locals: I,
         ret_ty: mir_ty::Ty<'tcx>,

src/rty.rs

@@ -52,7 +52,7 @@ mod clause_builder;
 pub use clause_builder::ClauseBuilderExt;
 
 mod subtyping;
-pub use subtyping::{relate_sub_param_types, ClauseScope, Subtyping};
+pub use subtyping::{ClauseScope, Subtyping};
 
 mod params;
 pub use params::{RefinedTypeArgs, TypeParamIdx, TypeParamSubst};

src/rty/subtyping.rs

@@ -1,11 +1,9 @@
 //! Translation of subtyping relations into CHC constraints.
 
-use rustc_index::IndexVec;
-
 use crate::chc;
 use crate::pretty::PrettyDisplayExt;
 
-use super::{ClauseBuilderExt as _, FunctionParamIdx, PointerKind, RefKind, RefinedType, Type};
+use super::{ClauseBuilderExt as _, PointerKind, RefKind, RefinedType, Type};
 
 /// A scope for building clauses.
 ///
@@ -170,28 +168,3 @@ where
         clauses
     }
 }
-
-#[must_use]
-pub fn relate_sub_param_types(
-    got: &IndexVec<FunctionParamIdx, RefinedType<FunctionParamIdx>>,
-    expected: &IndexVec<FunctionParamIdx, RefinedType<FunctionParamIdx>>,
-) -> Vec<chc::Clause> {
-    assert_eq!(got.len(), expected.len());
-
-    let mut clauses = Vec::new();
-    let mut builder = chc::ClauseBuilder::default();
-
-    for (param_idx, param_rty) in got.iter_enumerated() {
-        let param_sort = param_rty.ty.to_sort();
-        if !param_sort.is_singleton() {
-            builder.add_mapped_var(param_idx, param_sort);
-        }
-    }
-
-    for (got_ty, expected_ty) in got.iter().zip(expected.iter()) {
-        let cs = builder.relate_sub_refined_type(expected_ty, got_ty);
-        clauses.extend(cs);
-    }
-
-    clauses
-}