Check for overflow in evaluate_canonical_goal

Author: spastorino

compiler/rustc_trait_selection/src/solve/mod.rs

@@ -211,27 +211,16 @@ impl<'a, 'tcx> EvalCtxt<'a, 'tcx> {
         search_graph: &'a mut search_graph::SearchGraph<'tcx>,
         canonical_goal: CanonicalGoal<'tcx>,
     ) -> QueryResult<'tcx> {
-        match search_graph.try_push_stack(tcx, canonical_goal) {
-            Ok(()) => {}
-            // Our goal is already on the stack, eager return.
-            Err(response) => return response,
-        }
-
-        // We may have to repeatedly recompute the goal in case of coinductive cycles,
-        // check out the `cache` module for more information.
+        // Deal with overflow, caching, and coinduction.
         //
-        // FIXME: Similar to `evaluate_all`, this has to check for overflow.
-        loop {
+        // The actual solver logic happens in `ecx.compute_goal`.
+        search_graph.with_new_goal(tcx, canonical_goal, |search_graph| {
             let (ref infcx, goal, var_values) =
                 tcx.infer_ctxt().build_with_canonical(DUMMY_SP, &canonical_goal);
             let mut ecx =
                 EvalCtxt { infcx, var_values, search_graph, in_projection_eq_hack: false };
-            let result = ecx.compute_goal(goal);
-
-            if search_graph.try_finalize_goal(tcx, canonical_goal, result) {
-                return result;
-            }
-        }
+            ecx.compute_goal(goal)
+        })
     }
 
     fn make_canonical_response(&self, certainty: Certainty) -> QueryResult<'tcx> {
@@ -487,7 +476,7 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
     ) -> Result<Certainty, NoSolution> {
         let mut new_goals = Vec::new();
         self.repeat_while_none(
-            |_| Certainty::Maybe(MaybeCause::Overflow),
+            |_| Ok(Certainty::Maybe(MaybeCause::Overflow)),
             |this| {
                 let mut has_changed = Err(Certainty::Yes);
                 for goal in goals.drain(..) {

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

@@ -3,6 +3,7 @@ pub(crate) mod overflow;
 
 use self::cache::ProvisionalEntry;
 use super::{CanonicalGoal, Certainty, MaybeCause, QueryResult};
+use crate::solve::search_graph::overflow::OverflowHandler;
 use cache::ProvisionalCache;
 use overflow::OverflowData;
 use rustc_index::vec::IndexVec;
@@ -13,7 +14,7 @@ rustc_index::newtype_index! {
     pub struct StackDepth {}
 }
 
-struct StackElem<'tcx> {
+pub(crate) struct StackElem<'tcx> {
     goal: CanonicalGoal<'tcx>,
     has_been_used: bool,
 }
@@ -127,7 +128,8 @@ impl<'tcx> SearchGraph<'tcx> {
         actual_goal: CanonicalGoal<'tcx>,
         response: QueryResult<'tcx>,
     ) -> bool {
-        let StackElem { goal, has_been_used } = self.stack.pop().unwrap();
+        let stack_elem = self.stack.pop().unwrap();
+        let StackElem { goal, has_been_used } = stack_elem;
         assert_eq!(goal, actual_goal);
 
         let cache = &mut self.provisional_cache;
@@ -156,18 +158,19 @@ impl<'tcx> SearchGraph<'tcx> {
             self.stack.push(StackElem { goal, has_been_used: false });
             false
         } else {
-            self.try_move_finished_goal_to_global_cache(tcx, &goal);
+            self.try_move_finished_goal_to_global_cache(tcx, stack_elem);
             true
         }
     }
 
     pub(super) fn try_move_finished_goal_to_global_cache(
         &mut self,
         tcx: TyCtxt<'tcx>,
-        goal: &CanonicalGoal<'tcx>,
+        stack_elem: StackElem<'tcx>,
     ) {
+        let StackElem { goal, .. } = stack_elem;
         let cache = &mut self.provisional_cache;
-        let provisional_entry_index = *cache.lookup_table.get(goal).unwrap();
+        let provisional_entry_index = *cache.lookup_table.get(&goal).unwrap();
         let provisional_entry = &mut cache.entries[provisional_entry_index];
         let depth = provisional_entry.depth;
 
@@ -193,4 +196,34 @@ impl<'tcx> SearchGraph<'tcx> {
             }
         }
     }
+
+    pub(super) fn with_new_goal(
+        &mut self,
+        tcx: TyCtxt<'tcx>,
+        canonical_goal: CanonicalGoal<'tcx>,
+        mut loop_body: impl FnMut(&mut Self) -> QueryResult<'tcx>,
+    ) -> QueryResult<'tcx> {
+        match self.try_push_stack(tcx, canonical_goal) {
+            Ok(()) => {}
+            // Our goal is already on the stack, eager return.
+            Err(response) => return response,
+        }
+
+        self.repeat_while_none(
+            |this| {
+                let result = this.deal_with_overflow(tcx, canonical_goal);
+                let stack_elem = this.stack.pop().unwrap();
+                this.try_move_finished_goal_to_global_cache(tcx, stack_elem);
+                result
+            },
+            |this| {
+                let result = loop_body(this);
+                if this.try_finalize_goal(tcx, canonical_goal, result) {
+                    Some(result)
+                } else {
+                    None
+                }
+            },
+        )
+    }
 }

compiler/rustc_trait_selection/src/solve/search_graph/overflow.rs

@@ -55,7 +55,7 @@ pub(crate) trait OverflowHandler<'tcx> {
 
     fn repeat_while_none<T>(
         &mut self,
-        on_overflow: impl FnOnce(&mut Self) -> T,
+        on_overflow: impl FnOnce(&mut Self) -> Result<T, NoSolution>,
         mut loop_body: impl FnMut(&mut Self) -> Option<Result<T, NoSolution>>,
     ) -> Result<T, NoSolution> {
         let start_depth = self.search_graph().overflow_data.additional_depth;
@@ -70,7 +70,7 @@ pub(crate) trait OverflowHandler<'tcx> {
         }
         self.search_graph().overflow_data.additional_depth = start_depth;
         self.search_graph().overflow_data.deal_with_overflow();
-        Ok(on_overflow(self))
+        on_overflow(self)
     }
 }