5 files changed, 44 insertions, 55 deletions

diff --git a/crates/hir_ty/src/lower.rs b/crates/hir_ty/src/lower.rs
index 5fa83567b..d84ec9b7a 100644
--- a/crates/hir_ty/src/lower.rs
+++ b/crates/hir_ty/src/lower.rs
@@ -8,7 +8,7 @@
 use std::{iter, sync::Arc};
 
 use base_db::CrateId;
-use chalk_ir::Mutability;
+use chalk_ir::{cast::Cast, Mutability};
 use hir_def::{
     adt::StructKind,
     builtin_type::BuiltinType,
@@ -27,6 +27,7 @@ use stdx::impl_from;
 
 use crate::{
     db::HirDatabase,
+    traits::chalk::{Interner, ToChalk},
     utils::{
         all_super_trait_refs, associated_type_by_name_including_super_traits, generics,
         make_mut_slice, variant_data,
@@ -914,10 +915,21 @@ impl TraitEnvironment {
     pub fn lower(db: &dyn HirDatabase, resolver: &Resolver) -> Arc<TraitEnvironment> {
         let ctx = TyLoweringContext::new(db, &resolver)
             .with_type_param_mode(TypeParamLoweringMode::Placeholder);
-        let mut predicates = resolver
-            .where_predicates_in_scope()
-            .flat_map(|pred| GenericPredicate::from_where_predicate(&ctx, pred))
-            .collect::<Vec<_>>();
+        let mut traits_in_scope = Vec::new();
+        let mut clauses = Vec::new();
+        for pred in resolver.where_predicates_in_scope() {
+            for pred in GenericPredicate::from_where_predicate(&ctx, pred) {
+                if pred.is_error() {
+                    continue;
+                }
+                if let GenericPredicate::Implemented(tr) = &pred {
+                    traits_in_scope.push((tr.self_ty().clone(), tr.trait_));
+                }
+                let program_clause: chalk_ir::ProgramClause<Interner> =
+                    pred.clone().to_chalk(db).cast(&Interner);
+                clauses.push(program_clause.into_from_env_clause(&Interner));
+            }
+        }
 
         if let Some(def) = resolver.generic_def() {
             let container: Option<AssocContainerId> = match def {
@@ -938,12 +950,15 @@ impl TraitEnvironment {
                 let substs = Substs::type_params(db, trait_id);
                 let trait_ref = TraitRef { trait_: trait_id, substs };
                 let pred = GenericPredicate::Implemented(trait_ref);
-
-                predicates.push(pred);
+                let program_clause: chalk_ir::ProgramClause<Interner> =
+                    pred.clone().to_chalk(db).cast(&Interner);
+                clauses.push(program_clause.into_from_env_clause(&Interner));
             }
         }
 
-        Arc::new(TraitEnvironment { predicates })
+        let env = chalk_ir::Environment::new(&Interner).add_clauses(&Interner, clauses);
+
+        Arc::new(TraitEnvironment { traits_from_clauses: traits_in_scope, env })
     }
 }
 
diff --git a/crates/hir_ty/src/method_resolution.rs b/crates/hir_ty/src/method_resolution.rs
index ccc12c075..d57c6de70 100644
--- a/crates/hir_ty/src/method_resolution.rs
+++ b/crates/hir_ty/src/method_resolution.rs
@@ -528,8 +528,7 @@ fn iterate_trait_method_candidates(
         self_ty.value.dyn_trait().into_iter().flat_map(|t| all_super_traits(db.upcast(), t));
     let env_traits = if let Ty::Placeholder(_) = self_ty.value {
         // if we have `T: Trait` in the param env, the trait doesn't need to be in scope
-        env.trait_predicates_for_self_ty(&self_ty.value)
-            .map(|tr| tr.trait_)
+        env.traits_in_scope_from_clauses(&self_ty.value)
             .flat_map(|t| all_super_traits(db.upcast(), t))
             .collect()
     } else {
diff --git a/crates/hir_ty/src/traits.rs b/crates/hir_ty/src/traits.rs
index e4cdb6d53..27f350f70 100644
--- a/crates/hir_ty/src/traits.rs
+++ b/crates/hir_ty/src/traits.rs
@@ -38,22 +38,25 @@ fn create_chalk_solver() -> chalk_recursive::RecursiveSolver<Interner> {
 /// fn foo<T: Default>(t: T) {}
 /// ```
 /// we assume that `T: Default`.
-#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct TraitEnvironment {
-    pub predicates: Vec<GenericPredicate>,
+    // When we're using Chalk's Ty we can make this a BTreeMap since it's Ord,
+    // but for now it's too annoying...
+    pub(crate) traits_from_clauses: Vec<(Ty, TraitId)>,
+    pub(crate) env: chalk_ir::Environment<Interner>,
 }
 
 impl TraitEnvironment {
-    /// Returns trait refs with the given self type which are supposed to hold
-    /// in this trait env. E.g. if we are in `foo<T: SomeTrait>()`, this will
-    /// find that `T: SomeTrait` if we call it for `T`.
-    pub(crate) fn trait_predicates_for_self_ty<'a>(
+    pub(crate) fn traits_in_scope_from_clauses<'a>(
         &'a self,
         ty: &'a Ty,
-    ) -> impl Iterator<Item = &'a TraitRef> + 'a {
-        self.predicates.iter().filter_map(move |pred| match pred {
-            GenericPredicate::Implemented(tr) if tr.self_ty() == ty => Some(tr),
-            _ => None,
+    ) -> impl Iterator<Item = TraitId> + 'a {
+        self.traits_from_clauses.iter().filter_map(move |(self_ty, trait_id)| {
+            if self_ty == ty {
+                Some(*trait_id)
+            } else {
+                None
+            }
         })
     }
 }
diff --git a/crates/hir_ty/src/traits/chalk.rs b/crates/hir_ty/src/traits/chalk.rs
index 565672b6b..4e7d43f7e 100644
--- a/crates/hir_ty/src/traits/chalk.rs
+++ b/crates/hir_ty/src/traits/chalk.rs
@@ -33,13 +33,13 @@ pub(super) mod tls;
 mod interner;
 mod mapping;
 
-pub(super) trait ToChalk {
+pub(crate) trait ToChalk {
     type Chalk;
     fn to_chalk(self, db: &dyn HirDatabase) -> Self::Chalk;
     fn from_chalk(db: &dyn HirDatabase, chalk: Self::Chalk) -> Self;
 }
 
-pub(super) fn from_chalk<T, ChalkT>(db: &dyn HirDatabase, chalk: ChalkT) -> T
+pub(crate) fn from_chalk<T, ChalkT>(db: &dyn HirDatabase, chalk: ChalkT) -> T
 where
     T: ToChalk<Chalk = ChalkT>,
 {
diff --git a/crates/hir_ty/src/traits/chalk/mapping.rs b/crates/hir_ty/src/traits/chalk/mapping.rs
index 3a08b67e9..b0415e8b0 100644
--- a/crates/hir_ty/src/traits/chalk/mapping.rs
+++ b/crates/hir_ty/src/traits/chalk/mapping.rs
@@ -17,7 +17,7 @@ use crate::{
     primitive::UintTy,
     traits::{Canonical, Obligation},
     AliasTy, CallableDefId, FnPointer, FnSig, GenericPredicate, InEnvironment, OpaqueTy,
-    OpaqueTyId, ProjectionPredicate, ProjectionTy, Scalar, Substs, TraitEnvironment, TraitRef, Ty,
+    OpaqueTyId, ProjectionPredicate, ProjectionTy, Scalar, Substs, TraitRef, Ty,
 };
 
 use super::interner::*;
@@ -536,31 +536,6 @@ where
     }
 }
 
-impl ToChalk for Arc<TraitEnvironment> {
-    type Chalk = chalk_ir::Environment<Interner>;
-
-    fn to_chalk(self, db: &dyn HirDatabase) -> chalk_ir::Environment<Interner> {
-        let mut clauses = Vec::new();
-        for pred in &self.predicates {
-            if pred.is_error() {
-                // for env, we just ignore errors
-                continue;
-            }
-            let program_clause: chalk_ir::ProgramClause<Interner> =
-                pred.clone().to_chalk(db).cast(&Interner);
-            clauses.push(program_clause.into_from_env_clause(&Interner));
-        }
-        chalk_ir::Environment::new(&Interner).add_clauses(&Interner, clauses)
-    }
-
-    fn from_chalk(
-        _db: &dyn HirDatabase,
-        _env: chalk_ir::Environment<Interner>,
-    ) -> Arc<TraitEnvironment> {
-        unimplemented!()
-    }
-}
-
 impl<T: ToChalk> ToChalk for InEnvironment<T>
 where
     T::Chalk: chalk_ir::interner::HasInterner<Interner = Interner>,
@@ -569,19 +544,16 @@ where
 
     fn to_chalk(self, db: &dyn HirDatabase) -> chalk_ir::InEnvironment<T::Chalk> {
         chalk_ir::InEnvironment {
-            environment: self.environment.to_chalk(db),
+            environment: self.environment.env.clone(),
             goal: self.value.to_chalk(db),
         }
     }
 
     fn from_chalk(
-        db: &dyn HirDatabase,
-        in_env: chalk_ir::InEnvironment<T::Chalk>,
+        _db: &dyn HirDatabase,
+        _in_env: chalk_ir::InEnvironment<T::Chalk>,
     ) -> InEnvironment<T> {
-        InEnvironment {
-            environment: from_chalk(db, in_env.environment),
-            value: from_chalk(db, in_env.goal),
-        }
+        unimplemented!()
     }
 }