Make it possible to cache the result of env elaboration

`program_clauses_for_env` can take a huge amount of time -- rust-analyzer is
unusably slow on the Chalk codebase itself, spending most of that time in env
elaboration because one `I: Interner` results in ~300 clauses, which gets
repeated tens of times per query.

This adds a method to `RustIrDatabase` that should basically just be a
pass-through for `program_clauses_for_env`, but which makes it possible to cache
the results in Salsa.

The effect of this is a ~20% speedup of analysis of rust-analyzer on itself, and
turning analysis of the Chalk codebase from seemingly forever to ~1 minute.

Author: flodiebold

chalk-integration/src/db.rs

@@ -8,6 +8,7 @@ use chalk_ir::tls;
 use chalk_ir::AssocTypeId;
 use chalk_ir::Canonical;
 use chalk_ir::ConstrainedSubst;
+use chalk_ir::Environment;
 use chalk_ir::Goal;
 use chalk_ir::ImplId;
 use chalk_ir::InEnvironment;
@@ -17,7 +18,7 @@ use chalk_ir::ProgramClause;
 use chalk_ir::StructId;
 use chalk_ir::TraitId;
 use chalk_ir::TypeName;
-use chalk_ir::UCanonical;
+use chalk_ir::{ProgramClauses, UCanonical};
 use chalk_rust_ir::AssociatedTyDatum;
 use chalk_rust_ir::AssociatedTyValue;
 use chalk_rust_ir::AssociatedTyValueId;
@@ -150,6 +151,13 @@ impl RustIrDatabase<ChalkIr> for ChalkDatabase {
             .well_known_trait_id(well_known_trait)
     }
 
+    fn program_clauses_for_env(
+        &self,
+        environment: &Environment<ChalkIr>,
+    ) -> ProgramClauses<ChalkIr> {
+        chalk_solve::program_clauses_for_env(self, environment)
+    }
+
     fn interner(&self) -> &ChalkIr {
         &ChalkIr
     }

chalk-integration/src/program.rs

@@ -390,6 +390,13 @@ impl RustIrDatabase<ChalkIr> for Program {
         self.well_known_traits.get(&well_known_trait).map(|x| *x)
     }
 
+    fn program_clauses_for_env(
+        &self,
+        environment: &chalk_ir::Environment<ChalkIr>,
+    ) -> ProgramClauses<ChalkIr> {
+        chalk_solve::program_clauses_for_env(self, environment)
+    }
+
     fn interner(&self) -> &ChalkIr {
         &ChalkIr
     }

chalk-solve/src/clauses.rs

@@ -127,8 +127,13 @@ pub(crate) fn program_clauses_for_goal<'db, I: Interner>(
     let mut vec = vec![];
     vec.extend(db.custom_clauses());
     program_clauses_that_could_match(db, environment, goal, &mut vec)?;
-    program_clauses_for_env(db, environment, &mut vec);
     vec.retain(|c| c.could_match(interner, goal));
+    vec.extend(
+        db.program_clauses_for_env(environment)
+            .iter(interner)
+            .filter(|c| (*c).could_match(interner, goal))
+            .cloned(),
+    );
 
     debug!("vec = {:#?}", vec);
 
@@ -468,20 +473,16 @@ fn match_struct<I: Interner>(builder: &mut ClauseBuilder<'_, I>, struct_id: Stru
         .to_program_clauses(builder)
 }
 
-fn program_clauses_for_env<'db, I: Interner>(
+pub fn program_clauses_for_env<'db, I: Interner>(
     db: &'db dyn RustIrDatabase<I>,
     environment: &Environment<I>,
-    clauses: &mut Vec<ProgramClause<I>>,
-) {
-    clauses.extend(environment.clauses.iter(db.interner()).cloned());
-
-    let mut last_round = FxHashSet::default();
-    elaborate_env_clauses(
-        db,
-        environment.clauses.as_slice(db.interner()),
-        &mut last_round,
-    );
-
+) -> ProgramClauses<I> {
+    let mut last_round = environment
+        .clauses
+        .as_slice(db.interner())
+        .iter()
+        .cloned()
+        .collect::<FxHashSet<_>>();
     let mut closure = last_round.clone();
     let mut next_round = FxHashSet::default();
     while !last_round.is_empty() {
@@ -493,5 +494,5 @@ fn program_clauses_for_env<'db, I: Interner>(
         );
     }
 
-    clauses.extend(closure.drain())
+    ProgramClauses::from(db.interner(), closure)
 }

chalk-solve/src/lib.rs

@@ -87,9 +87,15 @@ pub trait RustIrDatabase<I: Interner>: Debug {
     /// Returns id of a trait lang item, if found
     fn well_known_trait_id(&self, well_known_trait: WellKnownTrait) -> Option<TraitId<I>>;
 
+    /// Calculates program clauses from an env. This is intended to call the
+    /// `program_clauses_for_env` function and then possibly cache the clauses.
+    fn program_clauses_for_env(&self, environment: &Environment<I>) -> ProgramClauses<I>;
+
     fn interner(&self) -> &I;
 }
 
+pub use clauses::program_clauses_for_env;
+
 pub use solve::Guidance;
 pub use solve::Solution;
 pub use solve::Solver;