Some cleanups and add a solve_limited

chalk-engine/src/context.rs

@@ -179,6 +179,26 @@ pub trait Context: Clone + Debug {
     fn identity_constrained_subst(
         goal: &Self::UCanonicalGoalInEnvironment,
     ) -> Self::CanonicalConstrainedSubst;
+
+    /// Convert the context's goal type into the `HhGoal` type that
+    /// the SLG solver understands. The expectation is that the
+    /// context's goal type has the same set of variants, but with
+    /// different names and a different setup. If you inspect
+    /// `HhGoal`, you will see that this is a "shallow" or "lazy"
+    /// conversion -- that is, we convert the outermost goal into an
+    /// `HhGoal`, but the goals contained within are left as context
+    /// goals.
+    fn into_hh_goal(goal: Self::Goal) -> HhGoal<Self>;
+
+    // Used by: simplify
+    fn add_clauses(env: &Self::Environment, clauses: Self::ProgramClauses) -> Self::Environment;
+
+    /// Upcast this domain goal into a more general goal.
+    fn into_goal(domain_goal: Self::DomainGoal) -> Self::Goal;
+
+    /// Selects the next appropriate subgoal index for evaluation.
+    /// Used by: logic
+    fn next_subgoal_index(ex_clause: &ExClause<Self>) -> usize;
 }
 
 pub trait ContextOps<C: Context>: Sized + Clone + Debug + AggregateOps<C> {
@@ -247,32 +267,13 @@ pub trait AggregateOps<C: Context> {
         &self,
         root_goal: &C::UCanonicalGoalInEnvironment,
         answers: impl AnswerStream<C>,
+        should_continue: impl Fn() -> bool,
     ) -> Option<C::Solution>;
 }
 
 /// An "inference table" contains the state to support unification and
 /// other operations on terms.
-pub trait InferenceTable<C: Context>: ResolventOps<C> + TruncateOps<C> + UnificationOps<C> {
-    /// Convert the context's goal type into the `HhGoal` type that
-    /// the SLG solver understands. The expectation is that the
-    /// context's goal type has the same set of variants, but with
-    /// different names and a different setup. If you inspect
-    /// `HhGoal`, you will see that this is a "shallow" or "lazy"
-    /// conversion -- that is, we convert the outermost goal into an
-    /// `HhGoal`, but the goals contained within are left as context
-    /// goals.
-    fn into_hh_goal(&mut self, goal: C::Goal) -> HhGoal<C>;
-
-    // Used by: simplify
-    fn add_clauses(&mut self, env: &C::Environment, clauses: C::ProgramClauses) -> C::Environment;
-
-    /// Upcast this domain goal into a more general goal.
-    fn into_goal(&self, domain_goal: C::DomainGoal) -> C::Goal;
-
-    /// Selects the next appropriate subgoal index for evaluation.
-    /// Used by: logic
-    fn next_subgoal_index(&mut self, ex_clause: &ExClause<C>) -> usize;
-}
+pub trait InferenceTable<C: Context>: ResolventOps<C> + TruncateOps<C> + UnificationOps<C> {}
 
 /// Error type for the `UnificationOps::program_clauses` method --
 /// indicates that the complete set of program clauses for this goal
@@ -388,6 +389,10 @@ pub enum AnswerResult<C: Context> {
 
     /// No answer could be returned because the goal has floundered.
     Floundered,
+
+    // No answer could be returned *yet*, because we exceeded our
+    // quantum (`should_continue` returned false).
+    QuantumExceeded,
 }
 
 impl<C: Context> AnswerResult<C> {
@@ -411,6 +416,13 @@ impl<C: Context> AnswerResult<C> {
             _ => false,
         }
     }
+
+    pub fn is_quantum_exceeded(&self) -> bool {
+        match self {
+            Self::QuantumExceeded => true,
+            _ => false,
+        }
+    }
 }
 
 impl<C: Context> Debug for AnswerResult<C> {
@@ -419,21 +431,21 @@ impl<C: Context> Debug for AnswerResult<C> {
             AnswerResult::Answer(answer) => write!(fmt, "{:?}", answer),
             AnswerResult::Floundered => write!(fmt, "Floundered"),
             AnswerResult::NoMoreSolutions => write!(fmt, "None"),
+            AnswerResult::QuantumExceeded => write!(fmt, "QuantumExceeded"),
         }
     }
 }
 
 pub trait AnswerStream<C: Context> {
     /// Gets the next answer for a given goal, but doesn't increment the answer index.
     /// Calling this or `next_answer` again will give the same answer.
-    fn peek_answer(&mut self) -> AnswerResult<C>;
+    fn peek_answer(&mut self, should_continue: impl Fn() -> bool) -> AnswerResult<C>;
 
     /// Gets the next answer for a given goal, incrementing the answer index.
     /// Calling this or `peek_answer` again will give the next answer.
-    fn next_answer(&mut self) -> AnswerResult<C>;
+    fn next_answer(&mut self, should_continue: impl Fn() -> bool) -> AnswerResult<C>;
 
     /// Invokes `test` with each possible future answer, returning true immediately
     /// if we find any answer for which `test` returns true.
-    fn any_future_answer(&mut self, test: impl FnMut(&C::InferenceNormalizedSubst) -> bool)
-        -> bool;
+    fn any_future_answer(&self, test: impl Fn(&C::InferenceNormalizedSubst) -> bool) -> bool;
 }

