Make SearchGraph fully generic

3 files changed, 33 insertions, 18 deletions

diff --git a/compiler/rustc_middle/src/traits/solve.rs b/compiler/rustc_middle/src/traits/solve.rs
index 9e979620a44..888e3aec5ea 100644
--- a/compiler/rustc_middle/src/traits/solve.rs
+++ b/compiler/rustc_middle/src/traits/solve.rs
@@ -10,7 +10,7 @@ use crate::ty::{
 
 mod cache;
 
-pub use cache::{CacheData, EvaluationCache};
+pub use cache::EvaluationCache;
 
 pub type Goal<'tcx, P> = ir::solve::Goal<TyCtxt<'tcx>, P>;
 pub type QueryInput<'tcx, P> = ir::solve::QueryInput<TyCtxt<'tcx>, P>;
diff --git a/compiler/rustc_middle/src/traits/solve/cache.rs b/compiler/rustc_middle/src/traits/solve/cache.rs
index dc31114b2c4..72a8d4eb405 100644
--- a/compiler/rustc_middle/src/traits/solve/cache.rs
+++ b/compiler/rustc_middle/src/traits/solve/cache.rs
@@ -5,6 +5,8 @@ use rustc_data_structures::sync::Lock;
 use rustc_query_system::cache::WithDepNode;
 use rustc_query_system::dep_graph::DepNodeIndex;
 use rustc_session::Limit;
+use rustc_type_ir::solve::CacheData;
+
 /// The trait solver cache used by `-Znext-solver`.
 ///
 /// FIXME(@lcnr): link to some official documentation of how
@@ -14,17 +16,9 @@ pub struct EvaluationCache<'tcx> {
     map: Lock<FxHashMap<CanonicalInput<'tcx>, CacheEntry<'tcx>>>,
 }
 
-#[derive(Debug, PartialEq, Eq)]
-pub struct CacheData<'tcx> {
-    pub result: QueryResult<'tcx>,
-    pub proof_tree: Option<&'tcx inspect::CanonicalGoalEvaluationStep<TyCtxt<'tcx>>>,
-    pub additional_depth: usize,
-    pub encountered_overflow: bool,
-}
-
-impl<'tcx> EvaluationCache<'tcx> {
+impl<'tcx> rustc_type_ir::inherent::EvaluationCache<TyCtxt<'tcx>> for &'tcx EvaluationCache<'tcx> {
     /// Insert a final result into the global cache.
-    pub fn insert(
+    fn insert(
         &self,
         tcx: TyCtxt<'tcx>,
         key: CanonicalInput<'tcx>,
@@ -48,7 +42,7 @@ impl<'tcx> EvaluationCache<'tcx> {
         if cfg!(debug_assertions) {
             drop(map);
             let expected = CacheData { result, proof_tree, additional_depth, encountered_overflow };
-            let actual = self.get(tcx, key, [], Limit(additional_depth));
+            let actual = self.get(tcx, key, [], additional_depth);
             if !actual.as_ref().is_some_and(|actual| expected == *actual) {
                 bug!("failed to lookup inserted element for {key:?}: {expected:?} != {actual:?}");
             }
@@ -59,13 +53,13 @@ impl<'tcx> EvaluationCache<'tcx> {
     /// and handling root goals of coinductive cycles.
     ///
     /// If this returns `Some` the cache result can be used.
-    pub fn get(
+    fn get(
         &self,
         tcx: TyCtxt<'tcx>,
         key: CanonicalInput<'tcx>,
         stack_entries: impl IntoIterator<Item = CanonicalInput<'tcx>>,
-        available_depth: Limit,
-    ) -> Option<CacheData<'tcx>> {
+        available_depth: usize,
+    ) -> Option<CacheData<TyCtxt<'tcx>>> {
         let map = self.map.borrow();
         let entry = map.get(&key)?;
 
@@ -76,7 +70,7 @@ impl<'tcx> EvaluationCache<'tcx> {
         }
 
         if let Some(ref success) = entry.success {
-            if available_depth.value_within_limit(success.additional_depth) {
+            if Limit(available_depth).value_within_limit(success.additional_depth) {
                 let QueryData { result, proof_tree } = success.data.get(tcx);
                 return Some(CacheData {
                     result,
@@ -87,12 +81,12 @@ impl<'tcx> EvaluationCache<'tcx> {
             }
         }
 
-        entry.with_overflow.get(&available_depth.0).map(|e| {
+        entry.with_overflow.get(&available_depth).map(|e| {
             let QueryData { result, proof_tree } = e.get(tcx);
             CacheData {
                 result,
                 proof_tree,
-                additional_depth: available_depth.0,
+                additional_depth: available_depth,
                 encountered_overflow: true,
             }
         })
diff --git a/compiler/rustc_middle/src/ty/context.rs b/compiler/rustc_middle/src/ty/context.rs
index e2f15dac019..eec7fa8db1d 100644
--- a/compiler/rustc_middle/src/ty/context.rs
+++ b/compiler/rustc_middle/src/ty/context.rs
@@ -71,6 +71,7 @@ use rustc_target::abi::{FieldIdx, Layout, LayoutS, TargetDataLayout, VariantIdx}
 use rustc_target::spec::abi;
 use rustc_type_ir::fold::TypeFoldable;
 use rustc_type_ir::lang_items::TraitSolverLangItem;
+use rustc_type_ir::solve::SolverMode;
 use rustc_type_ir::TyKind::*;
 use rustc_type_ir::{CollectAndApply, Interner, TypeFlags, WithCachedTypeInfo};
 use tracing::{debug, instrument};
@@ -139,10 +140,30 @@ impl<'tcx> Interner for TyCtxt<'tcx> {
     type Clause = Clause<'tcx>;
     type Clauses = ty::Clauses<'tcx>;
 
+    type DepNodeIndex = DepNodeIndex;
+    fn with_cached_task<T>(self, task: impl FnOnce() -> T) -> (T, DepNodeIndex) {
+        self.dep_graph.with_anon_task(self, crate::dep_graph::dep_kinds::TraitSelect, task)
+    }
+
+    type EvaluationCache = &'tcx solve::EvaluationCache<'tcx>;
+    fn evaluation_cache(self, mode: SolverMode) -> &'tcx solve::EvaluationCache<'tcx> {
+        match mode {
+            SolverMode::Normal => &self.new_solver_evaluation_cache,
+            SolverMode::Coherence => &self.new_solver_coherence_evaluation_cache,
+        }
+    }
+
     fn expand_abstract_consts<T: TypeFoldable<TyCtxt<'tcx>>>(self, t: T) -> T {
         self.expand_abstract_consts(t)
     }
 
+    fn mk_external_constraints(
+        self,
+        data: ExternalConstraintsData<Self>,
+    ) -> ExternalConstraints<'tcx> {
+        self.mk_external_constraints(data)
+    }
+
     fn mk_canonical_var_infos(self, infos: &[ty::CanonicalVarInfo<Self>]) -> Self::CanonicalVars {
         self.mk_canonical_var_infos(infos)
     }