Merge pull request #2857 from rtfeldman/abilities-mono

Codegen for abilities

compiler/can/src/abilities.rs

@@ -16,6 +16,7 @@ pub struct MemberVariables {
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct AbilityMemberData {
     pub parent_ability: Symbol,
+    pub signature_var: Variable,
     pub signature: Type,
     pub variables: MemberVariables,
     pub region: Region,
@@ -60,15 +61,16 @@ impl AbilitiesStore {
     pub fn register_ability(
         &mut self,
         ability: Symbol,
-        members: Vec<(Symbol, Region, Type, MemberVariables)>,
+        members: Vec<(Symbol, Region, Variable, Type, MemberVariables)>,
     ) {
         let mut members_vec = Vec::with_capacity(members.len());
-        for (member, region, signature, variables) in members.into_iter() {
+        for (member, region, signature_var, signature, variables) in members.into_iter() {
             members_vec.push(member);
             let old_member = self.ability_members.insert(
                 member,
                 AbilityMemberData {
                     parent_ability: ability,
+                    signature_var,
                     signature,
                     region,
                     variables,

compiler/can/src/def.rs

@@ -542,7 +542,13 @@ pub fn canonicalize_defs<'a>(
                 flex_vars: iv.collect_flex(),
             };
 
-            can_members.push((member_sym, name_region, member_annot.typ, variables));
+            can_members.push((
+                member_sym,
+                name_region,
+                var_store.fresh(),
+                member_annot.typ,
+                variables,
+            ));
         }
 
         // Store what symbols a type must define implementations for to have this ability.

compiler/constrain/src/module.rs

@@ -2,12 +2,14 @@ use roc_builtins::std::StdLib;
 use roc_can::abilities::AbilitiesStore;
 use roc_can::constraint::{Constraint, Constraints};
 use roc_can::def::Declaration;
+use roc_can::expected::Expected;
 use roc_collections::all::MutMap;
 use roc_error_macros::internal_error;
 use roc_module::symbol::{ModuleId, Symbol};
-use roc_region::all::Loc;
+use roc_region::all::{Loc, Region};
 use roc_types::solved_types::{FreeVars, SolvedType};
 use roc_types::subs::{VarStore, Variable};
+use roc_types::types::{Category, Type};
 
 /// The types of all exposed values/functions of a collection of modules
 #[derive(Clone, Debug, Default)]
@@ -105,27 +107,53 @@ pub fn constrain_module(
     declarations: &[Declaration],
     home: ModuleId,
 ) -> Constraint {
-    let mut constraint = crate::expr::constrain_decls(constraints, home, declarations);
+    let constraint = crate::expr::constrain_decls(constraints, home, declarations);
 
+    let constraint = frontload_ability_constraints(constraints, abilities_store, constraint);
+
+    // The module constraint should always save the environment at the end.
+    debug_assert!(constraints.contains_save_the_environment(&constraint));
+
+    constraint
+}
+
+pub fn frontload_ability_constraints(
+    constraints: &mut Constraints,
+    abilities_store: &AbilitiesStore,
+    mut constraint: Constraint,
+) -> Constraint {
     for (member_name, member_data) in abilities_store.root_ability_members().iter() {
+        // 1. Attach the type of member signature to the reserved signature_var. This is
+        //    infallible.
+        let unify_with_signature_var = constraints.equal_types_var(
+            member_data.signature_var,
+            Expected::NoExpectation(member_data.signature.clone()),
+            Category::Storage(std::file!(), std::column!()),
+            Region::zero(),
+        );
+
+        // 2. Store the member signature on the member symbol. This makes sure we generalize it on
+        //    the toplevel, as appropriate.
         let vars = &member_data.variables;
         let rigids = (vars.rigid_vars.iter())
             // For our purposes, in the let constraint, able vars are treated like rigids.
             .chain(vars.able_vars.iter())
             .copied();
         let flex = vars.flex_vars.iter().copied();
-        constraint = constraints.let_constraint(
+
+        let let_constr = constraints.let_constraint(
             rigids,
             flex,
-            [(*member_name, Loc::at_zero(member_data.signature.clone()))],
+            [(
+                *member_name,
+                Loc::at_zero(Type::Variable(member_data.signature_var)),
+            )],
             Constraint::True,
             constraint,
         );
+
+        constraint = constraints.and_constraint([unify_with_signature_var, let_constr]);
     }
-
-    // The module constraint should always save the environment at the end.
-    debug_assert!(constraints.contains_save_the_environment(&constraint));
-
     constraint
 }
 

compiler/load_internal/src/file.rs

@@ -271,6 +271,7 @@ fn start_phase<'a>(
                     solved_subs,
                     decls,
                     ident_ids,
+                    abilities_store,
                 } = typechecked;
 
                 let mut imported_module_thunks = bumpalo::collections::Vec::new_in(arena);
@@ -294,6 +295,7 @@ fn start_phase<'a>(
                     decls,
                     ident_ids,
                     exposed_to_host: state.exposed_to_host.clone(),
+                    abilities_store,
                 }
             }
             Phase::MakeSpecializations => {
@@ -316,6 +318,7 @@ fn start_phase<'a>(
                     procs_base,
                     layout_cache,
                     module_timing,
+                    abilities_store,
                 } = found_specializations;
 
                 BuildTask::MakeSpecializations {
@@ -326,6 +329,7 @@ fn start_phase<'a>(
                     layout_cache,
                     specializations_we_must_make,
                     module_timing,
+                    abilities_store,
                 }
             }
         }
@@ -419,6 +423,7 @@ pub struct TypeCheckedModule<'a> {
     pub solved_subs: Solved<Subs>,
     pub decls: Vec<Declaration>,
     pub ident_ids: IdentIds,
+    pub abilities_store: AbilitiesStore,
 }
 
 #[derive(Debug)]
@@ -429,6 +434,7 @@ struct FoundSpecializationsModule<'a> {
     procs_base: ProcsBase<'a>,
     subs: Subs,
     module_timing: ModuleTiming,
+    abilities_store: AbilitiesStore,
 }
 
 #[derive(Debug)]
