Merge pull request #3321 from rtfeldman/abilities-stuff

A few changes to abilities in support of derived impls

compiler/can/src/abilities.rs

@@ -1,3 +1,5 @@
+use std::num::NonZeroU32;
+
 use roc_collections::{all::MutMap, VecMap, VecSet};
 use roc_error_macros::internal_error;
 use roc_module::symbol::{ModuleId, Symbol};
@@ -71,6 +73,17 @@ pub type SpecializationsMap<Phase> = VecMap<(Symbol, Symbol), MemberSpecializati
 pub type PendingSpecializations = SpecializationsMap<Pending>;
 pub type ResolvedSpecializations = SpecializationsMap<Resolved>;
 
+/// Solved lambda sets for an ability member specialization. For example, if we have
+///
+///   Default has default : {} -[[] + a:default:1]-> a | a has Default
+///
+///   A := {}
+///   default = \{} -[[closA]]-> @A {}
+///
+/// and this [MemberSpecialization] is for `A`, then there is a mapping of
+/// `1` to the variable representing `[[closA]]`.
+pub type SpecializationLambdaSets = VecMap<u8, Variable>;
+
 /// A particular specialization of an ability member.
 #[derive(Debug, Clone)]
 pub struct MemberSpecialization<Phase: ResolvePhase> {
@@ -78,20 +91,11 @@ pub struct MemberSpecialization<Phase: ResolvePhase> {
 
     pub symbol: Symbol,
 
-    /// Solved lambda sets for an ability member specialization. For example, if we have
+    pub specialization_lambda_sets: SpecializationLambdaSets,
-    ///
-    ///   Default has default : {} -[[] + a:default:1]-> a | a has Default
-    ///   
-    ///   A := {}
-    ///   default = \{} -[[closA]]-> @A {}
-    ///
-    /// and this [MemberSpecialization] is for `A`, then there is a mapping of
-    /// `1` to the variable representing `[[closA]]`.
-    pub specialization_lambda_sets: VecMap<u8, Variable>,
 }
 
 impl MemberSpecialization<Resolved> {
-    pub fn new(symbol: Symbol, specialization_lambda_sets: VecMap<u8, Variable>) -> Self {
+    pub fn new(symbol: Symbol, specialization_lambda_sets: SpecializationLambdaSets) -> Self {
         Self {
             _phase: Default::default(),
             symbol,
@@ -101,14 +105,9 @@ impl MemberSpecialization<Resolved> {
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-pub struct SpecializationId(u32);
+pub struct SpecializationId(NonZeroU32);
 
-#[allow(clippy::derivable_impls)] // let's be explicit about this
+static_assertions::assert_eq_size!(SpecializationId, Option<SpecializationId>);
-impl Default for SpecializationId {
-    fn default() -> Self {
-        Self(0)
-    }
-}
 
 pub enum SpecializationLambdaSetError {}
 
@@ -117,7 +116,7 @@ pub enum SpecializationLambdaSetError {}
 // TODO(abilities): this should probably go on the Scope, I don't put it there for now because we
 // are only dealing with intra-module abilities for now.
 // TODO(abilities): many of these should be `VecMap`s. Do some benchmarking.
-#[derive(Default, Debug, Clone)]
+#[derive(Debug, Clone)]
 pub struct IAbilitiesStore<Phase: ResolvePhase> {
     /// Maps an ability to the members defining it.
     members_of_ability: MutMap<Symbol, Vec<Symbol>>,
@@ -139,13 +138,28 @@ pub struct IAbilitiesStore<Phase: ResolvePhase> {
     /// member `member`, to the exact symbol that implements the ability.
     declared_specializations: SpecializationsMap<Phase>,
 
-    next_specialization_id: u32,
+    next_specialization_id: NonZeroU32,
 
     /// Resolved specializations for a symbol. These might be ephemeral (known due to type solving),
     /// or resolved on-the-fly during mono.
     resolved_specializations: MutMap<SpecializationId, Symbol>,
 }
 
+impl<Phase: ResolvePhase> Default for IAbilitiesStore<Phase> {
+    fn default() -> Self {
+        Self {
+             members_of_ability: Default::default(),
+             ability_members: Default::default(),
+             specialization_to_root: Default::default(),
+             declared_specializations: Default::default(),
+             next_specialization_id:
+                 // Safety: 1 != 0
+                 unsafe { NonZeroU32::new_unchecked(1) },
+             resolved_specializations: Default::default(),
+         }
+    }
+}
+
 pub type AbilitiesStore = IAbilitiesStore<Resolved>;
 pub type PendingAbilitiesStore = IAbilitiesStore<Pending>;
 
@@ -214,10 +228,12 @@ impl<Phase: ResolvePhase> IAbilitiesStore<Phase> {
     }
 
     pub fn fresh_specialization_id(&mut self) -> SpecializationId {
-        debug_assert!(self.next_specialization_id != std::u32::MAX);
+        debug_assert!(self.next_specialization_id.get() != std::u32::MAX);
 
         let id = SpecializationId(self.next_specialization_id);
-        self.next_specialization_id += 1;
+        // Safety: we already checked this won't overflow, and we started > 0.
+        self.next_specialization_id =
+            unsafe { NonZeroU32::new_unchecked(self.next_specialization_id.get() + 1) };
         id
     }
 
@@ -431,8 +447,8 @@ impl IAbilitiesStore<Pending> {
             );
         }
 
-        debug_assert!(next_specialization_id == 0);
+        debug_assert_eq!(next_specialization_id.get(), 1);
-        debug_assert!(self.next_specialization_id == 0);
+        debug_assert_eq!(self.next_specialization_id.get(), 1);
         debug_assert!(resolved_specializations.is_empty());
         debug_assert!(self.resolved_specializations.is_empty());
     }

compiler/can/src/expr.rs

@@ -101,8 +101,10 @@ pub enum Expr {
     AbilityMember(
         /// Actual member name
         Symbol,
-        /// Specialization to use, and its variable
+        /// Specialization to use, and its variable.
-        SpecializationId,
+        /// The specialization id may be [`None`] if construction of an ability member usage can
+        /// prove the usage is polymorphic.
+        Option<SpecializationId>,
         Variable,
     ),
 
@@ -1353,7 +1355,7 @@ fn canonicalize_var_lookup(
                 if scope.abilities_store.is_ability_member_name(symbol) {
                     AbilityMember(
                         symbol,
-                        scope.abilities_store.fresh_specialization_id(),
+                        Some(scope.abilities_store.fresh_specialization_id()),
                         var_store.fresh(),
                     )
                 } else {
@@ -1376,7 +1378,7 @@ fn canonicalize_var_lookup(
                 if scope.abilities_store.is_ability_member_name(symbol) {
                     AbilityMember(
                         symbol,
-                        scope.abilities_store.fresh_specialization_id(),
+                        Some(scope.abilities_store.fresh_specialization_id()),
                         var_store.fresh(),
                     )
                 } else {

compiler/can/src/traverse.rs

@@ -550,7 +550,9 @@ pub fn find_ability_member_and_owning_type_at(
             if region == self.region {
                 if let &Expr::AbilityMember(member_symbol, specialization_id, _var) = expr {
                     debug_assert!(self.found.is_none());
-                    self.found = match self.abilities_store.get_resolved(specialization_id) {
+                    self.found = match specialization_id
+                        .and_then(|id| self.abilities_store.get_resolved(id))
+                    {
                         Some(spec_symbol) => {
                             let spec_type = find_specialization_type_of_symbol(
                                 spec_symbol,

compiler/constrain/src/expr.rs

@@ -434,7 +434,8 @@ pub fn constrain_expr(
                 region,
             );
 
-            // Make sure we attempt to resolve the specialization.
+            // Make sure we attempt to resolve the specialization, if we need to.
+            if let Some(specialization_id) = specialization_id {
                 env.resolutions_to_make.push(OpportunisticResolve {
                     specialization_variable: specialization_var,
                     specialization_expectation: constraints.push_expected_type(
@@ -443,6 +444,7 @@ pub fn constrain_expr(
                     member: symbol,
                     specialization_id,
                 });
+            }
 
             constraints.and_constraint([store_expected, lookup_constr])
         }

compiler/mono/src/ir.rs

@@ -3772,14 +3772,13 @@ pub fn with_hole<'a>(
 
             specialize_naked_symbol(env, variable, procs, layout_cache, assigned, hole, symbol)
         }
-        AbilityMember(_member, specialization_id, _) => {
+        AbilityMember(member, specialization_id, specialization_var) => {
-            let specialization_symbol =
+            let specialization_symbol = late_resolve_ability_specialization(
-                env.abilities
+                env,
-                    .with_module_abilities_store(env.home, |store| {
+                member,
-                        store
+                specialization_id,
-                            .get_resolved(specialization_id)
+                specialization_var,
-                            .expect("Specialization was never made!")
+            );
-                    });
 
             specialize_naked_symbol(
                 env,
@@ -5161,12 +5160,13 @@ pub fn with_hole<'a>(
 fn late_resolve_ability_specialization<'a>(
     env: &mut Env<'a, '_>,
     member: Symbol,
-    specialization_id: SpecializationId,
+    specialization_id: Option<SpecializationId>,
     specialization_var: Variable,
 ) -> Symbol {
-    let opt_resolved = env
+    let opt_resolved = specialization_id.and_then(|id| {
-        .abilities
+        env.abilities
-        .with_module_abilities_store(env.home, |store| store.get_resolved(specialization_id));
+            .with_module_abilities_store(env.home, |store| store.get_resolved(id))
+    });
 
     if let Some(spec_symbol) = opt_resolved {
         // Fast path: specialization is monomorphic, was found during solving.

compiler/solve/src/ability.rs

@@ -57,7 +57,7 @@ pub enum Unfulfilled {
 
 /// Indexes a deriving of an ability for an opaque type.
 #[derive(Debug, PartialEq, Clone, Copy)]
-pub struct DeriveKey {
+pub struct RequestedDeriveKey {
     pub opaque: Symbol,
     pub ability: Symbol,
 }
@@ -72,7 +72,7 @@ struct ImplKey {
 #[derive(Debug)]
 pub struct PendingDerivesTable(
     /// derive key -> (opaque type var to use for checking, derive region)
-    VecMap<DeriveKey, (Variable, Region)>,
+    VecMap<RequestedDeriveKey, (Variable, Region)>,
 );
 
 impl PendingDerivesTable {
@@ -89,7 +89,7 @@ impl PendingDerivesTable {
                     ability.is_builtin_ability(),
                     "Not a builtin - should have been caught during can"
                 );
-                let derive_key = DeriveKey { opaque, ability };
+                let derive_key = RequestedDeriveKey { opaque, ability };
 
                 // Neither rank nor pools should matter here.
                 let opaque_var =
@@ -117,7 +117,7 @@ pub struct DeferredObligations {
     /// Derives that are claimed, but have also been determined to have
     /// specializations. Maps to the first member specialization of the same
     /// ability.
-    dominated_derives: VecMap<DeriveKey, Region>,
+    dominated_derives: VecMap<RequestedDeriveKey, Region>,
 }
 
 impl DeferredObligations {
@@ -133,7 +133,7 @@ impl DeferredObligations {
         self.obligations.push((must_implement, on_error));
     }
 
-    pub fn dominate(&mut self, key: DeriveKey, impl_region: Region) {
+    pub fn dominate(&mut self, key: RequestedDeriveKey, impl_region: Region) {
         // Only builtin abilities can be derived, and hence dominated.
         if self.pending_derives.0.contains_key(&key) && !self.dominated_derives.contains_key(&key) {
             self.dominated_derives.insert(key, impl_region);
@@ -153,7 +153,7 @@ impl DeferredObligations {
         self,
         subs: &mut Subs,
         abilities_store: &AbilitiesStore,
-    ) -> (Vec<TypeError>, Vec<DeriveKey>) {
+    ) -> (Vec<TypeError>, Vec<RequestedDeriveKey>) {
         let mut problems = vec![];
 
         let Self {
@@ -282,11 +282,11 @@ type ObligationResult = Result<(), Unfulfilled>;
 
 struct ObligationCache<'a> {
     abilities_store: &'a AbilitiesStore,
-    dominated_derives: &'a VecMap<DeriveKey, Region>,
+    dominated_derives: &'a VecMap<RequestedDeriveKey, Region>,
     pending_derives: &'a PendingDerivesTable,
 
     impl_cache: VecMap<ImplKey, ObligationResult>,
-    derive_cache: VecMap<DeriveKey, ObligationResult>,
+    derive_cache: VecMap<RequestedDeriveKey, ObligationResult>,
 }
 
 enum ReadCache {
@@ -342,7 +342,7 @@ impl ObligationCache<'_> {
 
     fn check_opaque(&mut self, subs: &mut Subs, opaque: Symbol, ability: Symbol) -> ReadCache {
         let impl_key = ImplKey { opaque, ability };
-        let derive_key = DeriveKey { opaque, ability };
+        let derive_key = RequestedDeriveKey { opaque, ability };
 
         match self.pending_derives.0.get(&derive_key) {
             Some(&(opaque_real_var, derive_region)) => {
@@ -377,7 +377,7 @@ impl ObligationCache<'_> {
             ReadCache::Impl => self.impl_cache.get(&ImplKey { opaque, ability }).unwrap(),
             ReadCache::Derive => self
                 .derive_cache
-                .get(&DeriveKey { opaque, ability })
+                .get(&RequestedDeriveKey { opaque, ability })
                 .unwrap(),
         }
     }
@@ -418,7 +418,7 @@ impl ObligationCache<'_> {
     fn check_derive(
         &mut self,
         subs: &mut Subs,
-        derive_key: DeriveKey,
+        derive_key: RequestedDeriveKey,
         opaque_real_var: Variable,
         derive_region: Region,
     ) {

compiler/solve/src/solve.rs

@@ -1,6 +1,6 @@
 use crate::ability::{
     resolve_ability_specialization, type_implementing_specialization, AbilityImplError,
-    DeferredObligations, DeriveKey, PendingDerivesTable, Resolved, Unfulfilled,
+    DeferredObligations, PendingDerivesTable, RequestedDeriveKey, Resolved, Unfulfilled,
 };
 use bumpalo::Bump;
 use roc_can::abilities::{AbilitiesStore, MemberSpecialization};
@@ -1660,7 +1660,7 @@ fn check_ability_specialization(
                             .add(must_implement_ability, AbilityImplError::IncompleteAbility);
                         // This specialization dominates any derives that might be present.
                         deferred_obligations.dominate(
-                            DeriveKey {
+                            RequestedDeriveKey {
                                 opaque,
                                 ability: parent_ability,
                             },