2 files changed, 41 insertions, 22 deletions

diff --git a/crates/ra_hir/src/db.rs b/crates/ra_hir/src/db.rs
index 358365176..7b7974f5b 100644
--- a/crates/ra_hir/src/db.rs
+++ b/crates/ra_hir/src/db.rs
@@ -1,6 +1,5 @@
 use std::sync::Arc;
 
-use parking_lot::Mutex;
 use ra_db::{salsa, SourceDatabase};
 use ra_syntax::{ast, Parse, SmolStr, SyntaxNode};
 
@@ -147,6 +146,7 @@ pub trait DefDatabase: InternDatabase {
 }
 
 #[salsa::query_group(HirDatabaseStorage)]
+#[salsa::requires(salsa::Database)]
 pub trait HirDatabase: DefDatabase + AstDatabase {
     #[salsa::invoke(ExprScopes::expr_scopes_query)]
     fn expr_scopes(&self, def: DefWithBody) -> Arc<ExprScopes>;
@@ -187,11 +187,10 @@ pub trait HirDatabase: DefDatabase + AstDatabase {
     /// This provides the Chalk trait solver instance. Because Chalk always
     /// works from a specific crate, this query is keyed on the crate; and
     /// because Chalk does its own internal caching, the solver is wrapped in a
-    /// Mutex and the query is marked volatile, to make sure the cached state is
-    /// thrown away when input facts change.
+    /// Mutex and the query does an untracked read internally, to make sure the
+    /// cached state is thrown away when input facts change.
     #[salsa::invoke(crate::ty::traits::trait_solver_query)]
-    #[salsa::volatile]
-    fn trait_solver(&self, krate: Crate) -> Arc<Mutex<crate::ty::traits::Solver>>;
+    fn trait_solver(&self, krate: Crate) -> crate::ty::traits::TraitSolver;
 
     #[salsa::invoke(crate::ty::traits::chalk::associated_ty_data_query)]
     fn associated_ty_data(&self, id: chalk_ir::TypeId) -> Arc<chalk_rust_ir::AssociatedTyDatum>;
diff --git a/crates/ra_hir/src/ty/traits.rs b/crates/ra_hir/src/ty/traits.rs
index fde5d8a47..5406fb4a6 100644
--- a/crates/ra_hir/src/ty/traits.rs
+++ b/crates/ra_hir/src/ty/traits.rs
@@ -4,6 +4,7 @@ use std::sync::Arc;
 use chalk_ir::cast::Cast;
 use log::debug;
 use parking_lot::Mutex;
+use ra_db::salsa;
 use ra_prof::profile;
 use rustc_hash::FxHashSet;
 
@@ -14,7 +15,34 @@ use self::chalk::{from_chalk, ToChalk};
 
 pub(crate) mod chalk;
 
-pub(crate) type Solver = chalk_solve::Solver;
+#[derive(Debug, Clone)]
+pub struct TraitSolver {
+    krate: Crate,
+    inner: Arc<Mutex<chalk_solve::Solver>>,
+}
+
+/// We need eq for salsa
+impl PartialEq for TraitSolver {
+    fn eq(&self, other: &TraitSolver) -> bool {
+        Arc::ptr_eq(&self.inner, &other.inner)
+    }
+}
+
+impl Eq for TraitSolver {}
+
+impl TraitSolver {
+    fn solve(
+        &self,
+        db: &impl HirDatabase,
+        goal: &chalk_ir::UCanonical<chalk_ir::InEnvironment<chalk_ir::Goal>>,
+    ) -> Option<chalk_solve::Solution> {
+        let context = ChalkContext { db, krate: self.krate };
+        debug!("solve goal: {:?}", goal);
+        let solution = self.inner.lock().solve_with_fuel(&context, goal, Some(1000));
+        debug!("solve({:?}) => {:?}", goal, solution);
+        solution
+    }
+}
 
 /// This controls the maximum size of types Chalk considers. If we set this too
 /// high, we can run into slow edge cases; if we set it too low, Chalk won't
@@ -27,11 +55,15 @@ struct ChalkContext<'a, DB> {
     krate: Crate,
 }
 
-pub(crate) fn trait_solver_query(_db: &impl HirDatabase, _krate: Crate) -> Arc<Mutex<Solver>> {
+pub(crate) fn trait_solver_query(
+    db: &(impl HirDatabase + salsa::Database),
+    krate: Crate,
+) -> TraitSolver {
+    db.salsa_runtime().report_untracked_read();
     // krate parameter is just so we cache a unique solver per crate
     let solver_choice = chalk_solve::SolverChoice::SLG { max_size: CHALK_SOLVER_MAX_SIZE };
-    debug!("Creating new solver for crate {:?}", _krate);
-    Arc::new(Mutex::new(solver_choice.into_solver()))
+    debug!("Creating new solver for crate {:?}", krate);
+    TraitSolver { krate, inner: Arc::new(Mutex::new(solver_choice.into_solver())) }
 }
 
 /// Collects impls for the given trait in the whole dependency tree of `krate`.
@@ -54,18 +86,6 @@ pub(crate) fn impls_for_trait_query(
     impls.into_iter().collect::<Vec<_>>().into()
 }
 
-fn solve(
-    db: &impl HirDatabase,
-    krate: Crate,
-    goal: &chalk_ir::UCanonical<chalk_ir::InEnvironment<chalk_ir::Goal>>,
-) -> Option<chalk_solve::Solution> {
-    let context = ChalkContext { db, krate };
-    let solver = db.trait_solver(krate);
-    let solution = solver.lock().solve(&context, goal);
-    debug!("solve({:?}) => {:?}", goal, solution);
-    solution
-}
-
 /// A set of clauses that we assume to be true. E.g. if we are inside this function:
 /// ```rust
 /// fn foo<T: Default>(t: T) {}
@@ -127,7 +147,7 @@ pub(crate) fn trait_solve_query(
     // We currently don't deal with universes (I think / hope they're not yet
     // relevant for our use cases?)
     let u_canonical = chalk_ir::UCanonical { canonical, universes: 1 };
-    let solution = solve(db, krate, &u_canonical);
+    let solution = db.trait_solver(krate).solve(db, &u_canonical);
     solution.map(|solution| solution_from_chalk(db, solution))
 }