From 706d12047e377c77b311d7183d0328731b22d16f Mon Sep 17 00:00:00 2001
From: Jack <jack.huey@umassmed.edu>
Date: Thu, 9 Jan 2020 14:16:18 -0500
Subject: [PATCH 1/6] Move into_hh_goal into Context

---
 chalk-engine/src/context.rs  | 20 ++++++++++----------
 chalk-engine/src/logic.rs    |  2 +-
 chalk-engine/src/simplify.rs |  8 ++++----
 chalk-solve/src/solve/slg.rs | 34 +++++++++++++++++-----------------
 4 files changed, 32 insertions(+), 32 deletions(-)

diff --git a/chalk-engine/src/context.rs b/chalk-engine/src/context.rs
index bb59e2f2f48..d9752a0acbe 100644
--- a/chalk-engine/src/context.rs
+++ b/chalk-engine/src/context.rs
@@ -179,6 +179,16 @@ 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>;
 }
 
 pub trait ContextOps<C: Context>: Sized + Clone + Debug + AggregateOps<C> {
@@ -253,16 +263,6 @@ pub trait AggregateOps<C: Context> {
 /// 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;
 
diff --git a/chalk-engine/src/logic.rs b/chalk-engine/src/logic.rs
index f0ead340a46..9c6f1b94524 100644
--- a/chalk-engine/src/logic.rs
+++ b/chalk-engine/src/logic.rs
@@ -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) => {
diff --git a/chalk-engine/src/simplify.rs b/chalk-engine/src/simplify.rs
index 7d7aacb8ed0..83d35721fcf 100644
--- a/chalk-engine/src/simplify.rs
+++ b/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)));
+                    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) => {
diff --git a/chalk-solve/src/solve/slg.rs b/chalk-solve/src/solve/slg.rs
index 0259ef8b71d..8269024c51f 100644
--- a/chalk-solve/src/solve/slg.rs
+++ b/chalk-solve/src/solve/slg.rs
@@ -174,6 +174,23 @@ impl<TF: TypeFamily> context::Context for SlgContext<TF> {
             })
             .quantified
     }
