Canonicalize all nested goals together at once when making response

Author: compiler-errors

Diff for: compiler/rustc_middle/src/traits/solve/inspect.rs

@@ -116,7 +116,7 @@ impl Debug for Probe<'_> {
 pub enum ProbeStep<'tcx> {
     /// We added a goal to the `EvalCtxt` which will get proven
     /// the next time `EvalCtxt::try_evaluate_added_goals` is called.
-    AddGoal(GoalSource, CanonicalState<'tcx, Goal<'tcx, ty::Predicate<'tcx>>>),
+    AddGoal(GoalSource, Goal<'tcx, ty::Predicate<'tcx>>),
     /// The inside of a `EvalCtxt::try_evaluate_added_goals` call.
     EvaluateGoals(AddedGoalsEvaluation<'tcx>),
     /// A call to `probe` while proving the current goal. This is
@@ -128,7 +128,11 @@ pub enum ProbeStep<'tcx> {
     /// with the certainty of the `try_evaluate_added_goals` that is done within;
     /// if it's `Certainty::Yes`, then we can trust that the candidate is "finished"
     /// and we didn't force ambiguity for some reason.
-    MakeCanonicalResponse { shallow_certainty: Certainty },
+    MakeCanonicalResponse {
+        shallow_certainty: Certainty,
+        added_goals:
+            CanonicalState<'tcx, &'tcx ty::List<(GoalSource, Goal<'tcx, ty::Predicate<'tcx>>)>>,
+    },
 }
 
 /// What kind of probe we're in. In case the probe represents a candidate, or

Diff for: compiler/rustc_middle/src/traits/solve/inspect/format.rs

@@ -133,7 +133,7 @@ impl<'a, 'b> ProofTreeFormatter<'a, 'b> {
                     }
                     ProbeStep::EvaluateGoals(eval) => this.format_added_goals_evaluation(eval)?,
                     ProbeStep::NestedProbe(probe) => this.format_probe(probe)?,
-                    ProbeStep::MakeCanonicalResponse { shallow_certainty } => {
+                    ProbeStep::MakeCanonicalResponse { shallow_certainty, added_goals: _ } => {
                         writeln!(this.f, "EVALUATE GOALS AND MAKE RESPONSE: {shallow_certainty:?}")?
                     }
                 }

Diff for: compiler/rustc_middle/src/ty/context.rs

@@ -157,6 +157,8 @@ pub struct CtxtInterners<'tcx> {
     poly_existential_predicates: InternedSet<'tcx, List<PolyExistentialPredicate<'tcx>>>,
     predicate: InternedSet<'tcx, WithCachedTypeInfo<ty::Binder<'tcx, PredicateKind<'tcx>>>>,
     clauses: InternedSet<'tcx, ListWithCachedTypeInfo<Clause<'tcx>>>,
+    nested_goals:
+        InternedSet<'tcx, List<(solve::GoalSource, solve::Goal<'tcx, ty::Predicate<'tcx>>)>>,
     projs: InternedSet<'tcx, List<ProjectionKind>>,
     place_elems: InternedSet<'tcx, List<PlaceElem<'tcx>>>,
     const_: InternedSet<'tcx, WithCachedTypeInfo<ConstData<'tcx>>>,
@@ -186,6 +188,7 @@ impl<'tcx> CtxtInterners<'tcx> {
             canonical_var_infos: Default::default(),
             predicate: Default::default(),
             clauses: Default::default(),
+            nested_goals: Default::default(),
             projs: Default::default(),
             place_elems: Default::default(),
             const_: Default::default(),
@@ -1933,6 +1936,7 @@ slice_interners!(
     poly_existential_predicates: intern_poly_existential_predicates(PolyExistentialPredicate<'tcx>),
     projs: pub mk_projs(ProjectionKind),
     place_elems: pub mk_place_elems(PlaceElem<'tcx>),
+    nested_goals: pub mk_nested_goals((solve::GoalSource, solve::Goal<'tcx, ty::Predicate<'tcx>>)),
     bound_variable_kinds: pub mk_bound_variable_kinds(ty::BoundVariableKind),
     fields: pub mk_fields(FieldIdx),
     local_def_ids: intern_local_def_ids(LocalDefId),
@@ -2320,6 +2324,17 @@ impl<'tcx> TyCtxt<'tcx> {
         T::collect_and_apply(iter, |xs| self.mk_clauses(xs))
     }
 
+    pub fn mk_nested_goals_from_iter<I, T>(self, iter: I) -> T::Output
+    where
+        I: Iterator<Item = T>,
+        T: CollectAndApply<
+                (solve::GoalSource, solve::Goal<'tcx, ty::Predicate<'tcx>>),
+                &'tcx List<(solve::GoalSource, solve::Goal<'tcx, ty::Predicate<'tcx>>)>,
+            >,
+    {
+        T::collect_and_apply(iter, |xs| self.mk_nested_goals(xs))
+    }
+
     pub fn mk_type_list_from_iter<I, T>(self, iter: I) -> T::Output
     where
         I: Iterator<Item = T>,

Diff for: compiler/rustc_middle/src/ty/structural_impls.rs

@@ -4,6 +4,7 @@
 //! to help with the tedium.
 
 use crate::mir::interpret;
+use crate::traits::solve;
 use crate::ty::fold::{FallibleTypeFolder, TypeFoldable, TypeSuperFoldable};
 use crate::ty::print::{with_no_trimmed_paths, FmtPrinter, Printer};
 use crate::ty::visit::{TypeSuperVisitable, TypeVisitable, TypeVisitor};
@@ -560,6 +561,17 @@ impl<'tcx> TypeFoldable<TyCtxt<'tcx>> for &'tcx ty::List<ty::Const<'tcx>> {
     }
 }
 
+impl<'tcx> TypeFoldable<TyCtxt<'tcx>>
+    for &'tcx ty::List<(solve::GoalSource, solve::Goal<'tcx, ty::Predicate<'tcx>>)>
+{
+    fn try_fold_with<F: FallibleTypeFolder<TyCtxt<'tcx>>>(
+        self,
+        folder: &mut F,
+    ) -> Result<Self, F::Error> {
+        ty::util::fold_list(self, folder, |tcx, v| tcx.mk_nested_goals(v))
+    }
+}
+
 impl<'tcx> TypeFoldable<TyCtxt<'tcx>> for Pattern<'tcx> {
     fn try_fold_with<F: FallibleTypeFolder<TyCtxt<'tcx>>>(
         self,

Diff for: compiler/rustc_trait_selection/src/solve/eval_ctxt/canonical.rs

@@ -90,7 +90,7 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
         &mut self,
         certainty: Certainty,
     ) -> QueryResult<'tcx> {
-        self.inspect.make_canonical_response(certainty);
+        self.inspect.make_canonical_response(self.infcx, self.max_input_universe, certainty);
 
         let goals_certainty = self.try_evaluate_added_goals()?;
         assert_eq!(
@@ -444,5 +444,5 @@ pub(in crate::solve) fn instantiate_canonical_state<'tcx, T: TypeFoldable<TyCtxt
     let inspect::State { var_values, data } = state.instantiate(infcx.tcx, &instantiation);
 
     EvalCtxt::unify_query_var_values(infcx, param_env, orig_values, var_values);
-    data
+    infcx.resolve_vars_if_possible(data)
 }

Diff for: compiler/rustc_trait_selection/src/solve/eval_ctxt/mod.rs

@@ -462,13 +462,13 @@ impl<'a, 'tcx> EvalCtxt<'a, 'tcx> {
 
     #[instrument(level = "debug", skip(self))]
     pub(super) fn add_normalizes_to_goal(&mut self, goal: Goal<'tcx, ty::NormalizesTo<'tcx>>) {
-        self.inspect.add_normalizes_to_goal(self.infcx, self.max_input_universe, goal);
+        self.inspect.add_normalizes_to_goal(self.tcx(), goal);
         self.nested_goals.normalizes_to_goals.push(goal);
     }
 
     #[instrument(level = "debug", skip(self))]
     pub(super) fn add_goal(&mut self, source: GoalSource, goal: Goal<'tcx, ty::Predicate<'tcx>>) {
-        self.inspect.add_goal(self.infcx, self.max_input_universe, source, goal);
+        self.inspect.add_goal(source, goal);
         self.nested_goals.goals.push((source, goal));
     }
 

Diff for: compiler/rustc_trait_selection/src/solve/inspect/analyse.rs

@@ -93,7 +93,10 @@ impl<'tcx> NormalizesToTermHack<'tcx> {
 pub struct InspectCandidate<'a, 'tcx> {
     goal: &'a InspectGoal<'a, 'tcx>,
     kind: inspect::ProbeKind<'tcx>,
-    nested_goals: Vec<(GoalSource, inspect::CanonicalState<'tcx, Goal<'tcx, ty::Predicate<'tcx>>>)>,
+    nested_goals: inspect::CanonicalState<
+        'tcx,
+        &'tcx ty::List<(GoalSource, Goal<'tcx, ty::Predicate<'tcx>>)>,
+    >,
     final_state: inspect::CanonicalState<'tcx, ()>,
     result: QueryResult<'tcx>,
     shallow_certainty: Certainty,
@@ -141,23 +144,15 @@ impl<'a, 'tcx> InspectCandidate<'a, 'tcx> {
         let infcx = self.goal.infcx;
         let param_env = self.goal.goal.param_env;
         let mut orig_values = self.goal.orig_values.to_vec();
-        let instantiated_goals: Vec<_> = self
-            .nested_goals
-            .iter()
-            .map(|(source, goal)| {
-                (
-                    *source,
-                    canonical::instantiate_canonical_state(
-                        infcx,
-                        span,
-                        param_env,
-                        &mut orig_values,
-                        *goal,
-                    ),
-                )
-            })
-            .collect();
-
+        let instantiated_goals = canonical::instantiate_canonical_state(
+            infcx,
+            span,
+            param_env,
+            &mut orig_values,
+            self.nested_goals,
+        );
+        // FIXME: Should we truncate orig_values back down? Any new infer vars
+        // between `nested_goals` and `final_state` are definitely unrelated...
         let () = canonical::instantiate_canonical_state(
             infcx,
             span,
@@ -174,7 +169,7 @@ impl<'a, 'tcx> InspectCandidate<'a, 'tcx> {
         }
 
         instantiated_goals
-            .into_iter()
+            .iter()
             .map(|(source, goal)| match goal.predicate.kind().no_bound_vars() {
                 Some(ty::PredicateKind::NormalizesTo(ty::NormalizesTo { alias, term })) => {
                     let unconstrained_term = match term.unpack() {
@@ -240,27 +235,24 @@ impl<'a, 'tcx> InspectGoal<'a, 'tcx> {
     fn candidates_recur(
         &'a self,
         candidates: &mut Vec<InspectCandidate<'a, 'tcx>>,
-        nested_goals: &mut Vec<(
-            GoalSource,
-            inspect::CanonicalState<'tcx, Goal<'tcx, ty::Predicate<'tcx>>>,
-        )>,
         probe: &inspect::Probe<'tcx>,
     ) {
-        let mut shallow_certainty = None;
+        let mut shallow_certainty_and_goals = None;
         for step in &probe.steps {
-            match step {
-                &inspect::ProbeStep::AddGoal(source, goal) => nested_goals.push((source, goal)),
+            match *step {
                 inspect::ProbeStep::NestedProbe(ref probe) => {
                     // Nested probes have to prove goals added in their parent
                     // but do not leak them, so we truncate the added goals
                     // afterwards.
-                    let num_goals = nested_goals.len();
-                    self.candidates_recur(candidates, nested_goals, probe);
-                    nested_goals.truncate(num_goals);
+                    self.candidates_recur(candidates, probe);
                 }
-                inspect::ProbeStep::MakeCanonicalResponse { shallow_certainty: c } => {
-                    assert_eq!(shallow_certainty.replace(*c), None);
+                inspect::ProbeStep::MakeCanonicalResponse { shallow_certainty, added_goals } => {
+                    assert_eq!(
+                        shallow_certainty_and_goals.replace((shallow_certainty, added_goals)),
+                        None
+                    );
                 }
+                inspect::ProbeStep::AddGoal(..) => {}
                 inspect::ProbeStep::EvaluateGoals(_) => (),
             }
         }
@@ -278,11 +270,11 @@ impl<'a, 'tcx> InspectGoal<'a, 'tcx> {
             | inspect::ProbeKind::OpaqueTypeStorageLookup { result } => {
                 // We only add a candidate if `shallow_certainty` was set, which means
                 // that we ended up calling `evaluate_added_goals_and_make_canonical_response`.
-                if let Some(shallow_certainty) = shallow_certainty {
+                if let Some((shallow_certainty, nested_goals)) = shallow_certainty_and_goals {
                     candidates.push(InspectCandidate {
                         goal: self,
                         kind: probe.kind,
-                        nested_goals: nested_goals.clone(),
+                        nested_goals,
                         final_state: probe.final_state,
                         result,
                         shallow_certainty,
@@ -310,8 +302,7 @@ impl<'a, 'tcx> InspectGoal<'a, 'tcx> {
             }
         };
 
-        let mut nested_goals = vec![];
-        self.candidates_recur(&mut candidates, &mut nested_goals, &last_eval_step.evaluation);
+        self.candidates_recur(&mut candidates, &last_eval_step.evaluation);
 
         candidates
     }

Diff for: compiler/rustc_trait_selection/src/solve/inspect/build.rs

@@ -238,10 +238,16 @@ impl<'tcx> WipProbe<'tcx> {
 
 #[derive(Eq, PartialEq, Debug)]
 enum WipProbeStep<'tcx> {
-    AddGoal(GoalSource, inspect::CanonicalState<'tcx, Goal<'tcx, ty::Predicate<'tcx>>>),
+    AddGoal(GoalSource, Goal<'tcx, ty::Predicate<'tcx>>),
     EvaluateGoals(WipAddedGoalsEvaluation<'tcx>),
     NestedProbe(WipProbe<'tcx>),
-    MakeCanonicalResponse { shallow_certainty: Certainty },
+    MakeCanonicalResponse {
+        shallow_certainty: Certainty,
+        added_goals: inspect::CanonicalState<
+            'tcx,
+            &'tcx ty::List<(GoalSource, Goal<'tcx, ty::Predicate<'tcx>>)>,
+        >,
+    },
 }
 
 impl<'tcx> WipProbeStep<'tcx> {
@@ -250,8 +256,8 @@ impl<'tcx> WipProbeStep<'tcx> {
             WipProbeStep::AddGoal(source, goal) => inspect::ProbeStep::AddGoal(source, goal),
             WipProbeStep::EvaluateGoals(eval) => inspect::ProbeStep::EvaluateGoals(eval.finalize()),
             WipProbeStep::NestedProbe(probe) => inspect::ProbeStep::NestedProbe(probe.finalize()),
-            WipProbeStep::MakeCanonicalResponse { shallow_certainty } => {
-                inspect::ProbeStep::MakeCanonicalResponse { shallow_certainty }
+            WipProbeStep::MakeCanonicalResponse { shallow_certainty, added_goals } => {
+                inspect::ProbeStep::MakeCanonicalResponse { shallow_certainty, added_goals }
             }
         }
     }
@@ -500,47 +506,46 @@ impl<'tcx> ProofTreeBuilder<'tcx> {
 
     pub fn add_normalizes_to_goal(
         &mut self,
-        infcx: &InferCtxt<'tcx>,
-        max_input_universe: ty::UniverseIndex,
+        tcx: TyCtxt<'tcx>,
         goal: Goal<'tcx, ty::NormalizesTo<'tcx>>,
     ) {
-        self.add_goal(
-            infcx,
-            max_input_universe,
-            GoalSource::Misc,
-            goal.with(infcx.tcx, goal.predicate),
-        );
+        self.add_goal(GoalSource::Misc, goal.with(tcx, goal.predicate));
     }
 
-    pub fn add_goal(
-        &mut self,
-        infcx: &InferCtxt<'tcx>,
-        max_input_universe: ty::UniverseIndex,
-        source: GoalSource,
-        goal: Goal<'tcx, ty::Predicate<'tcx>>,
-    ) {
+    pub fn add_goal(&mut self, source: GoalSource, goal: Goal<'tcx, ty::Predicate<'tcx>>) {
         match self.as_mut() {
             None => {}
             Some(DebugSolver::GoalEvaluationStep(state)) => {
-                let goal = canonical::make_canonical_state(
-                    infcx,
-                    &state.var_values,
-                    max_input_universe,
-                    goal,
-                );
                 state.current_evaluation_scope().steps.push(WipProbeStep::AddGoal(source, goal))
             }
             _ => bug!(),
         }
     }
 
-    pub fn make_canonical_response(&mut self, shallow_certainty: Certainty) {
+    pub fn make_canonical_response(
+        &mut self,
+        infcx: &InferCtxt<'tcx>,
+        max_input_universe: ty::UniverseIndex,
+        shallow_certainty: Certainty,
+    ) {
         match self.as_mut() {
             Some(DebugSolver::GoalEvaluationStep(state)) => {
+                let added_goals = infcx.tcx.mk_nested_goals_from_iter(
+                    state.current_evaluation_scope().steps.iter().filter_map(|step| match *step {
+                        WipProbeStep::AddGoal(source, goal) => Some((source, goal)),
+                        _ => None,
+                    }),
+                );
+                let added_goals = canonical::make_canonical_state(
+                    infcx,
+                    &state.var_values,
+                    max_input_universe,
+                    added_goals,
+                );
                 state
                     .current_evaluation_scope()
                     .steps
-                    .push(WipProbeStep::MakeCanonicalResponse { shallow_certainty });
+                    .push(WipProbeStep::MakeCanonicalResponse { shallow_certainty, added_goals });
             }
             None => {}
             _ => {}