moved all crates into seperate folder + related path fixes

crates/compiler/can/src/abilities.rs

@@ -0,0 +1,554 @@
+use std::num::NonZeroU32;
+
+use roc_collections::{all::MutMap, VecMap, VecSet};
+use roc_error_macros::internal_error;
+use roc_module::symbol::{ModuleId, Symbol};
+use roc_region::all::Region;
+use roc_types::{subs::Variable, types::Type};
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct MemberVariables {
+    pub able_vars: Vec<Variable>,
+    /// This includes - named rigid vars, lambda sets, wildcards. See
+    /// [`IntroducedVariables::collect_rigid`](crate::annotation::IntroducedVariables::collect_rigid).
+    pub rigid_vars: Vec<Variable>,
+    pub flex_vars: Vec<Variable>,
+}
+
+/// The member and its signature is defined locally, in the module the store is created for.
+/// We need to instantiate and introduce this during solving.
+#[derive(Debug, Clone)]
+pub struct ResolvedMemberType(Variable);
+
+/// Member type information that needs to be resolved from imports.
+#[derive(Debug, Clone)]
+pub enum PendingMemberType {
+    /// The member and its signature is defined locally, in the module the store is created for.
+    /// We need to instantiate and introduce this during solving.
+    Local {
+        signature_var: Variable,
+        signature: Type,
+        variables: MemberVariables,
+    },
+    /// The member was defined in another module, so we'll import its variable when it's time to
+    /// solve. At that point we'll resolve `var` here.
+    Imported,
+}
+
+pub trait ResolvePhase: std::fmt::Debug + Clone + Copy {
+    type MemberType: std::fmt::Debug + Clone;
+}
+
+#[derive(Default, Debug, Clone, Copy)]
+pub struct Pending;
+impl ResolvePhase for Pending {
+    type MemberType = PendingMemberType;
+}
+
+#[derive(Default, Debug, Clone, Copy)]
+pub struct Resolved;
+impl ResolvePhase for Resolved {
+    type MemberType = ResolvedMemberType;
+}
+
+/// Stores information about an ability member definition, including the parent ability, the
+/// defining type, and what type variables need to be instantiated with instances of the ability.
+// TODO: SoA and put me in an arena
+#[derive(Debug, Clone)]
+pub struct AbilityMemberData<Phase: ResolvePhase> {
+    pub parent_ability: Symbol,
+    pub region: Region,
+    pub typ: Phase::MemberType,
+}
+
+impl AbilityMemberData<Resolved> {
+    pub fn signature_var(&self) -> Variable {
+        self.typ.0
+    }
+}
+
+/// (member, specialization type) -> specialization
+pub type SpecializationsMap<Phase> = VecMap<(Symbol, Symbol), MemberSpecialization<Phase>>;
+
+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> {
+    _phase: std::marker::PhantomData<Phase>,
+
+    pub symbol: Symbol,
+
+    pub specialization_lambda_sets: SpecializationLambdaSets,
+}
+
+impl MemberSpecialization<Resolved> {
+    pub fn new(symbol: Symbol, specialization_lambda_sets: SpecializationLambdaSets) -> Self {
+        Self {
+            _phase: Default::default(),
+            symbol,
+            specialization_lambda_sets,
+        }
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct SpecializationId(NonZeroU32);
+
+static_assertions::assert_eq_size!(SpecializationId, Option<SpecializationId>);
+
+pub enum SpecializationLambdaSetError {}
+
+/// Stores information about what abilities exist in a scope, what it means to implement an
+/// ability, and what types implement them.
+// 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(Debug, Clone)]
+pub struct IAbilitiesStore<Phase: ResolvePhase> {
+    /// Maps an ability to the members defining it.
+    members_of_ability: MutMap<Symbol, Vec<Symbol>>,
+
+    /// Information about all members composing abilities.
+    ability_members: MutMap<Symbol, AbilityMemberData<Phase>>,
+
+    /// Map of symbols that specialize an ability member to the root ability symbol name.
+    /// For example, for the program
+    ///   Hash has hash : a -> U64 | a has Hash
+    ///            ^^^^ gets the symbol "#hash"
+    ///   hash = \@Id n -> n
+    ///   ^^^^ gets the symbol "#hash1"
+    ///
+    /// We keep the mapping #hash1->#hash
+    specialization_to_root: MutMap<Symbol, Symbol>,
+
+    /// Maps a tuple (member, type) specifying that `type` declares an implementation of an ability
+    /// member `member`, to the exact symbol that implements the ability.
+    declared_specializations: SpecializationsMap<Phase>,
+
+    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>;
+
+impl<Phase: ResolvePhase> IAbilitiesStore<Phase> {
+    /// Records the definition of an ability, including its members.
+    pub fn register_ability<I>(&mut self, ability: Symbol, members: I)
+    where
+        I: IntoIterator<Item = (Symbol, AbilityMemberData<Phase>)>,
+        I::IntoIter: ExactSizeIterator,
+    {
+        let members = members.into_iter();
+        let mut members_vec = Vec::with_capacity(members.len());
+        for (member, member_data) in members {
+            members_vec.push(member);
+            let old_member = self.ability_members.insert(member, member_data);
+            debug_assert!(old_member.is_none(), "Replacing existing member definition");
+        }
+        let old_ability = self.members_of_ability.insert(ability, members_vec);
+        debug_assert!(
+            old_ability.is_none(),
+            "Replacing existing ability definition"
+        );
+    }
+
+    /// Checks if `name` is a root ability member symbol name.
+    /// Note that this will return `false` for specializations of an ability member, which have
+    /// different symbols from the root.
+    pub fn is_ability_member_name(&self, name: Symbol) -> bool {
+        self.ability_members.contains_key(&name)
+    }
+
+    pub fn is_ability(&self, ability: Symbol) -> bool {
+        self.members_of_ability.contains_key(&ability)
+    }
+
+    /// Iterator over all abilities and their members that this store knows about.
+    pub fn iter_abilities(&self) -> impl Iterator<Item = (Symbol, &[Symbol])> {
+        self.members_of_ability
+            .iter()
+            .map(|(k, v)| (*k, v.as_slice()))
+    }
+
+    /// Returns information about all known ability members and their root symbols.
+    pub fn root_ability_members(&self) -> &MutMap<Symbol, AbilityMemberData<Phase>> {
+        &self.ability_members
+    }
+
+    /// Returns whether a symbol is declared to specialize an ability member.
+    pub fn is_specialization_name(&self, symbol: Symbol) -> bool {
+        self.specialization_to_root.contains_key(&symbol)
+    }
+
+    /// Records that the symbol `specializing_symbol` claims to specialize `ability_member`; for
+    /// example the symbol of `hash : Id -> U64` specializing `hash : a -> U64 | a has Hash`.
+    pub fn register_specializing_symbol(
+        &mut self,
+        specializing_symbol: Symbol,
+        ability_member: Symbol,
+    ) {
+        self.specialization_to_root
+            .insert(specializing_symbol, ability_member);
+    }
+
+    pub fn members_of_ability(&self, ability: Symbol) -> Option<&[Symbol]> {
+        self.members_of_ability.get(&ability).map(|v| v.as_ref())
+    }
+
+    pub fn fresh_specialization_id(&mut self) -> SpecializationId {
+        debug_assert!(self.next_specialization_id.get() != std::u32::MAX);
+
+        let id = SpecializationId(self.next_specialization_id);
+        // 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
+    }
+
+    /// Creates a store from [`self`] that closes over the abilities/members given by the
+    /// imported `symbols`, and their specializations (if any).
+    pub fn closure_from_imported(&self, symbols: &VecSet<Symbol>) -> PendingAbilitiesStore {
+        let Self {
+            members_of_ability,
+            ability_members,
+            declared_specializations,
+
+            // Covered by `declared_specializations`
+            specialization_to_root: _,
+
+            // Taking closure for a new module, so specialization IDs can be fresh
+            next_specialization_id: _,
+            resolved_specializations: _,
+        } = self;
+
+        let mut new = PendingAbilitiesStore::default();
+
+        // 1. Figure out the abilities we need to introduce.
+        let mut abilities_to_introduce = VecSet::with_capacity(2);
+        symbols.iter().for_each(|symbol| {
+            if let Some(member_data) = ability_members.get(symbol) {
+                // If the symbol is member of an ability, we need to capture the entire ability.
+                abilities_to_introduce.insert(member_data.parent_ability);
+            }
+            if members_of_ability.contains_key(symbol) {
+                abilities_to_introduce.insert(*symbol);
+            }
+        });
+
+        // 2. Add each ability, and any specializations of its members we know about.
+        for ability in abilities_to_introduce.into_iter() {
+            let members = members_of_ability.get(&ability).unwrap();
+            let mut imported_member_data = Vec::with_capacity(members.len());
+            for member in members {
+                let AbilityMemberData {
+                    parent_ability,
+                    region,
+                    typ: _,
+                } = ability_members.get(member).unwrap().clone();
+                // All external members need to be marked as imported. We'll figure out their real
+                // type variables when it comes time to solve the module we're currently importing
+                // into.
+                let imported_data = AbilityMemberData {
+                    parent_ability,
+                    region,
+                    typ: PendingMemberType::Imported,
+                };
+
+                imported_member_data.push((*member, imported_data));
+            }
+
+            new.register_ability(ability, imported_member_data);
+
+            // Add any specializations of the ability's members we know about.
+            declared_specializations
+                .iter()
+                .filter(|((member, _), _)| members.contains(member))
+                .for_each(|(&(member, typ), specialization)| {
+                    new.register_specializing_symbol(specialization.symbol, member);
+                    new.import_specialization(member, typ, specialization);
+                });
+        }
+
+        new
+    }
+}
+
+impl IAbilitiesStore<Resolved> {
+    /// Finds the symbol name and ability member definition for a symbol specializing the ability
+    /// member, if it specializes any.
+    /// For example, suppose `hash : Id -> U64` has symbol #hash1 and specializes
+    /// `hash : a -> U64 | a has Hash` with symbol #hash. Calling this with #hash1 would retrieve
+    /// the ability member data for #hash.
+    pub fn root_name_and_def(
+        &self,
+        specializing_symbol: Symbol,
+    ) -> Option<(Symbol, &AbilityMemberData<Resolved>)> {
+        let root_symbol = self.specialization_to_root.get(&specializing_symbol)?;
+        debug_assert!(self.ability_members.contains_key(root_symbol));
+        let root_data = self.ability_members.get(root_symbol).unwrap();
+        Some((*root_symbol, root_data))
+    }
+
+    /// Finds the ability member definition for a member name.
+    pub fn member_def(&self, member: Symbol) -> Option<&AbilityMemberData<Resolved>> {
+        self.ability_members.get(&member)
+    }
+
+    /// Returns an iterator over pairs ((ability member, type), specialization) specifying that
+    /// "ability member" has a "specialization" for type "type".
+    pub fn iter_specializations(
+        &self,
+    ) -> impl Iterator<Item = ((Symbol, Symbol), &MemberSpecialization<Resolved>)> + '_ {
+        self.declared_specializations.iter().map(|(k, v)| (*k, v))
+    }
+
+    /// Retrieves the specialization of `member` for `typ`, if it exists.
+    pub fn get_specialization(
+        &self,
+        member: Symbol,
+        typ: Symbol,
+    ) -> Option<&MemberSpecialization<Resolved>> {
+        self.declared_specializations.get(&(member, typ))
+    }
+
+    /// Records a specialization of `ability_member` with specialized type `implementing_type`.
+    /// Entries via this function are considered a source of truth. It must be ensured that a
+    /// specialization is validated before being registered here.
+    pub fn register_specialization_for_type(
+        &mut self,
+        ability_member: Symbol,
+        implementing_type: Symbol,
+        specialization: MemberSpecialization<Resolved>,
+    ) {
+        let old_spec = self
+            .declared_specializations
+            .insert((ability_member, implementing_type), specialization);
+        debug_assert!(old_spec.is_none(), "Replacing existing specialization");
+    }
+
+    pub fn insert_resolved(&mut self, id: SpecializationId, specialization: Symbol) {
+        // May not be a thing in mono
+        // debug_assert!(self.is_specialization_name(specialization));
+
+        let old_specialization = self.resolved_specializations.insert(id, specialization);
+
+        debug_assert!(
+            old_specialization.is_none(),
+            "Existing resolution: {:?}",
+            old_specialization
+        );
+    }
+
+    pub fn remove_resolved(&mut self, id: SpecializationId) {
+        let old_specialization = self.resolved_specializations.remove(&id);
+
+        debug_assert!(
+            old_specialization.is_some(),
+            "Trying to remove a resolved specialization that was never there!",
+        );
+    }
+
+    pub fn get_resolved(&self, id: SpecializationId) -> Option<Symbol> {
+        self.resolved_specializations.get(&id).copied()
+    }
+}
+
+impl IAbilitiesStore<Pending> {
+    pub fn import_specialization(
+        &mut self,
+        ability_member: Symbol,
+        implementing_type: Symbol,
+        specialization: &MemberSpecialization<impl ResolvePhase>,
+    ) {
+        let MemberSpecialization {
+            _phase,
+            symbol,
+            specialization_lambda_sets,
+        } = specialization;
+
+        let old_spec = self.declared_specializations.insert(
+            (ability_member, implementing_type),
+            MemberSpecialization {
+                _phase: Default::default(),
+                symbol: *symbol,
+                specialization_lambda_sets: specialization_lambda_sets.clone(),
+            },
+        );
+        debug_assert!(old_spec.is_none(), "Replacing existing specialization");
+    }
+
+    pub fn union(&mut self, other: Self) {
+        let Self {
+            members_of_ability: other_members_of_ability,
+            ability_members: mut other_ability_members,
+            specialization_to_root,
+            declared_specializations,
+            next_specialization_id,
+            resolved_specializations,
+        } = other;
+
+        for (ability, members) in other_members_of_ability.into_iter() {
+            if let Some(my_members) = self.members_of_ability(ability) {
+                debug_assert!(
+                    my_members == members,
+                    "Two abilities have different definitions, definitely a bug"
+                );
+            }
+            let member_data = members
+                .into_iter()
+                .map(|member| (member, other_ability_members.remove(&member).unwrap()));
+            self.register_ability(ability, member_data);
+        }
+
+        for (specialization, member) in specialization_to_root.into_iter() {
+            let old_root = self.specialization_to_root.insert(specialization, member);
+            debug_assert!(old_root.is_none() || old_root.unwrap() == member);
+        }
+
+        for ((member, typ), specialization) in declared_specializations.into_iter() {
+            let old_specialization = self
+                .declared_specializations
+                .insert((member, typ), specialization.clone());
+            debug_assert!(
+                old_specialization.is_none()
+                    || old_specialization.unwrap().symbol == specialization.symbol
+            );
+        }
+
+        debug_assert_eq!(next_specialization_id.get(), 1);
+        debug_assert_eq!(self.next_specialization_id.get(), 1);
+        debug_assert!(resolved_specializations.is_empty());
+        debug_assert!(self.resolved_specializations.is_empty());
+    }
+
+    pub fn resolve_for_module<Ctx, VarOfSymbol, ImportVar>(
+        self,
+        my_module: ModuleId,
+        my_module_ctx: &mut Ctx,
+        mut variable_of_symbol: VarOfSymbol,
+        mut import_lambda_set_var_from_module: ImportVar,
+    ) -> AbilitiesStore
+    where
+        VarOfSymbol: FnMut(&mut Ctx, Symbol) -> Variable,
+        ImportVar: FnMut(&mut Ctx, ModuleId, Variable) -> Variable,
+    {
+        let Self {
+            members_of_ability,
+            ability_members,
+            specialization_to_root,
+            declared_specializations,
+            next_specialization_id,
+            resolved_specializations,
+        } = self;
+
+        let ability_members = ability_members
+            .into_iter()
+            .map(|(member_symbol, member_data)| {
+                let AbilityMemberData {
+                    parent_ability,
+                    region,
+                    typ,
+                } = member_data;
+
+                let typ = match typ {
+                    PendingMemberType::Local {
+                        signature_var,
+                        signature: _,
+                        variables: _,
+                    } => ResolvedMemberType(signature_var),
+                    PendingMemberType::Imported => {
+                        ResolvedMemberType(variable_of_symbol(my_module_ctx, member_symbol))
+                    }
+                };
+
+                let member_data = AbilityMemberData {
+                    parent_ability,
+                    region,
+                    typ,
+                };
+
+                (member_symbol, member_data)
+            })
+            .collect();
+
+        let declared_specializations = declared_specializations
+            .into_iter()
+            .map(
+                |(
+                    key,
+                    MemberSpecialization {
+                        _phase,
+                        symbol,
+                        specialization_lambda_sets,
+                    },
+                )| {
+                    let symbol_module = symbol.module_id();
+                    // NOTE: this totally assumes we're dealing with subs that belong to an
+                    // individual module, things would be badly broken otherwise
+                    let member_specialization = if symbol_module == my_module {
+                        internal_error!("Ability store may only be pending before module solving, \
+                            so there shouldn't be any known module specializations at this point, but we found one for {:?}", symbol);
+                        // MemberSpecialization::new(symbol, specialization_lambda_sets)
+                    } else {
+                        let specialization_lambda_sets = specialization_lambda_sets
+                            .into_iter()
+                            .map(|(region, variable)| {
+                                (
+                                    region,
+                                    import_lambda_set_var_from_module(
+                                        my_module_ctx,
+                                        symbol_module,
+                                        variable,
+                                    ),
+                                )
+                            })
+                            .collect();
+                        MemberSpecialization::new(symbol, specialization_lambda_sets)
+                    };
+                    (key, member_specialization)
+                },
+            )
+            .collect();
+
+        AbilitiesStore {
+            members_of_ability,
+            ability_members,
+            specialization_to_root,
+            declared_specializations,
+            next_specialization_id,
+            resolved_specializations,
+        }
+    }
+}