Implement repeat_while_none for both SearchGraph and EvalCtxt

3 files changed, 38 insertions, 24 deletions

diff --git a/compiler/rustc_trait_selection/src/solve/mod.rs b/compiler/rustc_trait_selection/src/solve/mod.rs
index e3f8f7cddab..358a2bcc7b9 100644
--- a/compiler/rustc_trait_selection/src/solve/mod.rs
+++ b/compiler/rustc_trait_selection/src/solve/mod.rs
@@ -31,6 +31,7 @@ use rustc_middle::ty::{
 };
 use rustc_span::DUMMY_SP;
 
+use crate::solve::search_graph::overflow::OverflowHandler;
 use crate::traits::ObligationCause;
 
 mod assembly;
diff --git a/compiler/rustc_trait_selection/src/solve/search_graph/mod.rs b/compiler/rustc_trait_selection/src/solve/search_graph/mod.rs
index c25a9cddfe7..438bcd9a7d6 100644
--- a/compiler/rustc_trait_selection/src/solve/search_graph/mod.rs
+++ b/compiler/rustc_trait_selection/src/solve/search_graph/mod.rs
@@ -1,5 +1,5 @@
 mod cache;
-mod overflow;
+pub(crate) mod overflow;
 
 use self::cache::ProvisionalEntry;
 use super::{CanonicalGoal, Certainty, MaybeCause, QueryResult};
@@ -18,7 +18,7 @@ struct StackElem<'tcx> {
     has_been_used: bool,
 }
 
-pub(super) struct SearchGraph<'tcx> {
+pub(crate) struct SearchGraph<'tcx> {
     /// The stack of goals currently being computed.
     ///
     /// An element is *deeper* in the stack if its index is *lower*.
diff --git a/compiler/rustc_trait_selection/src/solve/search_graph/overflow.rs b/compiler/rustc_trait_selection/src/solve/search_graph/overflow.rs
index c472dfe5a00..0d6863b1e81 100644
--- a/compiler/rustc_trait_selection/src/solve/search_graph/overflow.rs
+++ b/compiler/rustc_trait_selection/src/solve/search_graph/overflow.rs
@@ -50,36 +50,49 @@ impl OverflowData {
     }
 }
 
-impl<'tcx> SearchGraph<'tcx> {
-    pub fn deal_with_overflow(
-        &mut self,
-        tcx: TyCtxt<'tcx>,
-        goal: Canonical<'tcx, impl Sized>,
-    ) -> QueryResult<'tcx> {
-        self.overflow_data.deal_with_overflow();
-        response_no_constraints(tcx, goal, Certainty::Maybe(MaybeCause::Overflow))
-    }
-}
+pub(crate) trait OverflowHandler<'tcx> {
+    fn search_graph(&mut self) -> &mut SearchGraph<'tcx>;
 
-impl<'tcx> EvalCtxt<'_, 'tcx> {
-    /// A `while`-loop which tracks overflow.
-    pub fn repeat_while_none<T>(
+    fn repeat_while_none<T>(
         &mut self,
-        mut overflow_body: impl FnMut(&mut Self) -> T,
+        on_overflow: impl FnOnce(&mut Self) -> T,
         mut loop_body: impl FnMut(&mut Self) -> Option<Result<T, NoSolution>>,
     ) -> Result<T, NoSolution> {
-        let start_depth = self.search_graph.overflow_data.additional_depth;
-        let depth = self.search_graph.stack.len();
-        while !self.search_graph.overflow_data.has_overflow(depth) {
+        let start_depth = self.search_graph().overflow_data.additional_depth;
+        let depth = self.search_graph().stack.len();
+        while !self.search_graph().overflow_data.has_overflow(depth) {
             if let Some(result) = loop_body(self) {
-                self.search_graph.overflow_data.additional_depth = start_depth;
+                self.search_graph().overflow_data.additional_depth = start_depth;
                 return result;
             }
 
-            self.search_graph.overflow_data.additional_depth += 1;
+            self.search_graph().overflow_data.additional_depth += 1;
         }
-        self.search_graph.overflow_data.additional_depth = start_depth;
-        self.search_graph.overflow_data.deal_with_overflow();
-        Ok(overflow_body(self))
+        self.search_graph().overflow_data.additional_depth = start_depth;
+        self.search_graph().overflow_data.deal_with_overflow();
+        Ok(on_overflow(self))
+    }
+}
+
+impl<'tcx> OverflowHandler<'tcx> for EvalCtxt<'_, 'tcx> {
+    fn search_graph(&mut self) -> &mut SearchGraph<'tcx> {
+        &mut self.search_graph
+    }
+}
+
+impl<'tcx> OverflowHandler<'tcx> for SearchGraph<'tcx> {
+    fn search_graph(&mut self) -> &mut SearchGraph<'tcx> {
+        self
+    }
+}
+
+impl<'tcx> SearchGraph<'tcx> {
+    pub fn deal_with_overflow(
+        &mut self,
+        tcx: TyCtxt<'tcx>,
+        goal: Canonical<'tcx, impl Sized>,
+    ) -> QueryResult<'tcx> {
+        self.overflow_data.deal_with_overflow();
+        response_no_constraints(tcx, goal, Certainty::Maybe(MaybeCause::Overflow))
     }
 }