+
+    fn into_hh_goal(goal: Goal<TF>) -> HhGoal<SlgContext<TF>> {
+        match goal.data().clone() {
+            GoalData::Quantified(QuantifierKind::ForAll, binders_goal) => {
+                HhGoal::ForAll(binders_goal)
+            }
+            GoalData::Quantified(QuantifierKind::Exists, binders_goal) => {
+                HhGoal::Exists(binders_goal)
+            }
+            GoalData::Implies(dg, subgoal) => HhGoal::Implies(dg, subgoal),
+            GoalData::All(goals) => HhGoal::All(goals),
+            GoalData::Not(g1) => HhGoal::Not(g1),
+            GoalData::EqGoal(EqGoal { a, b }) => HhGoal::Unify((), a, b),
+            GoalData::DomainGoal(domain_goal) => HhGoal::DomainGoal(domain_goal),
+            GoalData::CannotProve(()) => HhGoal::CannotProve,
+        }
+    }
 }
 
 impl<'me, TF: TypeFamily> context::ContextOps<SlgContext<TF>> for SlgContextOps<'me, TF> {
@@ -323,23 +340,6 @@ impl<TF: TypeFamily> context::TruncateOps<SlgContext<TF>> for TruncatingInferenc
 }
 
 impl<TF: TypeFamily> context::InferenceTable<SlgContext<TF>> for TruncatingInferenceTable<TF> {
-    fn into_hh_goal(&mut self, goal: Goal<TF>) -> HhGoal<SlgContext<TF>> {
-        match goal.data().clone() {
-            GoalData::Quantified(QuantifierKind::ForAll, binders_goal) => {
-                HhGoal::ForAll(binders_goal)
-            }
-            GoalData::Quantified(QuantifierKind::Exists, binders_goal) => {
-                HhGoal::Exists(binders_goal)
-            }
-            GoalData::Implies(dg, subgoal) => HhGoal::Implies(dg, subgoal),
-            GoalData::All(goals) => HhGoal::All(goals),
-            GoalData::Not(g1) => HhGoal::Not(g1),
-            GoalData::EqGoal(EqGoal { a, b }) => HhGoal::Unify((), a, b),
-            GoalData::DomainGoal(domain_goal) => HhGoal::DomainGoal(domain_goal),
-            GoalData::CannotProve(()) => HhGoal::CannotProve,
-        }
-    }
-
     // Used by: simplify
     fn add_clauses(
         &mut self,

From 9f5d084d3cd39feba2a6264d879ffc4e1c09d9e8 Mon Sep 17 00:00:00 2001
From: Jack <jack.huey@umassmed.edu>
Date: Thu, 9 Jan 2020 14:23:48 -0500
Subject: [PATCH 2/6] Move methods out of InferenceTable

---
 chalk-engine/src/context.rs  | 22 ++++++++---------
 chalk-engine/src/logic.rs    |  2 +-
 chalk-engine/src/simplify.rs |  4 +--
 chalk-solve/src/solve/slg.rs | 48 +++++++++++++++++-------------------
 4 files changed, 36 insertions(+), 40 deletions(-)

diff --git a/chalk-engine/src/context.rs b/chalk-engine/src/context.rs
index d9752a0acbe..7fb6b3cac56 100644
--- a/chalk-engine/src/context.rs
+++ b/chalk-engine/src/context.rs
@@ -189,6 +189,16 @@ pub trait Context: Clone + Debug {
     /// `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> {
@@ -262,17 +272,7 @@ pub trait AggregateOps<C: Context> {
 
 /// 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> {
-    // 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
diff --git a/chalk-engine/src/logic.rs b/chalk-engine/src/logic.rs
index 9c6f1b94524..ab804a5df76 100644
--- a/chalk-engine/src/logic.rs
+++ b/chalk-engine/src/logic.rs
@@ -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(
diff --git a/chalk-engine/src/simplify.rs b/chalk-engine/src/simplify.rs
index 83d35721fcf..6a674508722 100644
--- a/chalk-engine/src/simplify.rs
+++ b/chalk-engine/src/simplify.rs
@@ -38,7 +38,7 @@ impl<C: Context> Forest<C> {
                     pending_goals.push((environment, C::into_hh_goal(subgoal)))
                 }
                 HhGoal::Implies(wc, subgoal) => {
-                    let new_environment = infer.add_clauses(&environment, wc);
+                    let new_environment = C::add_clauses(&environment, wc);
                     pending_goals.push((new_environment, C::into_hh_goal(subgoal)));
                 }
                 HhGoal::All(subgoals) => {
@@ -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 => {
diff --git a/chalk-solve/src/solve/slg.rs b/chalk-solve/src/solve/slg.rs
index 8269024c51f..8ff798c096a 100644
--- a/chalk-solve/src/solve/slg.rs
+++ b/chalk-solve/src/solve/slg.rs
@@ -191,6 +191,27 @@ impl<TF: TypeFamily> context::Context for SlgContext<TF> {
             GoalData::CannotProve(()) => HhGoal::CannotProve,
         }
     }
+
+    // Used by: simplify
+    fn add_clauses(env: &Environment<TF>, clauses: Vec<ProgramClause<TF>>) -> Environment<TF> {
+        Environment::add_clauses(env, clauses)
+    }
+
+    fn into_goal(domain_goal: DomainGoal<TF>) -> Goal<TF> {
+        domain_goal.cast()
+    }
+
+    // Used by: logic
+    fn next_subgoal_index(ex_clause: &ExClause<SlgContext<TF>>) -> usize {
+        // For now, we always pick the last subgoal in the
+        // list.
+        //
+        // FIXME(rust-lang-nursery/chalk#80) -- we should be more
+        // selective. For example, we don't want to pick a
+        // negative literal that will flounder, and we don't want
+        // to pick things like `?T: Sized` if we can help it.
+        ex_clause.subgoals.len() - 1
+    }
 }
 
 impl<'me, TF: TypeFamily> context::ContextOps<SlgContext<TF>> for SlgContextOps<'me, TF> {
@@ -339,32 +360,7 @@ impl<TF: TypeFamily> context::TruncateOps<SlgContext<TF>> for TruncatingInferenc
     }
 }
 
-impl<TF: TypeFamily> context::InferenceTable<SlgContext<TF>> for TruncatingInferenceTable<TF> {
-    // Used by: simplify
-    fn add_clauses(
-        &mut self,
-        env: &Environment<TF>,
-        clauses: Vec<ProgramClause<TF>>,
-    ) -> Environment<TF> {
-        Environment::add_clauses(env, clauses)
-    }
-
-    fn into_goal(&self, domain_goal: DomainGoal<TF>) -> Goal<TF> {
-        domain_goal.cast()
-    }
-
-    // Used by: logic
-    fn next_subgoal_index(&mut self, ex_clause: &ExClause<SlgContext<TF>>) -> usize {
-        // For now, we always pick the last subgoal in the
-        // list.
-        //
-        // FIXME(rust-lang-nursery/chalk#80) -- we should be more
-        // selective. For example, we don't want to pick a
-        // negative literal that will flounder, and we don't want
-        // to pick things like `?T: Sized` if we can help it.
-        ex_clause.subgoals.len() - 1
-    }
-}
+impl<TF: TypeFamily> context::InferenceTable<SlgContext<TF>> for TruncatingInferenceTable<TF> {}
 
 impl<TF: TypeFamily> context::UnificationOps<SlgContext<TF>> for TruncatingInferenceTable<TF> {
     fn instantiate_binders_universally(&mut self, arg: &Binders<Goal<TF>>) -> Goal<TF> {

From 22a990508b049929d6e74aa69308914d1a93c001 Mon Sep 17 00:00:00 2001
From: Jack <jack.huey@umassmed.edu>
Date: Fri, 10 Jan 2020 14:40:05 -0500
Subject: [PATCH 3/6] Add a should_continue argument to AnswerStream methods

---
 chalk-engine/src/context.rs            | 12 ++++++++----
 chalk-engine/src/forest.rs             | 22 ++++++++++++----------
 chalk-engine/src/logic.rs              |  6 +++---
 chalk-engine/src/table.rs              |  4 ++++
 chalk-solve/src/solve/slg/aggregate.rs |  8 +++++---
 5 files changed, 32 insertions(+), 20 deletions(-)

diff --git a/chalk-engine/src/context.rs b/chalk-engine/src/context.rs
index 7fb6b3cac56..8b4fcd570c0 100644
--- a/chalk-engine/src/context.rs
+++ b/chalk-engine/src/context.rs
@@ -388,6 +388,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> {
@@ -419,6 +423,7 @@ 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"),
         }
     }
 }
@@ -426,14 +431,13 @@ impl<C: Context> Debug for AnswerResult<C> {
 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;
 }
diff --git a/chalk-engine/src/forest.rs b/chalk-engine/src/forest.rs
index 5e0328be23e..27a1a96ea41 100644
--- a/chalk-engine/src/forest.rs
+++ b/chalk-engine/src/forest.rs
@@ -75,7 +75,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 +87,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 +158,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 +179,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 +197,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)
     }
 }
diff --git a/chalk-engine/src/logic.rs b/chalk-engine/src/logic.rs
index ab804a5df76..d926e625979 100644
--- a/chalk-engine/src/logic.rs
+++ b/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,
             ))
