1 files changed, 48 insertions, 3 deletions

diff --git a/compiler/rustc_trait_selection/src/traits/engine.rs b/compiler/rustc_trait_selection/src/traits/engine.rs
index 6c177f638..e0c8deec9 100644
--- a/compiler/rustc_trait_selection/src/traits/engine.rs
+++ b/compiler/rustc_trait_selection/src/traits/engine.rs
@@ -3,7 +3,8 @@ use std::cell::RefCell;
 use super::TraitEngine;
 use super::{ChalkFulfillmentContext, FulfillmentContext};
 use crate::infer::InferCtxtExt;
-use rustc_hir::def_id::DefId;
+use rustc_data_structures::fx::FxHashSet;
+use rustc_hir::def_id::{DefId, LocalDefId};
 use rustc_infer::infer::{InferCtxt, InferOk};
 use rustc_infer::traits::{
     FulfillmentError, Obligation, ObligationCause, PredicateObligation, TraitEngineExt as _,
@@ -12,9 +13,11 @@ use rustc_middle::ty::error::TypeError;
 use rustc_middle::ty::ToPredicate;
 use rustc_middle::ty::TypeFoldable;
 use rustc_middle::ty::{self, Ty, TyCtxt};
+use rustc_span::Span;
 
 pub trait TraitEngineExt<'tcx> {
     fn new(tcx: TyCtxt<'tcx>) -> Box<Self>;
+    fn new_in_snapshot(tcx: TyCtxt<'tcx>) -> Box<Self>;
 }
 
 impl<'tcx> TraitEngineExt<'tcx> for dyn TraitEngine<'tcx> {
@@ -25,20 +28,32 @@ impl<'tcx> TraitEngineExt<'tcx> for dyn TraitEngine<'tcx> {
             Box::new(FulfillmentContext::new())
         }
     }
+
+    fn new_in_snapshot(tcx: TyCtxt<'tcx>) -> Box<Self> {
+        if tcx.sess.opts.unstable_opts.chalk {
+            Box::new(ChalkFulfillmentContext::new())
+        } else {
+            Box::new(FulfillmentContext::new_in_snapshot())
+        }
+    }
 }
 
 /// Used if you want to have pleasant experience when dealing
 /// with obligations outside of hir or mir typeck.
 pub struct ObligationCtxt<'a, 'tcx> {
-    pub infcx: &'a InferCtxt<'a, 'tcx>,
+    pub infcx: &'a InferCtxt<'tcx>,
     engine: RefCell<Box<dyn TraitEngine<'tcx>>>,
 }
 
 impl<'a, 'tcx> ObligationCtxt<'a, 'tcx> {
-    pub fn new(infcx: &'a InferCtxt<'a, 'tcx>) -> Self {
+    pub fn new(infcx: &'a InferCtxt<'tcx>) -> Self {
         Self { infcx, engine: RefCell::new(<dyn TraitEngine<'_>>::new(infcx.tcx)) }
     }
 
+    pub fn new_in_snapshot(infcx: &'a InferCtxt<'tcx>) -> Self {
+        Self { infcx, engine: RefCell::new(<dyn TraitEngine<'_>>::new_in_snapshot(infcx.tcx)) }
+    }
+
     pub fn register_obligation(&self, obligation: PredicateObligation<'tcx>) {
         self.engine.borrow_mut().register_predicate_obligation(self.infcx, obligation);
     }
@@ -109,4 +124,34 @@ impl<'a, 'tcx> ObligationCtxt<'a, 'tcx> {
     pub fn select_all_or_error(&self) -> Vec<FulfillmentError<'tcx>> {
         self.engine.borrow_mut().select_all_or_error(self.infcx)
     }
+
+    pub fn assumed_wf_types(
+        &self,
+        param_env: ty::ParamEnv<'tcx>,
+        span: Span,
+        def_id: LocalDefId,
+    ) -> FxHashSet<Ty<'tcx>> {
+        let tcx = self.infcx.tcx;
+        let assumed_wf_types = tcx.assumed_wf_types(def_id);
+        let mut implied_bounds = FxHashSet::default();
+        let hir_id = tcx.hir().local_def_id_to_hir_id(def_id);
+        let cause = ObligationCause::misc(span, hir_id);
+        for ty in assumed_wf_types {
+            // FIXME(@lcnr): rustc currently does not check wf for types
+            // pre-normalization, meaning that implied bounds are sometimes
+            // incorrect. See #100910 for more details.
+            //
+            // Not adding the unnormalized types here mostly fixes that, except
+            // that there are projections which are still ambiguous in the item definition
+            // but do normalize successfully when using the item, see #98543.
+            //
+            // Anyways, I will hopefully soon change implied bounds to make all of this
+            // sound and then uncomment this line again.
+
+            // implied_bounds.insert(ty);
+            let normalized = self.normalize(cause.clone(), param_env, ty);
+            implied_bounds.insert(normalized);
+        }
+        implied_bounds
+    }
 }