Rollup merge of #107739 - spastorino:check-overflow-evaluate_canonical_goal, r=lcnr

Check for overflow in evaluate_canonical_goal

r? `@lcnr`

Author: matthiaskrgr

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::OverflowHandler;
 use crate::traits::ObligationCause;
 
 mod assembly;
@@ -210,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> {
@@ -485,35 +475,38 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
         mut goals: Vec<Goal<'tcx, ty::Predicate<'tcx>>>,
     ) -> Result<Certainty, NoSolution> {
         let mut new_goals = Vec::new();
-        self.repeat_while_none(|this| {
-            let mut has_changed = Err(Certainty::Yes);
-            for goal in goals.drain(..) {
-                let (changed, certainty) = match this.evaluate_goal(goal) {
-                    Ok(result) => result,
-                    Err(NoSolution) => return Some(Err(NoSolution)),
-                };
-
-                if changed {
-                    has_changed = Ok(());
-                }
+        self.repeat_while_none(
+            |_| Ok(Certainty::Maybe(MaybeCause::Overflow)),
+            |this| {
+                let mut has_changed = Err(Certainty::Yes);
+                for goal in goals.drain(..) {
+                    let (changed, certainty) = match this.evaluate_goal(goal) {
+                        Ok(result) => result,
+                        Err(NoSolution) => return Some(Err(NoSolution)),
+                    };
+
+                    if changed {
+                        has_changed = Ok(());
+                    }
 
-                match certainty {
-                    Certainty::Yes => {}
-                    Certainty::Maybe(_) => {
-                        new_goals.push(goal);
-                        has_changed = has_changed.map_err(|c| c.unify_and(certainty));
+                    match certainty {
+                        Certainty::Yes => {}
+                        Certainty::Maybe(_) => {
+                            new_goals.push(goal);
+                            has_changed = has_changed.map_err(|c| c.unify_and(certainty));
+                        }
                     }
                 }
-            }
 
-            match has_changed {
-                Ok(()) => {
-                    mem::swap(&mut new_goals, &mut goals);
-                    None
+                match has_changed {
+                    Ok(()) => {
+                        mem::swap(&mut new_goals, &mut goals);
+                        None
+                    }
+                    Err(certainty) => Some(Ok(certainty)),
                 }
-                Err(certainty) => Some(Ok(certainty)),
-            }
-        })
+            },
+        )
     }
 
     // Recursively evaluates a list of goals to completion, making a query response.

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

@@ -3,6 +3,7 @@ mod overflow;
 
 use self::cache::ProvisionalEntry;
 use super::{CanonicalGoal, Certainty, MaybeCause, QueryResult};
+pub(super) use crate::solve::search_graph::overflow::OverflowHandler;
 use cache::ProvisionalCache;
 use overflow::OverflowData;
 use rustc_index::vec::IndexVec;
@@ -46,7 +47,7 @@ impl<'tcx> SearchGraph<'tcx> {
     ///
     /// This correctly updates the provisional cache if there is a cycle.
     #[instrument(level = "debug", skip(self, tcx), ret)]
-    pub(super) fn try_push_stack(
+    fn try_push_stack(
         &mut self,
         tcx: TyCtxt<'tcx>,
         goal: CanonicalGoal<'tcx>,
@@ -121,19 +122,19 @@ impl<'tcx> SearchGraph<'tcx> {
     ///
     /// FIXME: Refer to the rustc-dev-guide entry once it exists.
     #[instrument(level = "debug", skip(self, tcx, actual_goal), ret)]
-    pub(super) fn try_finalize_goal(
+    fn try_finalize_goal(
         &mut self,
         tcx: TyCtxt<'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;
         let provisional_entry_index = *cache.lookup_table.get(&goal).unwrap();
         let provisional_entry = &mut cache.entries[provisional_entry_index];
-        let depth = provisional_entry.depth;
         // We eagerly update the response in the cache here. If we have to reevaluate
         // this goal we use the new response when hitting a cycle, and we definitely
         // want to access the final response whenever we look at the cache.
@@ -157,29 +158,72 @@ impl<'tcx> SearchGraph<'tcx> {
             self.stack.push(StackElem { goal, has_been_used: false });
             false
         } else {
-            // If not, we're done with this goal.
-            //
-            // Check whether that this goal doesn't depend on a goal deeper on the stack
-            // and if so, move it and all nested goals to the global cache.
-            //
-            // Note that if any nested goal were to depend on something deeper on the stack,
-            // this would have also updated the depth of the current goal.
-            if depth == self.stack.next_index() {
-                for (i, entry) in cache.entries.drain_enumerated(provisional_entry_index.index()..)
-                {
-                    let actual_index = cache.lookup_table.remove(&entry.goal);
-                    debug_assert_eq!(Some(i), actual_index);
-                    debug_assert!(entry.depth == depth);
-                    cache::try_move_finished_goal_to_global_cache(
-                        tcx,
-                        &mut self.overflow_data,
-                        &self.stack,
-                        entry.goal,
-                        entry.response,
-                    );
-                }
-            }
+            self.try_move_finished_goal_to_global_cache(tcx, stack_elem);
             true
         }
     }
+
+    fn try_move_finished_goal_to_global_cache(
+        &mut self,
+        tcx: TyCtxt<'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 = &mut cache.entries[provisional_entry_index];
+        let depth = provisional_entry.depth;
+
+        // If not, we're done with this goal.
+        //
+        // Check whether that this goal doesn't depend on a goal deeper on the stack
+        // and if so, move it and all nested goals to the global cache.
+        //
+        // Note that if any nested goal were to depend on something deeper on the stack,
+        // this would have also updated the depth of the current goal.
+        if depth == self.stack.next_index() {
+            for (i, entry) in cache.entries.drain_enumerated(provisional_entry_index.index()..) {
+                let actual_index = cache.lookup_table.remove(&entry.goal);
+                debug_assert_eq!(Some(i), actual_index);
+                debug_assert!(entry.depth == depth);
+                cache::try_move_finished_goal_to_global_cache(
+                    tcx,
+                    &mut self.overflow_data,
+                    &self.stack,
+                    entry.goal,
+                    entry.response,
+                );
+            }
+        }
+    }
+
+    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

@@ -50,6 +50,42 @@ impl OverflowData {
     }
 }
 
+pub(in crate::solve) trait OverflowHandler<'tcx> {
+    fn search_graph(&mut self) -> &mut SearchGraph<'tcx>;
+
+    fn repeat_while_none<T>(
+        &mut self,
+        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;
+        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;
+                return result;
+            }
+
+            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();
+        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,
@@ -60,25 +96,3 @@ impl<'tcx> SearchGraph<'tcx> {
         response_no_constraints(tcx, goal, Certainty::Maybe(MaybeCause::Overflow))
     }
 }
-
-impl<'tcx> EvalCtxt<'_, 'tcx> {
-    /// A `while`-loop which tracks overflow.
-    pub fn repeat_while_none(
-        &mut self,
-        mut loop_body: impl FnMut(&mut Self) -> Option<Result<Certainty, NoSolution>>,
-    ) -> Result<Certainty, 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) {
-            if let Some(result) = loop_body(self) {
-                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 = start_depth;
-        self.search_graph.overflow_data.deal_with_overflow();
-        Ok(Certainty::Maybe(MaybeCause::Overflow))
-    }
-}