Some cleanups and add a solve_limited

chalk-engine/src/context.rs

@@ -267,6 +267,7 @@ pub trait AggregateOps<C: Context> {
         &self,
         root_goal: &C::UCanonicalGoalInEnvironment,
         answers: impl AnswerStream<C>,
+        should_continue: impl Fn() -> bool,
     ) -> Option<C::Solution>;
 }
 
@@ -415,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> {

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

chalk-solve/src/solve.rs

@@ -135,7 +135,37 @@ 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
+    ///
+    /// # 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 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

chalk-solve/src/solve/slg/aggregate.rs

@@ -19,8 +19,9 @@ impl<TF: TypeFamily> context::AggregateOps<SlgContext<TF>> for SlgContextOps<'_,
         &self,
         root_goal: &UCanonical<InEnvironment<Goal<TF>>>,
         mut answers: impl context::AnswerStream<SlgContext<TF>>,
+        should_continue: impl Fn() -> bool,
     ) -> Option<Solution<TF>> {
-        let CompleteAnswer { subst, ambiguous } = match answers.next_answer(|| true) {
+        let CompleteAnswer { subst, ambiguous } = match answers.next_answer(|| should_continue()) {
             AnswerResult::NoMoreSolutions => {
                 // No answers at all
                 return None;
@@ -30,11 +31,19 @@ impl<TF: TypeFamily> context::AggregateOps<SlgContext<TF>> for SlgContextOps<'_,
                 subst: SlgContext::identity_constrained_subst(root_goal),
                 ambiguous: true,
             },
-            AnswerResult::QuantumExceeded => unreachable!(),
+            AnswerResult::QuantumExceeded => {
+                return Some(Solution::Ambig(Guidance::Unknown));
+            }
         };
 
         // Exactly 1 unconditional answer?
-        if answers.peek_answer(|| true).is_no_more_solutions() && !ambiguous {
+        let next_answer = answers.peek_answer(|| should_continue());
+        if next_answer.is_quantum_exceeded() {
+            return Some(Solution::Ambig(Guidance::Suggested(
+                subst.map(|cs| cs.subst),
+            )));
+        }
+        if next_answer.is_no_more_solutions() && !ambiguous {
             return Some(Solution::Unique(subst));
         }
 
@@ -70,7 +79,7 @@ impl<TF: TypeFamily> context::AggregateOps<SlgContext<TF>> for SlgContextOps<'_,
                 }
             }
 
-            let new_subst = match answers.next_answer(|| false) {
+            let new_subst = match answers.next_answer(|| should_continue()) {
                 AnswerResult::Answer(answer1) => answer1.subst,
                 AnswerResult::Floundered => {
                     // FIXME: this doesn't trigger for any current tests
@@ -79,7 +88,9 @@ impl<TF: TypeFamily> context::AggregateOps<SlgContext<TF>> for SlgContextOps<'_,
                 AnswerResult::NoMoreSolutions => {
                     break Guidance::Definite(subst);
                 }
-                AnswerResult::QuantumExceeded => continue,
+                AnswerResult::QuantumExceeded => {
+                    break Guidance::Suggested(subst);
+                }
             };
             subst = merge_into_guidance(SlgContext::canonical(root_goal), subst, &new_subst);
             num_answers += 1;