try_normalize_ty end with rigid alias on failure

Author: lcnr

compiler/rustc_middle/src/traits/solve/inspect.rs

@@ -122,6 +122,8 @@ pub enum ProbeStep<'tcx> {
     /// used whenever there are multiple candidates to prove the
     /// current goalby .
     NestedProbe(Probe<'tcx>),
+    CommitIfOkStart,
+    CommitIfOkSuccess,
 }
 
 /// What kind of probe we're in. In case the probe represents a candidate, or
@@ -142,6 +144,9 @@ pub enum ProbeKind<'tcx> {
     /// Used in the probe that wraps normalizing the non-self type for the unsize
     /// trait, which is also structurally matched on.
     UnsizeAssembly,
+    /// A call to `EvalCtxt::commit_if_ok` which failed, causing the work
+    /// to be discarded.
+    CommitIfOk,
     /// During upcasting from some source object to target object type, used to
     /// do a probe to find out what projection type(s) may be used to prove that
     /// the source type upholds all of the target type's object bounds.

compiler/rustc_middle/src/traits/solve/inspect/format.rs

@@ -109,6 +109,9 @@ impl<'a, 'b> ProofTreeFormatter<'a, 'b> {
             ProbeKind::UpcastProjectionCompatibility => {
                 writeln!(self.f, "PROBING FOR PROJECTION COMPATIBILITY FOR UPCASTING:")
             }
+            ProbeKind::CommitIfOk => {
+                writeln!(self.f, "COMMIT_IF_OK:")
+            }
             ProbeKind::MiscCandidate { name, result } => {
                 writeln!(self.f, "CANDIDATE {name}: {result:?}")
             }
@@ -123,6 +126,8 @@ impl<'a, 'b> ProofTreeFormatter<'a, 'b> {
                     ProbeStep::AddGoal(goal) => writeln!(this.f, "ADDED GOAL: {goal:?}")?,
                     ProbeStep::EvaluateGoals(eval) => this.format_added_goals_evaluation(eval)?,
                     ProbeStep::NestedProbe(probe) => this.format_probe(probe)?,
+                    ProbeStep::CommitIfOkStart => writeln!(this.f, "COMMIT_IF_OK START")?,
+                    ProbeStep::CommitIfOkSuccess => writeln!(this.f, "COMMIT_IF_OK SUCCESS")?,
                 }
             }
             Ok(())

compiler/rustc_trait_selection/src/solve/assembly/mod.rs