@@ -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
diff --git a/chalk-engine/src/table.rs b/chalk-engine/src/table.rs
index abed8edf181..5048c1f7fa8 100644
--- a/chalk-engine/src/table.rs
+++ b/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())
     }
diff --git a/chalk-solve/src/solve/slg/aggregate.rs b/chalk-solve/src/solve/slg/aggregate.rs
index 0263e1d7694..19181c1db1e 100644
--- a/chalk-solve/src/solve/slg/aggregate.rs
+++ b/chalk-solve/src/solve/slg/aggregate.rs
@@ -20,7 +20,7 @@ impl<TF: TypeFamily> context::AggregateOps<SlgContext<TF>> for SlgContextOps<'_,
         root_goal: &UCanonical<InEnvironment<Goal<TF>>>,
         mut answers: impl context::AnswerStream<SlgContext<TF>>,
     ) -> Option<Solution<TF>> {
-        let CompleteAnswer { subst, ambiguous } = match answers.next_answer() {
+        let CompleteAnswer { subst, ambiguous } = match answers.next_answer(|| true) {
             AnswerResult::NoMoreSolutions => {
                 // No answers at all
                 return None;
@@ -30,10 +30,11 @@ impl<TF: TypeFamily> context::AggregateOps<SlgContext<TF>> for SlgContextOps<'_,
                 subst: SlgContext::identity_constrained_subst(root_goal),
                 ambiguous: true,
             },
+            AnswerResult::QuantumExceeded => unreachable!(),
         };
 
         // Exactly 1 unconditional answer?
-        if answers.peek_answer().is_no_more_solutions() && !ambiguous {
+        if answers.peek_answer(|| true).is_no_more_solutions() && !ambiguous {
             return Some(Solution::Unique(subst));
         }
 
@@ -69,7 +70,7 @@ impl<TF: TypeFamily> context::AggregateOps<SlgContext<TF>> for SlgContextOps<'_,
                 }
             }
 
-            let new_subst = match answers.next_answer() {
+            let new_subst = match answers.next_answer(|| false) {
                 AnswerResult::Answer(answer1) => answer1.subst,
                 AnswerResult::Floundered => {
                     // FIXME: this doesn't trigger for any current tests
@@ -78,6 +79,7 @@ impl<TF: TypeFamily> context::AggregateOps<SlgContext<TF>> for SlgContextOps<'_,
                 AnswerResult::NoMoreSolutions => {
                     break Guidance::Definite(subst);
                 }
+                AnswerResult::QuantumExceeded => continue,
             };
             subst = merge_into_guidance(SlgContext::canonical(root_goal), subst, &new_subst);
             num_answers += 1;

From 257980bb5bd8f94eeeaaf8488efdbe1fb1a87aa8 Mon Sep 17 00:00:00 2001
From: Jack <jack.huey@umassmed.edu>
Date: Fri, 10 Jan 2020 18:14:35 -0500
Subject: [PATCH 4/6] Add new solve_limited which allows stopping of solving

---
 chalk-engine/src/context.rs            |  8 +++++++
 chalk-engine/src/forest.rs             |  3 ++-
 chalk-solve/src/solve.rs               | 32 +++++++++++++++++++++++++-
 chalk-solve/src/solve/slg/aggregate.rs | 21 +++++++++++++----
 4 files changed, 57 insertions(+), 7 deletions(-)

diff --git a/chalk-engine/src/context.rs b/chalk-engine/src/context.rs
index 8b4fcd570c0..2111362f394 100644
--- a/chalk-engine/src/context.rs
+++ b/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> {
diff --git a/chalk-engine/src/forest.rs b/chalk-engine/src/forest.rs
index 27a1a96ea41..d37ffeffa02 100644
--- a/chalk-engine/src/forest.rs
+++ b/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
diff --git a/chalk-solve/src/solve.rs b/chalk-solve/src/solve.rs
index 116e1b4e23f..18a420660b7 100644
--- a/chalk-solve/src/solve.rs
+++ b/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
diff --git a/chalk-solve/src/solve/slg/aggregate.rs b/chalk-solve/src/solve/slg/aggregate.rs
index 19181c1db1e..bcf490a6718 100644
--- a/chalk-solve/src/solve/slg/aggregate.rs
+++ b/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;

From 8cec7ba95a8e9df11651bada11472328de86d91b Mon Sep 17 00:00:00 2001
From: Jack <jack.huey@umassmed.edu>
Date: Sun, 12 Jan 2020 14:45:03 -0500
Subject: [PATCH 5/6] Use full path for Fn in aggregate and solve since Fn is
 also a chalk_ir type

---
 chalk-solve/src/solve.rs               | 2 +-
 chalk-solve/src/solve/slg/aggregate.rs | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/chalk-solve/src/solve.rs b/chalk-solve/src/solve.rs
index 18a420660b7..c9d2c1c094e 100644
--- a/chalk-solve/src/solve.rs
+++ b/chalk-solve/src/solve.rs
@@ -162,7 +162,7 @@ impl<TF: TypeFamily> Solver<TF> {
         &mut self,
         program: &dyn RustIrDatabase<TF>,
         goal: &UCanonical<InEnvironment<Goal<TF>>>,
-        should_continue: impl Fn() -> bool,
+        should_continue: impl std::ops::Fn() -> bool,
     ) -> Option<Solution<TF>> {
         let ops = self.forest.context().ops(program);
         self.forest.solve(&ops, goal, should_continue)
diff --git a/chalk-solve/src/solve/slg/aggregate.rs b/chalk-solve/src/solve/slg/aggregate.rs
index bcf490a6718..2c50930c6c1 100644
--- a/chalk-solve/src/solve/slg/aggregate.rs
+++ b/chalk-solve/src/solve/slg/aggregate.rs
@@ -19,7 +19,7 @@ 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,
+        should_continue: impl std::ops::Fn() -> bool,
     ) -> Option<Solution<TF>> {
         let CompleteAnswer { subst, ambiguous } = match answers.next_answer(|| should_continue()) {
             AnswerResult::NoMoreSolutions => {

From 8a30731416a7297a3ba72349236266f7e013500b Mon Sep 17 00:00:00 2001
From: Jack <jack.huey@umassmed.edu>
Date: Fri, 17 Jan 2020 12:09:17 -0500
Subject: [PATCH 6/6] Add nit about guidance in solve_limited docs

---
 chalk-solve/src/solve.rs | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/chalk-solve/src/solve.rs b/chalk-solve/src/solve.rs
index c9d2c1c094e..86f4ef19d36 100644
--- a/chalk-solve/src/solve.rs
+++ b/chalk-solve/src/solve.rs
@@ -151,7 +151,9 @@ impl<TF: TypeFamily> Solver<TF> {
     ///     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
+    /// - `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
     ///