Merge #1135

1135: Generics improvements r=Aurel300 a=Aurel300

- [x] Improved associated types resolution (thanks `@vl0w,` from #980). The hope is that this won't erase too many regions...
- [x] Normalise associated types when monomorphising MIR bodies for pure functions.
- [x] Cache encoded pure functions by `DefId` + (full) encoded signature, not just "type parameters".

Co-authored-by: Aurel Bílý <[email protected]>
Co-authored-by: Jonas <[email protected]>

Author: 3 people

prusti-interface/src/environment/mod.rs

@@ -9,7 +9,7 @@
 use prusti_rustc_interface::middle::mir;
 use prusti_rustc_interface::hir::hir_id::HirId;
 use prusti_rustc_interface::hir::def_id::{DefId, LocalDefId};
-use prusti_rustc_interface::middle::ty::{self, Binder, BoundConstness, ImplPolarity, TraitPredicate, TraitRef, TyCtxt};
+use prusti_rustc_interface::middle::ty::{self, Binder, BoundConstness, ImplPolarity, TraitPredicate, TraitRef, TyCtxt, TypeSuperFoldable};
 use prusti_rustc_interface::middle::ty::subst::{Subst, SubstsRef};
 use prusti_rustc_interface::trait_selection::infer::{InferCtxtExt, TyCtxtInferExt};
 use prusti_rustc_interface::trait_selection::traits::{ImplSource, Obligation, ObligationCause, SelectionContext};
@@ -55,6 +55,11 @@ struct CachedBody<'tcx> {
     monomorphised_bodies: HashMap<SubstsRef<'tcx>, Rc<mir::Body<'tcx>>>,
     /// Cached borrowck information.
     borrowck_facts: Rc<BorrowckFacts>,
+    /// Copies of the MIR body with the given substs applied, called from the
+    /// given caller. This also allows for associated types to be correctly
+    /// normalised.
+    /// TODO: merge more nicely with monomorphised_bodies?
+    monomorphised_bodies_with_caller: HashMap<(SubstsRef<'tcx>, LocalDefId), Rc<mir::Body<'tcx>>>,
 }
 
 struct CachedExternalBody<'tcx> {
@@ -303,12 +308,56 @@ impl<'tcx> Environment<'tcx> {
                     base_body: Rc::new(body),
                     monomorphised_bodies: HashMap::new(),
                     borrowck_facts: Rc::new(facts),
+                    monomorphised_bodies_with_caller: HashMap::new(),
                 }
             });
         body
             .monomorphised_bodies
             .entry(substs)
-            .or_insert_with(|| ty::EarlyBinder(body.base_body.clone()).subst(self.tcx, substs))
+            .or_insert_with(|| {
+                let param_env = self.tcx.param_env(def_id);
+                let substituted = ty::EarlyBinder(body.base_body.clone()).subst(self.tcx, substs);
+                self.resolve_assoc_types(substituted.clone(), param_env)
+            })
+            .clone()
+    }
+
+    pub fn local_mir_with_caller(
+        &self,
+        def_id: LocalDefId,
+        caller_def_id: LocalDefId,
+        substs: SubstsRef<'tcx>,
+    ) -> Rc<mir::Body<'tcx>> {
+        // TODO: duplication ...
+        let mut bodies = self.bodies.borrow_mut();
+        let body = bodies.entry(def_id)
+            .or_insert_with(|| {
+                // SAFETY: This is safe because we are feeding in the same `tcx`
+                // that was used to store the data.
+                let body_with_facts = unsafe {
+                    self::mir_storage::retrieve_mir_body(self.tcx, def_id)
+                };
+                let body = body_with_facts.body;
+                let facts = BorrowckFacts {
+                    input_facts: RefCell::new(Some(body_with_facts.input_facts)),
+                    output_facts: body_with_facts.output_facts,
+                    location_table: RefCell::new(Some(body_with_facts.location_table)),
+                };
+
+                CachedBody {
+                    base_body: Rc::new(body),
+                    monomorphised_bodies: HashMap::new(),
+                    borrowck_facts: Rc::new(facts),
+                    monomorphised_bodies_with_caller: HashMap::new(),
+                }
+            });
+        body
+            .monomorphised_bodies_with_caller
+            .entry((substs, caller_def_id))
+            .or_insert_with(|| {
+                let param_env = self.tcx.param_env(caller_def_id);
+                self.tcx.subst_and_normalize_erasing_regions(substs, param_env, body.base_body.clone())
+            })
             .clone()
     }
 
