Rollup merge of #114287 - lcnr:overflow, r=compiler-errors

update overflow handling in the new trait solver

implements https://hackmd.io/QY0dfEOgSNWwU4oiGnVRLw?view. I want to clean up this doc and add it to the rustc-dev-guide, but I think this PR is ready for merge as is, even without the dev-guide entry.

r? `@compiler-errors`

Author: compiler-errors

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

@@ -1,7 +1,6 @@
 use std::ops::ControlFlow;
 
 use rustc_data_structures::intern::Interned;
-use rustc_query_system::cache::Cache;
 
 use crate::infer::canonical::{CanonicalVarValues, QueryRegionConstraints};
 use crate::traits::query::NoSolution;
@@ -11,9 +10,10 @@ use crate::ty::{
     TypeVisitor,
 };
 
+mod cache;
 pub mod inspect;
 
-pub type EvaluationCache<'tcx> = Cache<CanonicalInput<'tcx>, QueryResult<'tcx>>;
+pub use cache::{CacheData, EvaluationCache};
 
 /// A goal is a statement, i.e. `predicate`, we want to prove
 /// given some assumptions, i.e. `param_env`.

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

@@ -0,0 +1,100 @@
+use super::{CanonicalInput, QueryResult};
+use crate::ty::TyCtxt;
+use rustc_data_structures::fx::{FxHashMap, FxHashSet};
+use rustc_data_structures::sync::Lock;
+use rustc_query_system::cache::WithDepNode;
+use rustc_query_system::dep_graph::DepNodeIndex;
+use rustc_session::Limit;
+/// The trait solver cache used by `-Ztrait-solver=next`.
+///
+/// FIXME(@lcnr): link to some official documentation of how
+/// this works.
+#[derive(Default)]
+pub struct EvaluationCache<'tcx> {
+    map: Lock<FxHashMap<CanonicalInput<'tcx>, CacheEntry<'tcx>>>,
+}
+
+pub struct CacheData<'tcx> {
+    pub result: QueryResult<'tcx>,
+    pub reached_depth: usize,
+    pub encountered_overflow: bool,
+}
+
+impl<'tcx> EvaluationCache<'tcx> {
+    /// Insert a final result into the global cache.
+    pub fn insert(
+        &self,
+        key: CanonicalInput<'tcx>,
+        reached_depth: usize,
+        did_overflow: bool,
+        cycle_participants: FxHashSet<CanonicalInput<'tcx>>,
+        dep_node: DepNodeIndex,
+        result: QueryResult<'tcx>,
+    ) {
+        let mut map = self.map.borrow_mut();
+        let entry = map.entry(key).or_default();
+        let data = WithDepNode::new(dep_node, result);
+        entry.cycle_participants.extend(cycle_participants);
+        if did_overflow {
+            entry.with_overflow.insert(reached_depth, data);
+        } else {
+            entry.success = Some(Success { data, reached_depth });
+        }
+    }
+
+    /// Try to fetch a cached result, checking the recursion limit
+    /// and handling root goals of coinductive cycles.
+    ///
+    /// If this returns `Some` the cache result can be used.
+    pub fn get(
+        &self,
+        tcx: TyCtxt<'tcx>,
+        key: CanonicalInput<'tcx>,
+        cycle_participant_in_stack: impl FnOnce(&FxHashSet<CanonicalInput<'tcx>>) -> bool,
+        available_depth: Limit,
+    ) -> Option<CacheData<'tcx>> {
+        let map = self.map.borrow();
+        let entry = map.get(&key)?;
+
+        if cycle_participant_in_stack(&entry.cycle_participants) {
+            return None;
+        }
+
+        if let Some(ref success) = entry.success {
+            if available_depth.value_within_limit(success.reached_depth) {
+                return Some(CacheData {
+                    result: success.data.get(tcx),
+                    reached_depth: success.reached_depth,
+                    encountered_overflow: false,
+                });
+            }
+        }
+
+        entry.with_overflow.get(&available_depth.0).map(|e| CacheData {
+            result: e.get(tcx),
+            reached_depth: available_depth.0,
+            encountered_overflow: true,
+        })
+    }
+}
+
+struct Success<'tcx> {
+    data: WithDepNode<QueryResult<'tcx>>,
+    reached_depth: usize,
+}
+
+/// The cache entry for a goal `CanonicalInput`.
+///
+/// This contains results whose computation never hit the
+/// recursion limit in `success`, and all results which hit
+/// the recursion limit in `with_overflow`.
+#[derive(Default)]
+struct CacheEntry<'tcx> {
+    success: Option<Success<'tcx>>,
+    /// We have to be careful when caching roots of cycles.
+    ///
+    /// See the doc comment of `StackEntry::cycle_participants` for more
+    /// details.
+    cycle_participants: FxHashSet<CanonicalInput<'tcx>>,
+    with_overflow: FxHashMap<usize, WithDepNode<QueryResult<'tcx>>>,
+}

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