chalk-engine/src/forest.rs

@@ -59,8 +59,9 @@ impl<C: Context> Forest<C> {
         &mut self,
         context: &impl ContextOps<C>,
         goal: &C::UCanonicalGoalInEnvironment,
+        should_continue: impl Fn() -> bool,
     ) -> Option<C::Solution> {
-        context.make_solution(&goal, self.iter_answers(context, goal))
+        context.make_solution(&goal, self.iter_answers(context, goal), should_continue)
     }
 
     /// Solves a given goal, producing the solution. This will do only
@@ -75,7 +76,7 @@ impl<C: Context> Forest<C> {
     ) -> bool {
         let mut answers = self.iter_answers(context, goal);
         loop {
-            let subst = match answers.next_answer() {
+            let subst = match answers.next_answer(|| true) {
                 AnswerResult::Answer(answer) => {
                     if !answer.ambiguous {
                         SubstitutionResult::Definite(context.constrained_subst_from_answer(answer))
@@ -87,9 +88,10 @@ impl<C: Context> Forest<C> {
                 AnswerResult::NoMoreSolutions => {
                     return true;
                 }
+                AnswerResult::QuantumExceeded => continue,
             };
 
-            if !f(subst, !answers.peek_answer().is_no_more_solutions()) {
+            if !f(subst, !answers.peek_answer(|| true).is_no_more_solutions()) {
                 return false;
             }
         }
@@ -157,7 +159,7 @@ impl<'me, C: Context, CO: ContextOps<C>> AnswerStream<C> for ForestSolver<'me, C
     /// # Panics
     ///
     /// Panics if a negative cycle was detected.
-    fn peek_answer(&mut self) -> AnswerResult<C> {
+    fn peek_answer(&mut self, should_continue: impl Fn() -> bool) -> AnswerResult<C> {
         loop {
             match self
                 .forest
@@ -178,7 +180,11 @@ impl<'me, C: Context, CO: ContextOps<C>> AnswerStream<C> for ForestSolver<'me, C
                     return AnswerResult::NoMoreSolutions;
                 }
 
-                Err(RootSearchFail::QuantumExceeded) => {}
+                Err(RootSearchFail::QuantumExceeded) => {
+                    if !should_continue() {
+                        return AnswerResult::QuantumExceeded;
+                    }
+                }
 
                 Err(RootSearchFail::NegativeCycle) => {
                     // Negative cycles *ought* to be avoided by construction. Hence panic
@@ -192,16 +198,13 @@ impl<'me, C: Context, CO: ContextOps<C>> AnswerStream<C> for ForestSolver<'me, C
         }
     }
 
-    fn next_answer(&mut self) -> AnswerResult<C> {
-        let answer = self.peek_answer();
+    fn next_answer(&mut self, should_continue: impl Fn() -> bool) -> AnswerResult<C> {
+        let answer = self.peek_answer(should_continue);
         self.answer.increment();
         answer
     }
 
-    fn any_future_answer(
-        &mut self,
-        test: impl FnMut(&C::InferenceNormalizedSubst) -> bool,
-    ) -> bool {
+    fn any_future_answer(&self, test: impl Fn(&C::InferenceNormalizedSubst) -> bool) -> bool {
         self.forest.any_future_answer(self.table, self.answer, test)
     }
 }

chalk-engine/src/logic.rs

@@ -204,7 +204,7 @@ impl<C: Context> Forest<C> {
     }
 
     pub(super) fn any_future_answer(
-        &mut self,
+        &self,
         table: TableIndex,
         answer: AnswerIndex,
         mut test: impl FnMut(&C::InferenceNormalizedSubst) -> bool,
@@ -214,7 +214,7 @@ impl<C: Context> Forest<C> {
             return test(C::inference_normalized_subst_from_subst(&answer.subst));
         }
 
-        self.tables[table].strands_mut().any(|strand| {
+        self.tables[table].strands().any(|strand| {
             test(C::inference_normalized_subst_from_ex_clause(
                 &strand.canonical_ex_clause,
             ))
@@ -966,7 +966,7 @@ impl<C: Context> Forest<C> {
                     continue;
                 }
 
-                let subgoal_index = strand.infer.next_subgoal_index(&strand.ex_clause);
+                let subgoal_index = C::next_subgoal_index(&strand.ex_clause);
 
                 // Get or create table for this subgoal.
                 match self.get_or_create_table_for_subgoal(
@@ -1228,7 +1228,7 @@ impl<C: Context> Forest<C> {
         goal: C::Goal,
     ) {
         let table_ref = &mut self.tables[table];
-        match infer.into_hh_goal(goal) {
+        match C::into_hh_goal(goal) {
             HhGoal::DomainGoal(domain_goal) => {
                 match context.program_clauses(&environment, &domain_goal, &mut infer) {
                     Ok(clauses) => {
@@ -1418,7 +1418,7 @@ impl<C: Context> Forest<C> {
         // untruncated literal.  Suppose that we truncate the selected
         // goal to:
         //
-        //     // Vec<Vec<T>: Sized
+        //     // Vec<Vec<T>>: Sized
         //
         // Clearly this table will have some solutions that don't
         // apply to us.  e.g., `Vec<Vec<u32>>: Sized` is a solution to

chalk-engine/src/simplify.rs

@@ -31,19 +31,19 @@ impl<C: Context> Forest<C> {
             match hh_goal {
                 HhGoal::ForAll(subgoal) => {
                     let subgoal = infer.instantiate_binders_universally(&subgoal);
-                    pending_goals.push((environment, infer.into_hh_goal(subgoal)));
+                    pending_goals.push((environment, C::into_hh_goal(subgoal)));
                 }
                 HhGoal::Exists(subgoal) => {
                     let subgoal = infer.instantiate_binders_existentially(&subgoal);
-                    pending_goals.push((environment, infer.into_hh_goal(subgoal)))
+                    pending_goals.push((environment, C::into_hh_goal(subgoal)))
                 }
                 HhGoal::Implies(wc, subgoal) => {
-                    let new_environment = infer.add_clauses(&environment, wc);
-                    pending_goals.push((new_environment, infer.into_hh_goal(subgoal)));
+                    let new_environment = C::add_clauses(&environment, wc);
+                    pending_goals.push((new_environment, C::into_hh_goal(subgoal)));
                 }
                 HhGoal::All(subgoals) => {
                     for subgoal in subgoals {
-                        pending_goals.push((environment.clone(), infer.into_hh_goal(subgoal)));
+                        pending_goals.push((environment.clone(), C::into_hh_goal(subgoal)));
                     }
                 }
                 HhGoal::Not(subgoal) => {
@@ -66,7 +66,7 @@ impl<C: Context> Forest<C> {
                         .subgoals
                         .push(Literal::Positive(C::goal_in_environment(
                             &environment,
-                            infer.into_goal(domain_goal),
+                            C::into_goal(domain_goal),
                         )));
                 }
                 HhGoal::CannotProve => {

chalk-engine/src/table.rs

@@ -70,6 +70,10 @@ impl<C: Context> Table<C> {
         self.strands.iter_mut()
     }
 
+    pub(crate) fn strands(&self) -> impl Iterator<Item = &CanonicalStrand<C>> {
+        self.strands.iter()
+    }
+
     pub(crate) fn take_strands(&mut self) -> VecDeque<CanonicalStrand<C>> {
         mem::replace(&mut self.strands, VecDeque::new())
     }

chalk-solve/src/solve.rs

@@ -135,7 +135,39 @@ impl<TF: TypeFamily> Solver<TF> {
         goal: &UCanonical<InEnvironment<Goal<TF>>>,
     ) -> Option<Solution<TF>> {
         let ops = self.forest.context().ops(program);
-        self.forest.solve(&ops, goal)
+        self.forest.solve(&ops, goal, || true)
+    }
+
+    /// Attempts to solve the given goal, which must be in canonical
+    /// form. Returns a unique solution (if one exists).  This will do
+    /// only as much work towards `goal` as it has to (and that work
+    /// is cached for future attempts). In addition, the solving of the
+    /// goal can be limited by returning `false` from `should_continue`.
+    ///
+    /// # Parameters
+    ///
+    /// - `program` -- defines the program clauses in scope.
+    ///   - **Important:** You must supply the same set of program clauses
+    ///     each time you invoke `solve`, as otherwise the cached data may be
+    ///     invalid.
+    /// - `goal` the goal to solve
+    /// - `should_continue` if `false` is returned, the no further solving
+    ///   will be done. A `Guidance(Suggested(...))` will be returned a
+    ///   `Solution`, using any answers that were generated up to that point.
+    ///
+    /// # Returns
+    ///
+    /// - `None` is the goal cannot be proven.
+    /// - `Some(solution)` if we succeeded in finding *some* answers,
+    ///   although `solution` may reflect ambiguity and unknowns.
+    pub fn solve_limited(
+        &mut self,
+        program: &dyn RustIrDatabase<TF>,
+        goal: &UCanonical<InEnvironment<Goal<TF>>>,
+        should_continue: impl std::ops::Fn() -> bool,
+    ) -> Option<Solution<TF>> {
+        let ops = self.forest.context().ops(program);
+        self.forest.solve(&ops, goal, should_continue)
     }
 
     /// Attempts to solve the given goal, which must be in canonical