@@ -538,10 +587,35 @@ impl<'tcx> Environment<'tcx> {
         // Normalize the type to account for associated types
         let ty = self.resolve_assoc_types(ty, param_env);
         let ty = self.tcx.erase_late_bound_regions(ty);
+
+        // Associated type was not normalised but it might still have a
+        // `Copy` bound declared on it.
+        // TODO: implement this more generally in `type_implements_trait` and
+        //       `type_implements_trait_with_trait_substs`.
+        if let ty::TyKind::Projection(proj) = ty.kind() {
+            let item_bounds = self.tcx.bound_item_bounds(proj.item_def_id);
+            if item_bounds.0.iter().any(|predicate| {
+                if let ty::PredicateKind::Trait(ty::TraitPredicate {
+                    trait_ref: ty::TraitRef { def_id: trait_def_id, .. },
+                    polarity: ty::ImplPolarity::Positive,
+                    ..
+                }) = predicate.kind().skip_binder() {
+                    trait_def_id == self.tcx.lang_items()
+                        .copy_trait()
+                        .unwrap()
+                } else {
+                    false
+                }
+            }) {
+                return true;
+            }
+        }
+
         ty.is_copy_modulo_regions(self.tcx.at(prusti_rustc_interface::span::DUMMY_SP), param_env)
     }
 
     /// Checks whether the given type implements the trait with the given DefId.
+    /// The trait should have no type parameters.
     pub fn type_implements_trait(&self, ty: ty::Ty<'tcx>, trait_def_id: DefId, param_env: ty::ParamEnv<'tcx>) -> bool {
         self.type_implements_trait_with_trait_substs(ty, trait_def_id, ty::List::empty(), param_env)
     }
@@ -583,22 +657,51 @@ impl<'tcx> Environment<'tcx> {
 
     /// Normalizes associated types in foldable types,
     /// i.e. this resolves projection types ([ty::TyKind::Projection]s)
-    /// **Important:** Regions while be erased during this process
-    pub fn resolve_assoc_types<T: ty::TypeFoldable<'tcx> + std::fmt::Debug + Copy>(&self, normalizable: T, param_env: ty::ParamEnv<'tcx>) -> T {
-        let norm_res = self.tcx.try_normalize_erasing_regions(
-            param_env,
-            normalizable
-        );
+    pub fn resolve_assoc_types<T: ty::TypeFoldable<'tcx> + std::fmt::Debug>(&self, normalizable: T, param_env: ty::ParamEnv<'tcx>) -> T {
+        struct Normalizer<'a, 'tcx> {
+            tcx: &'a ty::TyCtxt<'tcx>,
+            param_env: ty::ParamEnv<'tcx>,
+        }
+        impl<'a, 'tcx> ty::fold::TypeFolder<'tcx> for Normalizer<'a, 'tcx> {
+            fn tcx(&self) -> ty::TyCtxt<'tcx> {
+                *self.tcx
+            }
 
-        match norm_res {
-            Ok(normalized) => {
-                debug!("Normalized {:?}: {:?}", normalizable, normalized);
-                normalized
-            },
-            Err(err) => {
-                debug!("Error while resolving associated types for {:?}: {:?}", normalizable, err);
-                normalizable
+            fn fold_mir_const(&mut self, c: mir::ConstantKind<'tcx>) -> mir::ConstantKind<'tcx> {
+                // rustc by default panics when we execute this TypeFolder on a mir::* type,
+                // but we want to resolve associated types when we retrieve a local mir::Body
+                c
+            }
+
+            fn fold_ty(&mut self, ty: ty::Ty<'tcx>) -> ty::Ty<'tcx> {
+                match ty.kind() {
+                    ty::TyKind::Projection(_) => {
+                        let normalized = self.tcx.infer_ctxt().enter(|infcx| {
+                            use prusti_rustc_interface::trait_selection::traits::{fully_normalize, FulfillmentContext};
+
+                            let normalization_result = fully_normalize(
+                                &infcx,
+                                FulfillmentContext::new(),
+                                ObligationCause::dummy(),
+                                self.param_env,
+                                ty
+                            );
+
+                            match normalization_result {
+                                Ok(res) => res,
+                                Err(errors) => {
+                                    debug!("Error while resolving associated types: {:?}", errors);
+                                    ty
+                                }
+                            }
+                        });
+                        normalized.super_fold_with(self)
+                    }
+                    _ => ty.super_fold_with(self)
+                }
             }
         }
+
+        normalizable.fold_with(&mut Normalizer {tcx: &self.tcx, param_env})
     }
 }

prusti-tests/tests/verify_overflow/pass/generic/associated-copy.rs

@@ -0,0 +1,15 @@
+use prusti_contracts::*;
+
+trait Trait {
+    type Assoc: Copy;
+
+    #[pure] fn output_copy(&self) -> Self::Assoc;
+    #[pure] fn input_copy(&self, x: Self::Assoc) -> bool;
+}
+
+#[requires(x.output_copy() === y)]
+#[requires(x.input_copy(z))]
+fn test<U: Copy, T: Trait<Assoc = U>>(x: &mut T, y: U, z: U) {}
+
+#[trusted]
+fn main() {}