correctly handle normalization in implied bounds

aliemjay · aliemjay · commit 795401e9a695 · 2023-03-23T07:54:47.000+03:00
diff --git a/compiler/rustc_trait_selection/src/solve/mod.rs b/compiler/rustc_trait_selection/src/solve/mod.rs
@@ -140,19 +140,13 @@ impl<'a, 'tcx> EvalCtxt<'a, 'tcx> {
     #[instrument(level = "debug", skip(self))]
     fn compute_well_formed_goal(
         &mut self,
-        goal: Goal<'tcx, ty::GenericArg<'tcx>>,
+        Goal { predicate, param_env }: Goal<'tcx, ty::GenericArg<'tcx>>,
     ) -> QueryResult<'tcx> {
-        match crate::traits::wf::unnormalized_obligations(
-            self.infcx,
-            goal.param_env,
-            goal.predicate,
-        ) {
-            Some(obligations) => {
-                self.add_goals(obligations.into_iter().map(|o| o.into()));
-                self.evaluate_added_goals_and_make_canonical_response(Certainty::Yes)
-            }
-            None => self.evaluate_added_goals_and_make_canonical_response(Certainty::AMBIGUOUS),
-        }
+        debug_assert_eq!(predicate, self.infcx.resolve_vars_if_possible(predicate));
+        let obligations =
+            crate::traits::wf::unnormalized_obligations(self.tcx(), param_env, predicate);
+        self.add_goals(obligations.into_iter().map(|o| o.into()));
+        self.evaluate_added_goals_and_make_canonical_response(Certainty::Yes)
     }
 
     #[instrument(level = "debug", skip(self), ret)]
diff --git a/compiler/rustc_trait_selection/src/traits/wf.rs b/compiler/rustc_trait_selection/src/traits/wf.rs
@@ -77,21 +77,20 @@ pub fn obligations<'tcx>(
 
 /// Compute the predicates that are required for a type to be well-formed.
 ///
-/// This is only intended to be used in the new solver, since it does not
-/// take into account recursion depth or proper error-reporting spans.
+/// This is only intended to be used in implied bounds computation and in
+/// the new solver, since it does not take into account recursion depth or
+/// proper error-reporting spans.
 pub fn unnormalized_obligations<'tcx>(
-    infcx: &InferCtxt<'tcx>,
+    tcx: TyCtxt<'tcx>,
     param_env: ty::ParamEnv<'tcx>,
     arg: GenericArg<'tcx>,
-) -> Option<Vec<traits::PredicateObligation<'tcx>>> {
+) -> Vec<traits::PredicateObligation<'tcx>> {
     if let ty::GenericArgKind::Lifetime(..) = arg.unpack() {
-        return Some(vec![]);
+        return vec![];
     }
 
-    debug_assert_eq!(arg, infcx.resolve_vars_if_possible(arg));
-
     let mut wf = WfPredicates {
-        tcx: infcx.tcx,
+        tcx,
         param_env,
         body_id: CRATE_DEF_ID,
         span: DUMMY_SP,
@@ -100,7 +99,7 @@ pub fn unnormalized_obligations<'tcx>(
         item: None,
     };
     wf.compute(arg);
-    Some(wf.out)
+    wf.out
 }
 
 /// Returns the obligations that make this trait reference
diff --git a/compiler/rustc_traits/src/implied_outlives_bounds.rs b/compiler/rustc_traits/src/implied_outlives_bounds.rs
@@ -2,18 +2,16 @@
 //! Do not call this query directory. See
 //! [`rustc_trait_selection::traits::query::type_op::implied_outlives_bounds`].
 
-use rustc_infer::infer::canonical::{self, Canonical};
+use rustc_infer::infer::canonical;
 use rustc_infer::infer::outlives::components::{push_outlives_components, Component};
 use rustc_infer::infer::TyCtxtInferExt;
 use rustc_infer::traits::query::OutlivesBound;
 use rustc_middle::ty::query::Providers;
 use rustc_middle::ty::{self, Ty, TyCtxt, TypeVisitableExt};
-use rustc_span::def_id::CRATE_DEF_ID;
-use rustc_span::source_map::DUMMY_SP;
 use rustc_trait_selection::infer::InferCtxtBuilderExt;
 use rustc_trait_selection::traits::query::{CanonicalTyGoal, Fallible, NoSolution};
 use rustc_trait_selection::traits::wf;