@@ -864,23 +864,18 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
 
         let result = self.probe_misc_candidate("coherence unknowable").enter(|ecx| {
             let trait_ref = goal.predicate.trait_ref(tcx);
-
             #[derive(Debug)]
-            enum FailureKind {
-                Overflow,
-                NoSolution(NoSolution),
-            }
+            struct Overflow;
             let lazily_normalize_ty = |ty| match ecx.try_normalize_ty(goal.param_env, ty) {
-                Ok(Some(ty)) => Ok(ty),
-                Ok(None) => Err(FailureKind::Overflow),
-                Err(e) => Err(FailureKind::NoSolution(e)),
+                Some(ty) => Ok(ty),
+                None => Err(Overflow),
             };
 
             match coherence::trait_ref_is_knowable(tcx, trait_ref, lazily_normalize_ty) {
-                Err(FailureKind::Overflow) => {
+                Err(Overflow) => {
                     ecx.evaluate_added_goals_and_make_canonical_response(Certainty::OVERFLOW)
                 }
-                Err(FailureKind::NoSolution(NoSolution)) | Ok(Ok(())) => Err(NoSolution),
+                Ok(Ok(())) => Err(NoSolution),
                 Ok(Err(_)) => {
                     ecx.evaluate_added_goals_and_make_canonical_response(Certainty::AMBIGUOUS)
                 }

compiler/rustc_trait_selection/src/solve/eval_ctxt/commit_if_ok.rs

@@ -0,0 +1,45 @@
+use super::EvalCtxt;
+use crate::solve::inspect;
+use rustc_middle::traits::query::NoSolution;
+
+impl<'a, 'tcx> EvalCtxt<'a, 'tcx> {
+    pub(in crate::solve) fn commit_if_ok<T>(
+        &mut self,
+        f: impl FnOnce(&mut EvalCtxt<'_, 'tcx>) -> Result<T, NoSolution>,
+    ) -> Result<T, NoSolution> {
+        let mut nested_ecx = EvalCtxt {
+            infcx: self.infcx,
+            variables: self.variables,
+            var_values: self.var_values,
+            predefined_opaques_in_body: self.predefined_opaques_in_body,
+            max_input_universe: self.max_input_universe,
+            search_graph: self.search_graph,
+            nested_goals: self.nested_goals.clone(),
+            tainted: self.tainted,
+            inspect: self.inspect.new_probe(),
+        };
+
+        let result = nested_ecx.infcx.commit_if_ok(|_| f(&mut nested_ecx));
+        if result.is_ok() {
+            let EvalCtxt {
+                infcx: _,
+                variables: _,
+                var_values: _,
+                predefined_opaques_in_body: _,
+                max_input_universe: _,
+                search_graph: _,
+                nested_goals,
+                tainted,
+                inspect,
+            } = nested_ecx;
+            self.nested_goals = nested_goals;
+            self.tainted = tainted;
+            self.inspect.integrate_snapshot(inspect);
+        } else {
+            nested_ecx.inspect.probe_kind(inspect::ProbeKind::CommitIfOk);
+            self.inspect.finish_probe(nested_ecx.inspect);
+        }
+
+        result
+    }
+}

compiler/rustc_trait_selection/src/solve/eval_ctxt/mod.rs

@@ -34,6 +34,7 @@ use super::{search_graph::SearchGraph, Goal};
 pub use select::InferCtxtSelectExt;
 
 mod canonical;
+mod commit_if_ok;
 mod probe;
 mod select;
 

compiler/rustc_trait_selection/src/solve/inspect/analyse.rs

@@ -120,7 +120,6 @@ impl<'a, 'tcx> InspectGoal<'a, 'tcx> {
         for step in &probe.steps {
             match step {
                 &inspect::ProbeStep::AddGoal(goal) => nested_goals.push(goal),
-                inspect::ProbeStep::EvaluateGoals(_) => (),
                 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
@@ -129,13 +128,17 @@ impl<'a, 'tcx> InspectGoal<'a, 'tcx> {
                     self.candidates_recur(candidates, nested_goals, probe);
                     nested_goals.truncate(num_goals);
                 }
+                inspect::ProbeStep::EvaluateGoals(_)
+                | inspect::ProbeStep::CommitIfOkStart
+                | inspect::ProbeStep::CommitIfOkSuccess => (),
             }
         }
 
         match probe.kind {
             inspect::ProbeKind::NormalizedSelfTyAssembly
             | inspect::ProbeKind::UnsizeAssembly
-            | inspect::ProbeKind::UpcastProjectionCompatibility => (),
+            | inspect::ProbeKind::UpcastProjectionCompatibility
+            | inspect::ProbeKind::CommitIfOk => (),
             // We add a candidate for the root evaluation if there
             // is only one way to prove a given goal, e.g. for `WellFormed`.
             //

compiler/rustc_trait_selection/src/solve/inspect/build.rs

@@ -219,6 +219,8 @@ enum WipProbeStep<'tcx> {
     AddGoal(inspect::CanonicalState<'tcx, Goal<'tcx, ty::Predicate<'tcx>>>),
     EvaluateGoals(WipAddedGoalsEvaluation<'tcx>),
     NestedProbe(WipProbe<'tcx>),
+    CommitIfOkStart,
+    CommitIfOkSuccess,
 }
 
 impl<'tcx> WipProbeStep<'tcx> {
@@ -227,6 +229,8 @@ impl<'tcx> WipProbeStep<'tcx> {
             WipProbeStep::AddGoal(goal) => inspect::ProbeStep::AddGoal(goal),
             WipProbeStep::EvaluateGoals(eval) => inspect::ProbeStep::EvaluateGoals(eval.finalize()),
             WipProbeStep::NestedProbe(probe) => inspect::ProbeStep::NestedProbe(probe.finalize()),
+            WipProbeStep::CommitIfOkStart => inspect::ProbeStep::CommitIfOkStart,
+            WipProbeStep::CommitIfOkSuccess => inspect::ProbeStep::CommitIfOkSuccess,
         }
     }
 }
@@ -459,6 +463,29 @@ impl<'tcx> ProofTreeBuilder<'tcx> {
         }
     }
 
+    /// Used by `EvalCtxt::commit_if_ok` to flatten the work done inside
+    /// of the probe into the parent.
+    pub fn integrate_snapshot(&mut self, probe: ProofTreeBuilder<'tcx>) {
+        if let Some(this) = self.as_mut() {
+            match (this, probe.state.unwrap().tree) {
+                (
+                    DebugSolver::Probe(WipProbe { steps, .. })
+                    | DebugSolver::GoalEvaluationStep(WipGoalEvaluationStep {
+                        evaluation: WipProbe { steps, .. },
+                        ..
+                    }),
+                    DebugSolver::Probe(probe),
+                ) => {
+                    steps.push(WipProbeStep::CommitIfOkStart);
+                    assert_eq!(probe.kind, None);
+                    steps.extend(probe.steps);
+                    steps.push(WipProbeStep::CommitIfOkSuccess);
+                }
+                _ => unreachable!(),
+            }
+        }
+    }
+
     pub fn new_evaluate_added_goals(&mut self) -> ProofTreeBuilder<'tcx> {
         self.nested(|| WipAddedGoalsEvaluation { evaluations: vec![], result: None })
     }

compiler/rustc_trait_selection/src/solve/mod.rs

@@ -297,25 +297,40 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
     fn try_normalize_ty(
         &mut self,
         param_env: ty::ParamEnv<'tcx>,
-        mut ty: Ty<'tcx>,
-    ) -> Result<Option<Ty<'tcx>>, NoSolution> {
-        for _ in 0..self.local_overflow_limit() {
-            let ty::Alias(_, projection_ty) = *ty.kind() else {
-                return Ok(Some(ty));
-            };
-
-            let normalized_ty = self.next_ty_infer();
+        ty: Ty<'tcx>,
+    ) -> Option<Ty<'tcx>> {
+        self.try_normalize_ty_recur(param_env, 0, ty)
+    }
+
+    fn try_normalize_ty_recur(
+        &mut self,
+        param_env: ty::ParamEnv<'tcx>,
+        depth: usize,
+        ty: Ty<'tcx>,
+    ) -> Option<Ty<'tcx>> {
+        if depth >= self.local_overflow_limit() {
+            return None;
+        }
+
+        let ty::Alias(_, projection_ty) = *ty.kind() else {
+            return Some(ty);
+        };
+
+        match self.commit_if_ok(|this| {
+            let normalized_ty = this.next_ty_infer();
             let normalizes_to_goal = Goal::new(
-                self.tcx(),
+                this.tcx(),
                 param_env,
                 ty::ProjectionPredicate { projection_ty, term: normalized_ty.into() },
             );
-            self.add_goal(normalizes_to_goal);
-            self.try_evaluate_added_goals()?;
-            ty = self.resolve_vars_if_possible(normalized_ty);
+            this.add_goal(normalizes_to_goal);
+            this.try_evaluate_added_goals()?;
+            let ty = this.resolve_vars_if_possible(normalized_ty);
+            Ok(this.try_normalize_ty_recur(param_env, depth + 1, ty))
+        }) {
+            Ok(ty) => ty,
+            Err(NoSolution) => Some(ty),
         }
-
-        Ok(None)
     }
 }
 

compiler/rustc_trait_selection/src/solve/trait_goals.rs

@@ -471,7 +471,7 @@ impl<'tcx> assembly::GoalKind<'tcx> for TraitPredicate<'tcx> {
             let a_ty = goal.predicate.self_ty();
             // We need to normalize the b_ty since it's destructured as a `dyn Trait`.
             let Some(b_ty) =
-                ecx.try_normalize_ty(goal.param_env, goal.predicate.trait_ref.args.type_at(1))?
+                ecx.try_normalize_ty(goal.param_env, goal.predicate.trait_ref.args.type_at(1))
             else {
                 return ecx.evaluate_added_goals_and_make_canonical_response(Certainty::OVERFLOW);
             };
@@ -538,9 +538,8 @@ impl<'tcx> assembly::GoalKind<'tcx> for TraitPredicate<'tcx> {
             let b_ty = match ecx
                 .try_normalize_ty(goal.param_env, goal.predicate.trait_ref.args.type_at(1))
             {
-                Ok(Some(b_ty)) => b_ty,
-                Ok(None) => return vec![misc_candidate(ecx, Certainty::OVERFLOW)],
-                Err(_) => return vec![],
+                Some(b_ty) => b_ty,
+                None => return vec![misc_candidate(ecx, Certainty::OVERFLOW)],
             };
 
             let goal = goal.with(ecx.tcx(), (a_ty, b_ty));