Use ImplKey in abilities store public API

crates/compiler/can/src/abilities.rs

@@ -70,9 +70,6 @@ impl AbilityMemberData<Resolved> {
     }
 }
 
-/// (member, specialization type) -> specialization
-pub type ImplMap = VecMap<(Symbol, Symbol), MemberImpl>;
-
 /// Solved lambda sets for an ability member specialization. For example, if we have
 ///
 ///   Default has default : {} -[[] + a:default:1]-> a | a has Default
@@ -152,7 +149,7 @@ pub struct IAbilitiesStore<Phase: ResolvePhase> {
 
     /// Maps a tuple (member, type) specifying that `type` has an implementation of an ability
     /// member `member`, to how that implementation is defined.
-    declared_implementations: ImplMap,
+    declared_implementations: MutMap<ImplKey, MemberImpl>,
 
     /// Information about specialized ability member implementations for a type.
     specializations: MutMap<Symbol, MemberSpecializationInfo<Phase>>,
@@ -233,8 +230,7 @@ impl<Phase: ResolvePhase> IAbilitiesStore<Phase> {
             self.specialization_to_root
                 .insert(specialization_symbol, impl_key);
         }
-        self.declared_implementations
-            .insert((impl_key.ability_member, impl_key.opaque), member_impl);
+        self.declared_implementations.insert(impl_key, member_impl);
     }
 
     /// Records the implementations of an ability an opaque type declares to have.