-use rustc_trait_selection::traits::ObligationCtxt;
+use rustc_trait_selection::traits::ObligationCause;
 use smallvec::{smallvec, SmallVec};
 
 pub(crate) fn provide(p: &mut Providers) {
@@ -23,79 +21,52 @@ pub(crate) fn provide(p: &mut Providers) {
 fn implied_outlives_bounds<'tcx>(
     tcx: TyCtxt<'tcx>,
     goal: CanonicalTyGoal<'tcx>,
-) -> Result<
-    &'tcx Canonical<'tcx, canonical::QueryResponse<'tcx, Vec<OutlivesBound<'tcx>>>>,
-    NoSolution,
-> {
+) -> Fallible<canonical::CanonicalQueryResponse<'tcx, Vec<OutlivesBound<'tcx>>>> {
     tcx.infer_ctxt().enter_canonical_trait_query(&goal, |ocx, key| {
         let (param_env, ty) = key.into_parts();
-        compute_implied_outlives_bounds(ocx, param_env, ty)
+
+        compute_implied_outlives_bounds(tcx, param_env, ty, |ty| {
+            let ty = ocx.normalize(&ObligationCause::dummy(), param_env, ty);
+            if !ocx.select_all_or_error().is_empty() {
+                return Err(NoSolution);
+            }
+            let ty = ocx.infcx.resolve_vars_if_possible(ty);
+            assert!(!ty.has_non_region_infer());
+            Ok(ty)
+        })
     })
 }
 
+/// For the sake of completeness, we should be careful when dealing with inference artifacts:
+/// - This function shouldn't access an InferCtxt.
+/// - `ty` must be fully resolved.
+/// - `normalize_op` must return a fully resolved type.
 fn compute_implied_outlives_bounds<'tcx>(
-    ocx: &ObligationCtxt<'_, 'tcx>,
+    tcx: TyCtxt<'tcx>,
     param_env: ty::ParamEnv<'tcx>,
     ty: Ty<'tcx>,
+    normalize_op: impl Fn(Ty<'tcx>) -> Fallible<Ty<'tcx>>,
 ) -> Fallible<Vec<OutlivesBound<'tcx>>> {
-    let tcx = ocx.infcx.tcx;
-
     // Sometimes when we ask what it takes for T: WF, we get back that
     // U: WF is required; in that case, we push U onto this stack and
     // process it next. Because the resulting predicates aren't always
     // guaranteed to be a subset of the original type, so we need to store the
     // WF args we've computed in a set.
     let mut checked_wf_args = rustc_data_structures::fx::FxHashSet::default();
-    let mut wf_args = vec![ty.into()];
+    let mut wf_args = vec![ty.into(), normalize_op(ty)?.into()];
 
-    let mut outlives_bounds: Vec<ty::OutlivesPredicate<ty::GenericArg<'tcx>, ty::Region<'tcx>>> =
-        vec![];
+    let mut outlives_bounds: Vec<OutlivesBound<'tcx>> = vec![];
 
     while let Some(arg) = wf_args.pop() {
         if !checked_wf_args.insert(arg) {
             continue;
         }
 
-        // Compute the obligations for `arg` to be well-formed. If `arg` is
-        // an unresolved inference variable, just substituted an empty set
-        // -- because the return type here is going to be things we *add*
-        // to the environment, it's always ok for this set to be smaller
-        // than the ultimate set. (Note: normally there won't be
-        // unresolved inference variables here anyway, but there might be
-        // during typeck under some circumstances.)
-        //
-        // FIXME(@lcnr): It's not really "always fine", having fewer implied
-        // bounds can be backward incompatible, e.g. #101951 was caused by
-        // us not dealing with inference vars in `TypeOutlives` predicates.
-        let obligations = wf::obligations(ocx.infcx, param_env, CRATE_DEF_ID, 0, arg, DUMMY_SP)
-            .unwrap_or_default();
-
-        for obligation in obligations {
-            debug!(?obligation);
+        // From the full set of obligations, just filter down to the region relationships.
+        for obligation in wf::unnormalized_obligations(tcx, param_env, arg) {
             assert!(!obligation.has_escaping_bound_vars());
-
-            // While these predicates should all be implied by other parts of
-            // the program, they are still relevant as they may constrain
-            // inference variables, which is necessary to add the correct
-            // implied bounds in some cases, mostly when dealing with projections.
-            //
-            // Another important point here: we only register `Projection`
-            // predicates, since otherwise we might register outlives
-            // predicates containing inference variables, and we don't
-            // learn anything new from those.
-            if obligation.predicate.has_non_region_infer() {
-                match obligation.predicate.kind().skip_binder() {
-                    ty::PredicateKind::Clause(ty::Clause::Projection(..))
-                    | ty::PredicateKind::AliasEq(..) => {
-                        ocx.register_obligation(obligation.clone());
-                    }
-                    _ => {}
-                }
-            }
-
-            let pred = match obligation.predicate.kind().no_bound_vars() {
-                None => continue,
-                Some(pred) => pred,
+            let Some(pred) = obligation.predicate.kind().no_bound_vars() else {
+                continue;
             };
             match pred {
                 ty::PredicateKind::Clause(ty::Clause::Trait(..))
@@ -114,48 +85,27 @@ fn compute_implied_outlives_bounds<'tcx>(
                 | ty::PredicateKind::TypeWellFormedFromEnv(..) => {}
 
                 // We need to search through *all* WellFormed predicates
-                ty::PredicateKind::WellFormed(arg) => {
-                    wf_args.push(arg);
-                }
+                ty::PredicateKind::WellFormed(arg) => wf_args.push(arg),
 
                 // We need to register region relationships
-                ty::PredicateKind::Clause(ty::Clause::RegionOutlives(ty::OutlivesPredicate(
-                    r_a,
-                    r_b,
-                ))) => outlives_bounds.push(ty::OutlivesPredicate(r_a.into(), r_b)),
+                ty::PredicateKind::Clause(ty::Clause::RegionOutlives(
+                    ty::OutlivesPredicate(r_a, r_b),
+                )) => outlives_bounds.push(OutlivesBound::RegionSubRegion(r_b, r_a)),
 
                 ty::PredicateKind::Clause(ty::Clause::TypeOutlives(ty::OutlivesPredicate(
                     ty_a,
                     r_b,
-                ))) => outlives_bounds.push(ty::OutlivesPredicate(ty_a.into(), r_b)),
+                ))) => {
+                    let ty_a = normalize_op(ty_a)?;
+                    let mut components = smallvec![];
+                    push_outlives_components(tcx, ty_a, &mut components);
+                    outlives_bounds.extend(implied_bounds_from_components(r_b, components))
+                }
             }
         }
     }
 