@@ -529,6 +535,7 @@ enum Msg<'a> {
         problems: Vec<roc_mono::ir::MonoProblem>,
         solved_subs: Solved<Subs>,
         module_timing: ModuleTiming,
+        abilities_store: AbilitiesStore,
     },
     MadeSpecializations {
         module_id: ModuleId,
@@ -767,6 +774,7 @@ enum BuildTask<'a> {
         ident_ids: IdentIds,
         decls: Vec<Declaration>,
         exposed_to_host: ExposedToHost,
+        abilities_store: AbilitiesStore,
     },
     MakeSpecializations {
         module_id: ModuleId,
@@ -776,6 +784,7 @@ enum BuildTask<'a> {
         layout_cache: LayoutCache<'a>,
         specializations_we_must_make: Vec<ExternalSpecializations>,
         module_timing: ModuleTiming,
+        abilities_store: AbilitiesStore,
     },
 }
 
@@ -1880,6 +1889,7 @@ fn update<'a>(
                         solved_subs,
                         decls,
                         ident_ids,
+                        abilities_store,
                     };
 
                     state
@@ -1904,6 +1914,7 @@ fn update<'a>(
             layout_cache,
             problems,
             module_timing,
+            abilities_store,
         } => {
             log!("found specializations for {:?}", module_id);
 
@@ -1925,6 +1936,7 @@ fn update<'a>(
                 procs_base,
                 subs,
                 module_timing,
+                abilities_store,
             };
 
             state
@@ -3642,6 +3654,7 @@ fn make_specializations<'a>(
     specializations_we_must_make: Vec<ExternalSpecializations>,
     mut module_timing: ModuleTiming,
     target_info: TargetInfo,
+    abilities_store: AbilitiesStore,
 ) -> Msg<'a> {
     let make_specializations_start = SystemTime::now();
     let mut mono_problems = Vec::new();
@@ -3657,6 +3670,7 @@ fn make_specializations<'a>(
         update_mode_ids: &mut update_mode_ids,
         // call_specialization_counter=0 is reserved
         call_specialization_counter: 1,
+        abilities_store: &abilities_store,
     };
 
     let mut procs = Procs::new_in(arena);
@@ -3727,6 +3741,7 @@ fn build_pending_specializations<'a>(
     target_info: TargetInfo,
     // TODO remove
     exposed_to_host: ExposedToHost,
+    abilities_store: AbilitiesStore,
 ) -> Msg<'a> {
     let find_specializations_start = SystemTime::now();
 
@@ -3753,6 +3768,7 @@ fn build_pending_specializations<'a>(
         update_mode_ids: &mut update_mode_ids,
         // call_specialization_counter=0 is reserved
         call_specialization_counter: 1,
+        abilities_store: &abilities_store,
     };
 
     // Add modules' decls to Procs