@@ -1338,12 +1338,25 @@ impl<'tcx> ToPredicate<'tcx> for PolyTypeOutlivesPredicate<'tcx> {
     }
 }
 
+impl<'tcx> ToPredicate<'tcx> for ProjectionPredicate<'tcx> {
+    fn to_predicate(self, tcx: TyCtxt<'tcx>) -> Predicate<'tcx> {
+        ty::Binder::dummy(PredicateKind::Clause(ClauseKind::Projection(self))).to_predicate(tcx)
+    }
+}
+
 impl<'tcx> ToPredicate<'tcx> for PolyProjectionPredicate<'tcx> {
     fn to_predicate(self, tcx: TyCtxt<'tcx>) -> Predicate<'tcx> {
         self.map_bound(|p| PredicateKind::Clause(ClauseKind::Projection(p))).to_predicate(tcx)
     }
 }
 
+impl<'tcx> ToPredicate<'tcx, Clause<'tcx>> for ProjectionPredicate<'tcx> {
+    fn to_predicate(self, tcx: TyCtxt<'tcx>) -> Clause<'tcx> {
+        let p: Predicate<'tcx> = self.to_predicate(tcx);
+        p.expect_clause()
+    }
+}
+
 impl<'tcx> ToPredicate<'tcx, Clause<'tcx>> for PolyProjectionPredicate<'tcx> {
     fn to_predicate(self, tcx: TyCtxt<'tcx>) -> Clause<'tcx> {
         let p: Predicate<'tcx> = self.to_predicate(tcx);

Diff for: compiler/rustc_trait_selection/src/solve/alias_relate.rs

@@ -162,7 +162,7 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
         self.add_goal(Goal::new(
             self.tcx(),
             param_env,
-            ty::Binder::dummy(ty::ProjectionPredicate { projection_ty: alias, term: other }),
+            ty::ProjectionPredicate { projection_ty: alias, term: other },
         ));
 
         Ok(())

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

@@ -1,6 +1,5 @@
 //! Code shared by trait and projection goals for candidate assembly.
 
-use super::search_graph::OverflowHandler;
 use super::{EvalCtxt, SolverMode};
 use crate::traits::coherence;
 use rustc_hir::def_id::DefId;
@@ -315,7 +314,7 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
             return ambig;
         }
 
-        let mut candidates = self.assemble_candidates_via_self_ty(goal);
+        let mut candidates = self.assemble_candidates_via_self_ty(goal, 0);
 
         self.assemble_blanket_impl_candidates(goal, &mut candidates);
 
@@ -351,6 +350,7 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
     fn assemble_candidates_via_self_ty<G: GoalKind<'tcx>>(
         &mut self,
         goal: Goal<'tcx, G>,
+        num_steps: usize,
     ) -> Vec<Candidate<'tcx>> {
         debug_assert_eq!(goal, self.resolve_vars_if_possible(goal));
         if let Some(ambig) = self.assemble_self_ty_infer_ambiguity_response(goal) {
@@ -369,7 +369,7 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
 
         self.assemble_coherence_unknowable_candidates(goal, &mut candidates);
 
-        self.assemble_candidates_after_normalizing_self_ty(goal, &mut candidates);
+        self.assemble_candidates_after_normalizing_self_ty(goal, &mut candidates, num_steps);
 
         candidates
     }
@@ -393,49 +393,40 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
         &mut self,
         goal: Goal<'tcx, G>,
         candidates: &mut Vec<Candidate<'tcx>>,
+        num_steps: usize,
     ) {
         let tcx = self.tcx();
         let &ty::Alias(_, projection_ty) = goal.predicate.self_ty().kind() else { return };
 
-        let normalized_self_candidates: Result<_, NoSolution> =
-            self.probe(|_| CandidateKind::NormalizedSelfTyAssembly).enter(|ecx| {
-                ecx.with_incremented_depth(
-                    |ecx| {
-                        let result = ecx.evaluate_added_goals_and_make_canonical_response(
-                            Certainty::OVERFLOW,
-                        )?;
-                        Ok(vec![Candidate {
-                            source: CandidateSource::BuiltinImpl(BuiltinImplSource::Misc),
-                            result,
-                        }])
-                    },
-                    |ecx| {
-                        let normalized_ty = ecx.next_ty_infer();
-                        let normalizes_to_goal = goal.with(
-                            tcx,
-                            ty::Binder::dummy(ty::ProjectionPredicate {
-                                projection_ty,
-                                term: normalized_ty.into(),
-                            }),
-                        );
-                        ecx.add_goal(normalizes_to_goal);
-                        let _ = ecx.try_evaluate_added_goals().inspect_err(|_| {
-                            debug!("self type normalization failed");
-                        })?;
-                        let normalized_ty = ecx.resolve_vars_if_possible(normalized_ty);
-                        debug!(?normalized_ty, "self type normalized");
-                        // NOTE: Alternatively we could call `evaluate_goal` here and only
-                        // have a `Normalized` candidate. This doesn't work as long as we
-                        // use `CandidateSource` in winnowing.
-                        let goal = goal.with(tcx, goal.predicate.with_self_ty(tcx, normalized_ty));
-                        Ok(ecx.assemble_candidates_via_self_ty(goal))
-                    },
-                )
-            });
-
-        if let Ok(normalized_self_candidates) = normalized_self_candidates {
-            candidates.extend(normalized_self_candidates);
-        }
+        candidates.extend(self.probe(|_| CandidateKind::NormalizedSelfTyAssembly).enter(|ecx| {
+            if num_steps < ecx.local_overflow_limit() {
+                let normalized_ty = ecx.next_ty_infer();
+                let normalizes_to_goal = goal.with(
+                    tcx,
+                    ty::ProjectionPredicate { projection_ty, term: normalized_ty.into() },
+                );
+                ecx.add_goal(normalizes_to_goal);
+                if let Err(NoSolution) = ecx.try_evaluate_added_goals() {
+                    debug!("self type normalization failed");
+                    return vec![];
+                }
+                let normalized_ty = ecx.resolve_vars_if_possible(normalized_ty);
+                debug!(?normalized_ty, "self type normalized");
+                // NOTE: Alternatively we could call `evaluate_goal` here and only
+                // have a `Normalized` candidate. This doesn't work as long as we
+                // use `CandidateSource` in winnowing.
+                let goal = goal.with(tcx, goal.predicate.with_self_ty(tcx, normalized_ty));
+                ecx.assemble_candidates_via_self_ty(goal, num_steps + 1)
+            } else {
+                match ecx.evaluate_added_goals_and_make_canonical_response(Certainty::OVERFLOW) {
+                    Ok(result) => vec![Candidate {
+                        source: CandidateSource::BuiltinImpl(BuiltinImplSource::Misc),
+                        result,
+                    }],
+                    Err(NoSolution) => vec![],
+                }
+            }
+        }));
     }
 
     #[instrument(level = "debug", skip_all)]
@@ -533,7 +524,7 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
             ty::Alias(_, _) | ty::Placeholder(..) | ty::Error(_) => (),
 
             // FIXME: These should ideally not exist as a self type. It would be nice for
-            // the builtin auto trait impls of generators should instead directly recurse
+            // the builtin auto trait impls of generators to instead directly recurse
             // into the witness.
             ty::GeneratorWitness(_) | ty::GeneratorWitnessMIR(_, _) => (),