-    // This call to `select_all_or_error` is necessary to constrain inference variables, which we
-    // use further down when computing the implied bounds.
-    match ocx.select_all_or_error().as_slice() {
-        [] => (),
-        _ => return Err(NoSolution),
-    }
-
-    // We lazily compute the outlives components as
-    // `select_all_or_error` constrains inference variables.
-    let implied_bounds = outlives_bounds
-        .into_iter()
-        .flat_map(|ty::OutlivesPredicate(a, r_b)| match a.unpack() {
-            ty::GenericArgKind::Lifetime(r_a) => vec![OutlivesBound::RegionSubRegion(r_b, r_a)],
-            ty::GenericArgKind::Type(ty_a) => {
-                let ty_a = ocx.infcx.resolve_vars_if_possible(ty_a);
-                let mut components = smallvec![];
-                push_outlives_components(tcx, ty_a, &mut components);
-                implied_bounds_from_components(r_b, components)
-            }
-            ty::GenericArgKind::Const(_) => unreachable!(),
-        })
-        .collect();
-
-    Ok(implied_bounds)
+    Ok(outlives_bounds)
 }
 
 /// When we have an implied bound that `T: 'a`, we can further break
@@ -165,28 +115,22 @@ fn compute_implied_outlives_bounds<'tcx>(
 fn implied_bounds_from_components<'tcx>(
     sub_region: ty::Region<'tcx>,
     sup_components: SmallVec<[Component<'tcx>; 4]>,
-) -> Vec<OutlivesBound<'tcx>> {
-    sup_components
-        .into_iter()
-        .filter_map(|component| {
-            match component {
-                Component::Region(r) => Some(OutlivesBound::RegionSubRegion(sub_region, r)),
-                Component::Param(p) => Some(OutlivesBound::RegionSubParam(sub_region, p)),
-                Component::Alias(p) => Some(OutlivesBound::RegionSubAlias(sub_region, p)),
-                Component::EscapingAlias(_) =>
-                // If the projection has escaping regions, don't
-                // try to infer any implied bounds even for its
-                // free components. This is conservative, because
-                // the caller will still have to prove that those
-                // free components outlive `sub_region`. But the
-                // idea is that the WAY that the caller proves
-                // that may change in the future and we want to
-                // give ourselves room to get smarter here.
-                {
-                    None
-                }
-                Component::UnresolvedInferenceVariable(..) => None,
-            }
-        })
-        .collect()
+) -> impl Iterator<Item = OutlivesBound<'tcx>> {
+    sup_components.into_iter().filter_map(move |component| {
+        match component {
+            Component::Region(r) => Some(OutlivesBound::RegionSubRegion(sub_region, r)),
+            Component::Param(p) => Some(OutlivesBound::RegionSubParam(sub_region, p)),
+            Component::Alias(p) => Some(OutlivesBound::RegionSubAlias(sub_region, p)),
+            // If the projection has escaping regions, don't
+            // try to infer any implied bounds even for its
+            // free components. This is conservative, because
+            // the caller will still have to prove that those
+            // free components outlive `sub_region`. But the
+            // idea is that the WAY that the caller proves
+            // that may change in the future and we want to
+            // give ourselves room to get smarter here.
+            Component::EscapingAlias(_) => None,
+            Component::UnresolvedInferenceVariable(..) => bug!("inference var in implied bounds"),
+        }
+    })
 }