@@ -3806,6 +3822,7 @@ fn build_pending_specializations<'a>(
         procs_base,
         problems,
         module_timing,
+        abilities_store,
     }
 }
 
@@ -3823,7 +3840,11 @@ fn add_def_to_module<'a>(
     use roc_can::pattern::Pattern::*;
 
     match def.loc_pattern.value {
-        Identifier(symbol) => {
+        Identifier(symbol)
+        | AbilityMemberSpecialization {
+            ident: symbol,
+            specializes: _,
+        } => {
             let is_host_exposed = exposed_to_host.contains_key(&symbol);
 
             match def.loc_expr.value {
@@ -4026,6 +4047,7 @@ fn run_task<'a>(
             solved_subs,
             imported_module_thunks,
             exposed_to_host,
+            abilities_store,
         } => Ok(build_pending_specializations(
             arena,
             solved_subs,
@@ -4037,6 +4059,7 @@ fn run_task<'a>(
             layout_cache,
             target_info,
             exposed_to_host,
+            abilities_store,
         )),
         MakeSpecializations {
             module_id,
@@ -4046,6 +4069,7 @@ fn run_task<'a>(
             layout_cache,
             specializations_we_must_make,
             module_timing,
+            abilities_store,
         } => Ok(make_specializations(
             arena,
             module_id,
@@ -4056,6 +4080,7 @@ fn run_task<'a>(
             specializations_we_must_make,
             module_timing,
             target_info,
+            abilities_store,
         )),
     }?;
 

compiler/mono/src/ir.rs

@@ -7,6 +7,7 @@ use crate::layout::{
 use bumpalo::collections::Vec;
 use bumpalo::Bump;
 use roc_builtins::bitcode::{FloatWidth, IntWidth};
+use roc_can::abilities::AbilitiesStore;
 use roc_can::expr::{ClosureData, IntValue};
 use roc_collections::all::{default_hasher, BumpMap, BumpMapDefault, MutMap};
 use roc_exhaustive::{Ctor, Guard, RenderAs, TagId};
@@ -251,11 +252,22 @@ impl<'a> PartialProc<'a> {
 }
 
 #[derive(Clone, Debug)]
-enum PartialExprLink {
-    Aliases(Symbol),
+enum PolymorphicExpr {
+    /// A root ability member, which must be specialized at a call site, for example
+    /// "hash" which must be specialized to an exact symbol implementing "hash" for a type.
+    AbilityMember(Symbol),
+    /// A polymorphic expression we inline at the usage site.
     Expr(roc_can::expr::Expr, Variable),
 }
 
+#[derive(Clone, Debug)]
+enum PartialExprLink {
+    /// The root polymorphic expression
+    Sink(PolymorphicExpr),
+    /// A hop in a partial expression alias chain
+    Aliases(Symbol),
+}
+
 /// A table of symbols to polymorphic expressions. For example, in the program
 ///
 ///   n = 1
@@ -281,8 +293,8 @@ impl PartialExprs {
         Self(BumpMap::new_in(arena))
     }
 
-    fn insert(&mut self, symbol: Symbol, expr: roc_can::expr::Expr, expr_var: Variable) {
-        self.0.insert(symbol, PartialExprLink::Expr(expr, expr_var));
+    fn insert(&mut self, symbol: Symbol, expr: PolymorphicExpr) {
+        self.0.insert(symbol, PartialExprLink::Sink(expr));
     }
 
     fn insert_alias(&mut self, symbol: Symbol, aliases: Symbol) {
@@ -293,7 +305,7 @@ impl PartialExprs {
         self.0.contains_key(&symbol)
     }
 
-    fn get(&mut self, mut symbol: Symbol) -> Option<(&roc_can::expr::Expr, Variable)> {
+    fn get(&mut self, mut symbol: Symbol) -> Option<&PolymorphicExpr> {
         // In practice the alias chain is very short
         loop {
             match self.0.get(&symbol) {
@@ -303,8 +315,8 @@ impl PartialExprs {
                 Some(&PartialExprLink::Aliases(real_symbol)) => {
                     symbol = real_symbol;
                 }
-                Some(PartialExprLink::Expr(expr, var)) => {
-                    return Some((expr, *var));
+                Some(PartialExprLink::Sink(expr)) => {
+                    return Some(expr);
                 }
             }
         }
@@ -1119,6 +1131,7 @@ pub struct Env<'a, 'i> {
     pub target_info: TargetInfo,
     pub update_mode_ids: &'i mut UpdateModeIds,
     pub call_specialization_counter: u32,
+    pub abilities_store: &'i AbilitiesStore,
 }
 
 impl<'a, 'i> Env<'a, 'i> {
@@ -4218,10 +4231,14 @@ pub fn with_hole<'a>(
                 // a proc in this module, or an imported symbol
                 procs.partial_procs.contains_key(key)
                     || (env.is_imported_symbol(key) && !procs.is_imported_module_thunk(key))
+                    || env.abilities_store.is_ability_member_name(key)
             };
 
             match loc_expr.value {
                 roc_can::expr::Expr::Var(proc_name) if is_known(proc_name) => {
+                    // This might be an ability member - if so, use the appropriate specialization.
+                    let proc_name = get_specialization(env, fn_var, proc_name).unwrap_or(proc_name);
+
                     // a call by a known name
                     call_by_name(
                         env,
@@ -4330,8 +4347,30 @@ pub fn with_hole<'a>(
                                 unreachable!("calling a non-closure layout")
                             }
                         },
-                        UnspecializedExpr(symbol) => match full_layout {
-                            RawFunctionLayout::Function(arg_layouts, lambda_set, ret_layout) => {
+                        UnspecializedExpr(symbol) => match procs.partial_exprs.get(symbol).unwrap()
+                        {
+                            &PolymorphicExpr::AbilityMember(member) => {
+                                let proc_name = get_specialization(env, fn_var, member).expect("Recorded as an ability member, but it doesn't have a specialization");
+
+                                // a call by a known name
+                                return call_by_name(
+                                    env,
+                                    procs,
+                                    fn_var,
+                                    proc_name,
+                                    loc_args,
+                                    layout_cache,
+                                    assigned,
+                                    hole,
+                                );
+                            }
+                            PolymorphicExpr::Expr(lambda_expr, lambda_expr_var) => {
+                                match full_layout {
+                                    RawFunctionLayout::Function(
+                                        arg_layouts,
+                                        lambda_set,
+                                        ret_layout,
+                                    ) => {
                                         let closure_data_symbol = env.unique_symbol();
 
                                         result = match_on_lambda_set(
@@ -4345,15 +4384,12 @@ pub fn with_hole<'a>(
                                             hole,
                                         );
 
-                                let (lambda_expr, lambda_expr_var) =
-                                    procs.partial_exprs.get(symbol).unwrap();
-
                                         let snapshot = env.subs.snapshot();
                                         let cache_snapshot = layout_cache.snapshot();
                                         let _unified = roc_unify::unify::unify(
                                             env.subs,
                                             fn_var,
-                                    lambda_expr_var,
+                                            *lambda_expr_var,
                                             roc_unify::unify::Mode::EQ,
                                         );
 
@@ -4372,6 +4408,8 @@ pub fn with_hole<'a>(
                                     RawFunctionLayout::ZeroArgumentThunk(_) => {
                                         unreachable!("calling a non-closure layout")
                                     }
+                                }
+                            }
                         },
                         NotASymbol => {
                             // the expression is not a symbol. That means it's an expression
@@ -4705,6 +4743,43 @@ pub fn with_hole<'a>(
     }
 }
 
+#[inline(always)]
+fn get_specialization<'a>(
+    env: &mut Env<'a, '_>,
+    symbol_var: Variable,
+    symbol: Symbol,
+) -> Option<Symbol> {
+    use roc_solve::ability::type_implementing_member;
+    use roc_unify::unify::unify;
+
+    match env.abilities_store.member_def(symbol) {
+        None => {
+            // This is not an ability member, it doesn't need specialization.
+            None
+        }
+        Some(member) => {
+            let snapshot = env.subs.snapshot();
+            let (_, must_implement_ability) = unify(
+                env.subs,
+                symbol_var,
+                member.signature_var,
+                roc_unify::unify::Mode::EQ,
+            )
+            .expect_success("This typechecked previously");
+            env.subs.rollback_to(snapshot);
+            let specializing_type =
+                type_implementing_member(&must_implement_ability, member.parent_ability);
+
+            let specialization = env
+                .abilities_store
+                .get_specialization(symbol, specializing_type)
+                .expect("No specialization is recorded - I thought there would only be a type error here.");
+
+            Some(specialization.symbol)
+        }
+    }
+}
+
 #[allow(clippy::too_many_arguments)]
 fn construct_closure_data<'a>(
     env: &mut Env<'a, '_>,
@@ -5621,9 +5696,10 @@ pub fn from_can<'a>(
                         // At the definition site `n = 1` we only know `1` to have the type `[Int *]`,
                         // which won't be refined until the call `asU8 n`. Add it as a partial expression
                         // that will be specialized at each concrete usage site.
-                        procs
-                            .partial_exprs
-                            .insert(*symbol, def.loc_expr.value, def.expr_var);
+                        procs.partial_exprs.insert(
+                            *symbol,
+                            PolymorphicExpr::Expr(def.loc_expr.value, def.expr_var),
+                        );
 
                         let result = from_can(env, variable, cont.value, procs, layout_cache);
 
@@ -6329,7 +6405,7 @@ fn store_pattern_help<'a>(
 
     match can_pat {
         Identifier(symbol) => {
-            if let Some((_, var)) = procs.partial_exprs.get(outer_symbol) {
+            if let Some(&PolymorphicExpr::Expr(_, var)) = procs.partial_exprs.get(outer_symbol) {
                 // It might be the case that symbol we're storing hasn't been reified to a value
                 // yet, if it's polymorphic. Do that now.
                 stmt = specialize_symbol(
@@ -6727,6 +6803,13 @@ fn handle_variable_aliasing<'a, BuildRest>(
 where
     BuildRest: FnOnce(&mut Env<'a, '_>, &mut Procs<'a>, &mut LayoutCache<'a>) -> Stmt<'a>,
 {
+    if env.abilities_store.is_ability_member_name(right) {
+        procs
+            .partial_exprs
+            .insert(left, PolymorphicExpr::AbilityMember(right));
+        return build_rest(env, procs, layout_cache);
+    }
+
     if procs.partial_exprs.contains(right) {
         // If `right` links to a partial expression, make sure we link `left` to it as well, so
         // that usages of it will be specialized when building the rest of the program.
@@ -6804,7 +6887,7 @@ fn specialize_symbol<'a>(
     result: Stmt<'a>,
     original: Symbol,
 ) -> Stmt<'a> {
-    if let Some((expr, expr_var)) = procs.partial_exprs.get(original) {
+    if let Some(PolymorphicExpr::Expr(expr, expr_var)) = procs.partial_exprs.get(original) {
         // Specialize the expression type now, based off the `arg_var` we've been given.
         // TODO: cache the specialized result
         let snapshot = env.subs.snapshot();
@@ -6812,14 +6895,14 @@ fn specialize_symbol<'a>(
         let _unified = roc_unify::unify::unify(
             env.subs,
             arg_var.unwrap(),
-            expr_var,
+            *expr_var,
             roc_unify::unify::Mode::EQ,
         );
 
         let result = with_hole(
             env,
             expr.clone(),
-            expr_var,
+            *expr_var,
             procs,
             layout_cache,
             symbol,

compiler/solve/src/ability.rs

@@ -1,9 +1,11 @@
 use roc_can::abilities::AbilitiesStore;
+use roc_module::symbol::Symbol;
 use roc_region::all::{Loc, Region};
 use roc_types::subs::Subs;
 use roc_types::subs::Variable;
 use roc_types::types::{Category, PatternCategory};
 use roc_unify::unify::MustImplementAbility;
+use roc_unify::unify::MustImplementConstraints;
 
 use crate::solve::{IncompleteAbilityImplementation, TypeError};
 
@@ -18,17 +20,14 @@ pub enum AbilityImplError {
 }
 
 #[derive(Default)]
-pub struct DeferredMustImplementAbility(Vec<(Vec<MustImplementAbility>, AbilityImplError)>);
+pub struct DeferredMustImplementAbility(Vec<(MustImplementConstraints, AbilityImplError)>);
 
 impl DeferredMustImplementAbility {
-    pub fn add(&mut self, must_implement: Vec<MustImplementAbility>, on_error: AbilityImplError) {
+    pub fn add(&mut self, must_implement: MustImplementConstraints, on_error: AbilityImplError) {
         self.0.push((must_implement, on_error));
     }
 
     pub fn check(self, subs: &mut Subs, abilities_store: &AbilitiesStore) -> Vec<TypeError> {
-        // Two passes here. First up let's build up records of what types fully implement
-        // abilities, and what specializations are available/missing for the ones that don't.
-        // Use a vec since these lists should usually be pretty small.
         let mut good = vec![];
         let mut bad = vec![];
 
@@ -45,8 +44,21 @@ impl DeferredMustImplementAbility {
             };
         }
 
-        for (mias, _) in self.0.iter() {
-            for &mia @ MustImplementAbility { typ, ability } in mias {
+        let mut problems = vec![];
+
+        // Keep track of which types that have an incomplete ability were reported as part of
+        // another type error (from an expression or pattern). If we reported an error for a type
+        // that doesn't implement an ability in that context, we don't want to repeat the error
+        // message.
+        let mut reported_in_context = vec![];
+        let mut incomplete_not_in_context = vec![];
+
+        for (constraints, on_error) in self.0.into_iter() {
+            let must_implement = constraints.get_unique();
+
+            // First off, make sure we populate information about which of the "must implement"
+            // constraints are met, and which aren't.
+            for &mia @ MustImplementAbility { typ, ability } in must_implement.iter() {
                 if is_good!(&mia) || get_bad!(mia).is_some() {
                     continue;
                 }
@@ -80,22 +92,16 @@ impl DeferredMustImplementAbility {
                     ));
                 }
             }
-        }
 
-        // Now figure out what errors we need to report.
-        let mut problems = vec![];
-
-        // Keep track of which types that have an incomplete ability were reported as part of
-        // another type error (from an expression or pattern). If we reported an error for a type
-        // that doesn't implement an ability in that context, we don't want to repeat the error
-        // message.
-        let mut reported_in_context = vec![];
-        let mut incomplete_not_in_context = vec![];
-
-        for (must_implement, on_error) in self.0.into_iter() {
+            // Now, figure out what errors we need to report.
             use AbilityImplError::*;
             match on_error {
                 IncompleteAbility => {
+                    // These aren't attached to another type error, so if these must_implement
+                    // constraints aren't met, we'll emit a generic "this type doesn't implement an
+                    // ability" error message at the end. We only want to do this if it turns out
+                    // the "must implement" constraint indeed wasn't part of a more specific type
+                    // error.
                     incomplete_not_in_context.extend(must_implement);
                 }
                 BadExpr(region, category, var) => {
@@ -138,9 +144,11 @@ impl DeferredMustImplementAbility {
                         reported_in_context.extend(must_implement);
                     }
                 }
-            };
+            }
         }
 
+        // Go through and attach generic "type does not implement ability" errors, if they were not
+        // part of a larger context.
         for mia in incomplete_not_in_context.into_iter() {
             if let Some(must_implement) = get_bad!(mia) {
                 if !reported_in_context.contains(&mia) {
@@ -154,3 +162,29 @@ impl DeferredMustImplementAbility {
         problems
     }
 }
+
+/// Determines what type implements an ability member of a specialized signature, given the
+/// [MustImplementAbility] constraints of the signature.
+pub fn type_implementing_member(
+    specialization_must_implement_constraints: &MustImplementConstraints,
+    ability: Symbol,
+) -> Symbol {
+    debug_assert_eq!({
+            let ability_implementations_for_specialization =
+                specialization_must_implement_constraints
+                    .clone()
+                    .get_unique();
+
+            ability_implementations_for_specialization.len()
+        },
+        1,
+        "Multiple variables bound to an ability - this is ambiguous and should have been caught in canonicalization: {:?}",
+        specialization_must_implement_constraints
+    );
+
+    specialization_must_implement_constraints
+        .iter_for_ability(ability)
+        .next()
+        .unwrap()
+        .typ
+}

compiler/solve/src/lib.rs

@@ -2,6 +2,6 @@
 // See github.com/rtfeldman/roc/issues/800 for discussion of the large_enum_variant check.
 #![allow(clippy::large_enum_variant)]
 
-mod ability;
+pub mod ability;
 pub mod module;
 pub mod solve;

compiler/solve/src/solve.rs

@@ -1,4 +1,4 @@
-use crate::ability::{AbilityImplError, DeferredMustImplementAbility};
+use crate::ability::{type_implementing_member, AbilityImplError, DeferredMustImplementAbility};
 use bumpalo::Bump;
 use roc_can::abilities::{AbilitiesStore, MemberSpecialization};
 use roc_can::constraint::Constraint::{self, *};
@@ -703,7 +703,6 @@ fn solve(
                     check_ability_specialization(
                         arena,
                         subs,
-                        &new_env,
                         pools,
                         rank,
                         abilities_store,
@@ -812,7 +811,6 @@ fn solve(
                     check_ability_specialization(
                         arena,
                         subs,
-                        &new_env,
                         pools,
                         rank,
                         abilities_store,
@@ -1282,7 +1280,6 @@ fn solve(
 fn check_ability_specialization(
     arena: &Bump,
     subs: &mut Subs,
-    env: &Env,
     pools: &mut Pools,
     rank: Rank,
     abilities_store: &mut AbilitiesStore,
@@ -1295,9 +1292,7 @@ fn check_ability_specialization(
     // inferred type for the specialization actually aligns with the expected
     // implementation.
     if let Some((root_symbol, root_data)) = abilities_store.root_name_and_def(symbol) {
-        let root_signature_var = env
-            .get_var_by_symbol(&root_symbol)
-            .expect("Ability should be registered in env by now!");
+        let root_signature_var = root_data.signature_var;
 
         // Check if they unify - if they don't, then the claimed specialization isn't really one,
         // and that's a type error!
@@ -1351,16 +1346,8 @@ fn check_ability_specialization(
 
                 // First, figure out and register for what type does this symbol specialize
                 // the ability member.
-                let mut ability_implementations_for_specialization = must_implement_ability
-                    .iter()
-                    .filter(|mia| mia.ability == root_data.parent_ability)
-                    .collect::<Vec<_>>();
-                ability_implementations_for_specialization.dedup();
-
-                debug_assert!(ability_implementations_for_specialization.len() == 1, "Multiple variables bound to an ability - this is ambiguous and should have been caught in canonicalization");
-
-                // This is a valid specialization! Record it.
-                let specialization_type = ability_implementations_for_specialization[0].typ;
+                let specialization_type =
+                    type_implementing_member(&must_implement_ability, root_data.parent_ability);
                 let specialization = MemberSpecialization {
                     symbol,
                     region: symbol_loc_var.region,

compiler/solve/tests/solve_expr.rs

@@ -5913,4 +5913,23 @@ mod solve_expr {
             "U64",
         )
     }
+
+    #[test]
+    fn alias_ability_member() {
+        infer_eq_without_problem(
+            indoc!(
+                r#"
+                app "test" provides [ thething ] to "./platform"
+
+                Hash has
+                    hash : a -> U64 | a has Hash
+
+                thething =
+                    itis = hash
+                    itis
+                "#
+            ),
+            "a -> U64 | a has Hash",
+        )
+    }
 }

compiler/test_gen/src/gen_abilities.rs

@@ -0,0 +1,86 @@
+#[cfg(feature = "gen-llvm")]
+use crate::helpers::llvm::assert_evals_to;
+
+#[cfg(feature = "gen-dev")]
+use crate::helpers::dev::assert_evals_to;
+
+#[cfg(feature = "gen-wasm")]
+use crate::helpers::wasm::assert_evals_to;
+
+#[cfg(test)]
+use indoc::indoc;
+
+#[test]
+#[cfg(any(feature = "gen-llvm", feature = "gen-wasm"))]
+fn hash_specialization() {
+    assert_evals_to!(
+        indoc!(
+            r#"
+            app "test" provides [ main ] to "./platform"
+
+            Hash has
+                hash : a -> U64 | a has Hash
+
+            Id := U64
+
+            hash = \$Id n -> n
+
+            main = hash ($Id 1234)
+            "#
+        ),
+        1234,
+        u64
+    );
+}
+
+#[test]
+#[cfg(any(feature = "gen-llvm", feature = "gen-wasm"))]
+fn hash_specialization_multiple_add() {
+    assert_evals_to!(
+        indoc!(
+            r#"
+            app "test" provides [ main ] to "./platform"
+
+            Hash has
+                hash : a -> U64 | a has Hash
+
+            Id := U64
+
+            hash = \$Id n -> n
+
+            One := {}
+
+            hash = \$One _ -> 1
+
+            main = hash ($Id 1234) + hash ($One {})
+            "#
+        ),
+        1235,
+        u64
+    );
+}
+
+#[test]
+#[cfg(any(feature = "gen-llvm", feature = "gen-wasm"))]
+fn alias_member_specialization() {
+    assert_evals_to!(
+        indoc!(
+            r#"
+            app "test" provides [ main ] to "./platform"
+
+            Hash has
+                hash : a -> U64 | a has Hash
+
+            Id := U64
+
+            hash = \$Id n -> n
+
+            main =
+                aliasedHash = hash
+                aliasedHash ($Id 1234)
+            "#
+        ),
+        1234,
+        u64
+    );
+}

compiler/test_gen/src/tests.rs

@@ -4,6 +4,7 @@
 // we actually want to compare against the literal float bits
 #![allow(clippy::float_cmp)]
 
+pub mod gen_abilities;
 pub mod gen_compare;
 pub mod gen_dict;
 pub mod gen_list;

compiler/test_mono/generated/specialize_ability_call.txt

@@ -0,0 +1,7 @@
+procedure Test.5 (Test.8):
+    ret Test.8;
+
+procedure Test.0 ():
+    let Test.10 : U64 = 1234i64;
+    let Test.9 : U64 = CallByName Test.5 Test.10;
+    ret Test.9;

compiler/test_mono/src/tests.rs

@@ -1294,6 +1294,25 @@ fn issue_2811() {
     )
 }
 
+#[mono_test]
+fn specialize_ability_call() {
+    indoc!(
+        r#"
+        app "test" provides [ main ] to "./platform"
+
+        Hash has
+            hash : a -> U64 | a has Hash
+
+        Id := U64
+
+        hash : Id -> U64
+        hash = \$Id n -> n
+
+        main = hash ($Id 1234)
+        "#
+    )
+}
+
 // #[ignore]
 // #[mono_test]
 // fn static_str_closure() {

compiler/unify/src/unify.rs

@@ -1,5 +1,5 @@
 use bitflags::bitflags;
-use roc_error_macros::todo_abilities;
+use roc_error_macros::{internal_error, todo_abilities};
 use roc_module::ident::{Lowercase, TagName};
 use roc_module::symbol::Symbol;
 use roc_types::subs::Content::{self, *};
@@ -134,14 +134,26 @@ pub struct Context {
 pub enum Unified {
     Success {
         vars: Pool,
-        must_implement_ability: Vec<MustImplementAbility>,
+        must_implement_ability: MustImplementConstraints,
     },
     Failure(Pool, ErrorType, ErrorType, DoesNotImplementAbility),
     BadType(Pool, roc_types::types::Problem),
 }
 
+impl Unified {
+    pub fn expect_success(self, err_msg: &'static str) -> (Pool, MustImplementConstraints) {
+        match self {
+            Unified::Success {
+                vars,
+                must_implement_ability,
+            } => (vars, must_implement_ability),
+            _ => internal_error!("{}", err_msg),
+        }
+    }
+}
+
 /// Specifies that `type` must implement the ability `ability`.
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
 pub struct MustImplementAbility {
     // This only points to opaque type names currently.
     // TODO(abilities) support structural types in general
@@ -149,12 +161,39 @@ pub struct MustImplementAbility {
     pub ability: Symbol,
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Default)]
+pub struct MustImplementConstraints(Vec<MustImplementAbility>);
+
+impl MustImplementConstraints {
+    pub fn push(&mut self, must_implement: MustImplementAbility) {
+        self.0.push(must_implement)
+    }
+
+    pub fn extend(&mut self, other: Self) {
+        self.0.extend(other.0)
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.0.is_empty()
+    }
+
+    pub fn get_unique(mut self) -> Vec<MustImplementAbility> {
+        self.0.sort();
+        self.0.dedup();
+        self.0
+    }
+
+    pub fn iter_for_ability(&self, ability: Symbol) -> impl Iterator<Item = &MustImplementAbility> {
+        self.0.iter().filter(move |mia| mia.ability == ability)
+    }
+}
+
 #[derive(Debug, Default)]
 pub struct Outcome {
     mismatches: Vec<Mismatch>,
     /// We defer these checks until the end of a solving phase.
     /// NOTE: this vector is almost always empty!
-    must_implement_ability: Vec<MustImplementAbility>,
+    must_implement_ability: MustImplementConstraints,
 }
 
 impl Outcome {

reporting/tests/test_reporting.rs

@@ -125,6 +125,7 @@ mod test_reporting {
             mut solved,
             exposed_to_host,
             mut declarations_by_id,
+            abilities_store,
             ..
         } = result?;
 
@@ -177,6 +178,7 @@ mod test_reporting {
                 target_info,
                 // call_specialization_counter=0 is reserved
                 call_specialization_counter: 1,
+                abilities_store: &abilities_store,
             };
             let _mono_expr = Stmt::new(
                 &mut mono_env,
@@ -334,6 +336,7 @@ mod test_reporting {
                 target_info,
                 // call_specialization_counter=0 is reserved
                 call_specialization_counter: 1,
+                abilities_store: &abilities_store,
             };
             let _mono_expr = Stmt::new(
                 &mut mono_env,