@@ -344,12 +340,8 @@ impl<Phase: ResolvePhase> IAbilitiesStore<Phase> {
             // Add any specializations of the ability's members we know about.
             declared_implementations
                 .iter()
-                .filter(|((member, _), _)| members.contains(member))
-                .for_each(|(&(member, typ), member_impl)| {
-                    let impl_key = ImplKey {
-                        ability_member: member,
-                        opaque: typ,
-                    };
+                .filter(|(impl_key, _)| members.contains(&impl_key.ability_member))
+                .for_each(|(&impl_key, member_impl)| {
                     new.register_one_declared_impl(impl_key, *member_impl);
 
                     if let MemberImpl::Impl(spec_symbol) = member_impl {
@@ -398,26 +390,22 @@ impl IAbilitiesStore<Resolved> {
     /// the give type has an implementation of an ability member.
     pub fn iter_declared_implementations(
         &self,
-    ) -> impl Iterator<Item = ((Symbol, Symbol), &MemberImpl)> + '_ {
+    ) -> impl Iterator<Item = (ImplKey, &MemberImpl)> + '_ {
         self.declared_implementations.iter().map(|(k, v)| (*k, v))
     }
 
     /// Retrieves the declared implementation of `member` for `typ`, if it exists.
-    pub fn get_implementation(&self, member: Symbol, typ: Symbol) -> Option<&MemberImpl> {
-        self.declared_implementations.get(&(member, typ))
+    pub fn get_implementation(&self, impl_key: ImplKey) -> Option<&MemberImpl> {
+        self.declared_implementations.get(&impl_key)
     }
 
     /// Marks a declared implementation as either properly specializing, or as erroring.
     pub fn mark_implementation(
         &mut self,
-        ability_member: Symbol,
-        typ: Symbol,
+        impl_key: ImplKey,
         mark: Result<MemberSpecializationInfo<Resolved>, ()>,
     ) -> Result<(), MarkError> {
-        match self
-            .declared_implementations
-            .get_mut(&(ability_member, typ))
-        {
+        match self.declared_implementations.get_mut(&impl_key) {
             Some(member_impl) => match *member_impl {
                 MemberImpl::Impl(specialization_symbol) => {
                     debug_assert!(!self.specializations.contains_key(&specialization_symbol));
@@ -466,11 +454,6 @@ impl IAbilitiesStore<Resolved> {
 
 impl IAbilitiesStore<Pending> {
     pub fn import_implementation(&mut self, impl_key: ImplKey, resolved_impl: &ResolvedImpl) {
-        let ImplKey {
-            opaque,
-            ability_member,
-        } = impl_key;
-
         let member_impl = match resolved_impl {
             ResolvedImpl::Impl(specialization) => {
                 self.import_specialization(specialization);
@@ -480,9 +463,7 @@ impl IAbilitiesStore<Pending> {
             ResolvedImpl::Error => MemberImpl::Error,
         };
 
-        let old_declared_impl = self
-            .declared_implementations
-            .insert((ability_member, opaque), member_impl);
+        let old_declared_impl = self.declared_implementations.insert(impl_key, member_impl);
         debug_assert!(
             old_declared_impl.is_none(),
             "Replacing existing declared impl!"
@@ -539,8 +520,8 @@ impl IAbilitiesStore<Pending> {
             debug_assert!(old_root.is_none() || old_root.unwrap() == member);
         }
 
-        for ((member, typ), impl_) in declared_implementations.into_iter() {
-            let old_impl = self.declared_implementations.insert((member, typ), impl_);
+        for (impl_key, impl_) in declared_implementations.into_iter() {
+            let old_impl = self.declared_implementations.insert(impl_key, impl_);
             debug_assert!(old_impl.is_none() || old_impl.unwrap() == impl_);
         }
 

crates/compiler/can/src/traverse.rs

@@ -593,7 +593,7 @@ pub fn find_ability_member_and_owning_type_at(
         abilities_store
             .iter_declared_implementations()
             .find(|(_, member_impl)| matches!(member_impl, MemberImpl::Impl(sym) if *sym == symbol))
-            .map(|(spec, _)| spec.1)
+            .map(|(impl_key, _)| impl_key.opaque)
     }
 }
 

crates/compiler/solve/src/ability.rs

@@ -394,7 +394,10 @@ impl ObligationCache<'_> {
         for &member in members_of_ability {
             if self
                 .abilities_store
-                .get_implementation(member, opaque)
+                .get_implementation(roc_can::abilities::ImplKey {
+                    opaque,
+                    ability_member: ability,
+                })
                 .is_none()
             {
                 let root_data = self.abilities_store.member_def(member).unwrap();
@@ -671,7 +674,12 @@ pub fn resolve_ability_specialization(
 
     let resolved = match obligated {
         Obligated::Opaque(symbol) => {
-            match abilities_store.get_implementation(ability_member, symbol)? {
+            let impl_key = roc_can::abilities::ImplKey {
+                opaque: symbol,
+                ability_member,
+            };
+
+            match abilities_store.get_implementation(impl_key)? {
                 roc_types::types::MemberImpl::Impl(spec_symbol) => {
                     Resolved::Specialization(*spec_symbol)
                 }

crates/compiler/solve/src/module.rs

@@ -1,5 +1,5 @@
 use crate::solve::{self, Aliases};
-use roc_can::abilities::{AbilitiesStore, ImplKey, ResolvedImpl};
+use roc_can::abilities::{AbilitiesStore, ResolvedImpl};
 use roc_can::constraint::{Constraint as ConstraintSoa, Constraints};
 use roc_can::expr::PendingDerives;
 use roc_can::module::{ExposedByModule, ResolvedImplementations, RigidVariables};
@@ -190,16 +190,14 @@ pub fn extract_module_owned_implementations(
 ) -> ResolvedImplementations {
     abilities_store
         .iter_declared_implementations()
-        .filter_map(|((member, typ), member_impl)| {
+        .filter_map(|(impl_key, member_impl)| {
             // This module solved this specialization if either the member or the type comes from the
             // module.
-            if member.module_id() != module_id && typ.module_id() != module_id {
+            if impl_key.ability_member.module_id() != module_id
+                && impl_key.opaque.module_id() != module_id
+            {
                 return None;
             }
-            let impl_key = ImplKey {
-                opaque: typ,
-                ability_member: member,
-            };
 
             let resolved_impl = match member_impl {
                 MemberImpl::Impl(impl_symbol) => {

crates/compiler/solve/src/solve.rs

@@ -1855,7 +1855,7 @@ fn check_ability_specialization(
         };
 
         abilities_store
-            .mark_implementation(impl_key.ability_member, impl_key.opaque, resolved_mark)
+            .mark_implementation(impl_key, resolved_mark)
             .expect("marked as a custom implementation, but not recorded as such");
     }
 }
@@ -2310,8 +2310,12 @@ fn get_specialization_lambda_set_ambient_function<P: Phase>(
             let opaque_home = opaque.module_id();
             let external_specialized_lset =
                 phase.with_module_abilities_store(opaque_home, |abilities_store| {
+            let impl_key = roc_can::abilities::ImplKey {
+                opaque,
+                ability_member,
+            };
                     let opt_specialization =
-                        abilities_store.get_implementation(ability_member, opaque);
+                        abilities_store.get_implementation(impl_key);
                     match (P::IS_LATE, opt_specialization) {
                         (false, None) => {
                             // doesn't specialize, we'll have reported an error for this

crates/compiler/solve/tests/solve_expr.rs

@@ -12,7 +12,10 @@ mod solve_expr {
     use crate::helpers::with_larger_debug_stack;
     use lazy_static::lazy_static;
     use regex::Regex;
-    use roc_can::traverse::{find_ability_member_and_owning_type_at, find_type_at};
+    use roc_can::{
+        abilities::ImplKey,
+        traverse::{find_ability_member_and_owning_type_at, find_type_at},
+    };
     use roc_load::LoadedModule;
     use roc_module::symbol::{Interns, ModuleId};
     use roc_problem::can::Problem;
@@ -367,12 +370,12 @@ mod solve_expr {
         }
 
         let known_specializations = abilities_store.iter_declared_implementations().filter_map(
-            |((member, typ), member_impl)| match member_impl {
+            |(impl_key, member_impl)| match member_impl {
                 MemberImpl::Impl(impl_symbol) => {
                     let specialization = abilities_store.specialization_info(*impl_symbol).expect(
                         "declared implementations should be resolved conclusively after solving",
                     );
-                    Some((member, typ, specialization.clone()))
+                    Some((impl_key, specialization.clone()))
                 }
                 MemberImpl::Derived | MemberImpl::Error => None,
             },
@@ -381,13 +384,17 @@ mod solve_expr {
         use std::collections::HashSet;
         let pretty_specializations = known_specializations
             .into_iter()
-            .map(|(member, typ, _)| {
-                let member_data = abilities_store.member_def(member).unwrap();
-                let member_str = member.as_str(&interns);
+            .map(|(impl_key, _)| {
+                let ImplKey {
+                    opaque,
+                    ability_member,
+                } = impl_key;
+                let member_data = abilities_store.member_def(ability_member).unwrap();
+                let member_str = ability_member.as_str(&interns);
                 let ability_str = member_data.parent_ability.as_str(&interns);
                 (
                     format!("{}:{}", ability_str, member_str),
-                    typ.as_str(&interns),
+                    opaque.as_str(&interns),
                 )
             })
             .collect::<HashSet<_>>();