World abilities: big world vs tiny world

Normally a module can see all of its abilities, and any module it
depends on (this is the tiny world, i.e. only the world it is aware of).
But when making specializations, a module may need to see all abilities
in all modules (big world). This patch supports big-world viewing of
abilities.

compiler/late_solve/src/lib.rs

@@ -49,4 +49,4 @@ pub fn unify(
 pub use roc_solve::solve::instantiate_rigids;
 
 pub use roc_solve::ability::resolve_ability_specialization;
-pub use roc_solve::ability::Resolved;
+pub use roc_solve::ability::{Resolved, WorldAbilities};

compiler/load_internal/src/file.rs

@@ -41,6 +41,7 @@ use roc_parse::module::module_defs;
 use roc_parse::parser::{FileError, Parser, SyntaxError};
 use roc_region::all::{LineInfo, Loc, Region};
 use roc_reporting::report::RenderTarget;
+use roc_solve::ability::{AllModuleAbilities, WorldAbilities};
 use roc_solve::module::SolvedModule;
 use roc_solve::solve;
 use roc_target::TargetInfo;
@@ -446,6 +447,13 @@ fn start_phase<'a>(
                     abilities_store,
                 } = found_specializations;
 
+                let old_store = &state
+                    .inverse_solved_abilities
+                    .write()
+                    .unwrap()
+                    .insert(module_id, abilities_store);
+                debug_assert!(old_store.is_none(), "{:?} abilities not new", module_id);
+
                 BuildTask::MakeSpecializations {
                     module_id,
                     ident_ids,
@@ -454,7 +462,7 @@ fn start_phase<'a>(
                     layout_cache,
                     specializations_we_must_make,
                     module_timing,
-                    abilities_store,
+                    world_abilities: Arc::clone(&state.inverse_solved_abilities),
                 }
             }
         }
@@ -745,6 +753,9 @@ struct State<'a> {
 
     pub render: RenderTarget,
 
+    /// At a given point in time, this map contains the abilities stores in reverse order
+    pub inverse_solved_abilities: AllModuleAbilities,
+
     // cached subs (used for builtin modules, could include packages in the future too)
     cached_subs: CachedSubs,
 }
@@ -789,6 +800,7 @@ impl<'a> State<'a> {
             layout_caches: std::vec::Vec::with_capacity(number_of_workers),
             cached_subs: Arc::new(Mutex::new(cached_subs)),
             render,
+            inverse_solved_abilities: Default::default(),
         }
     }
 }
@@ -916,7 +928,7 @@ enum BuildTask<'a> {
         layout_cache: LayoutCache<'a>,
         specializations_we_must_make: Vec<ExternalSpecializations>,
         module_timing: ModuleTiming,
-        abilities_store: AbilitiesStore,
+        world_abilities: AllModuleAbilities,
     },
 }
 
@@ -4166,7 +4178,7 @@ fn make_specializations<'a>(
     specializations_we_must_make: Vec<ExternalSpecializations>,
     mut module_timing: ModuleTiming,
     target_info: TargetInfo,
-    abilities_store: AbilitiesStore,
+    world_abilities: AllModuleAbilities,
 ) -> Msg<'a> {
     let make_specializations_start = SystemTime::now();
     let mut update_mode_ids = UpdateModeIds::new();
@@ -4180,7 +4192,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,
+        abilities: WorldAbilities::BigWorld(world_abilities),
     };
 
     let mut procs = Procs::new_in(arena);
@@ -4275,7 +4287,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,
+        abilities: WorldAbilities::TinyWorld(&abilities_store),
     };
 
     // Add modules' decls to Procs
@@ -4587,7 +4599,7 @@ fn run_task<'a>(
             layout_cache,
             specializations_we_must_make,
             module_timing,
-            abilities_store,
+            world_abilities,
         } => Ok(make_specializations(
             arena,
             module_id,
@@ -4598,7 +4610,7 @@ fn run_task<'a>(
             specializations_we_must_make,
             module_timing,
             target_info,
-            abilities_store,
+            world_abilities,
         )),
     }?;
 

compiler/mono/src/ir.rs

@@ -7,7 +7,7 @@ use crate::layout::{
 use bumpalo::collections::{CollectIn, Vec};
 use bumpalo::Bump;
 use roc_builtins::bitcode::{FloatWidth, IntWidth};
-use roc_can::abilities::{AbilitiesStore, SpecializationId};
+use roc_can::abilities::SpecializationId;
 use roc_can::expr::{AnnotatedMark, ClosureData, IntValue};
 use roc_collections::all::{default_hasher, BumpMap, BumpMapDefault, MutMap};
 use roc_collections::VecMap;
@@ -19,7 +19,7 @@ use roc_debug_flags::{
 use roc_error_macros::todo_abilities;
 use roc_exhaustive::{Ctor, CtorName, Guard, RenderAs, TagId};
 use roc_late_solve::{
-    instantiate_rigids, resolve_ability_specialization, Resolved, UnificationFailed,
+    instantiate_rigids, resolve_ability_specialization, Resolved, UnificationFailed, WorldAbilities,
 };
 use roc_module::ident::{ForeignSymbol, Lowercase, TagName};
 use roc_module::low_level::LowLevel;
@@ -1242,7 +1242,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,
+    pub abilities: WorldAbilities<'i>,
 }
 
 impl<'a, 'i> Env<'a, 'i> {
@@ -1273,7 +1273,9 @@ impl<'a, 'i> Env<'a, 'i> {
     /// Unifies two variables and performs lambda set compaction.
     /// Use this rather than [roc_unify::unify] directly!
     fn unify(&mut self, left: Variable, right: Variable) -> Result<(), UnificationFailed> {
-        roc_late_solve::unify(self.arena, self.subs, self.abilities_store, left, right)
+        self.abilities.with_module_store(self.home, |store| {
+            roc_late_solve::unify(self.arena, self.subs, store, left, right)
+        })
     }
 }
 
@@ -3742,10 +3744,11 @@ pub fn with_hole<'a>(
             specialize_naked_symbol(env, variable, procs, layout_cache, assigned, hole, symbol)
         }
         AbilityMember(_member, specialization_id, _) => {
-            let specialization_symbol = env
+            let specialization_symbol = env.abilities.with_module_store(env.home, |store| {
-                .abilities_store
+                store
-                .get_resolved(specialization_id)
+                    .get_resolved(specialization_id)
-                .expect("Specialization was never made!");
+                    .expect("Specialization was never made!")
+            });
 
             specialize_naked_symbol(
                 env,
@@ -4752,7 +4755,10 @@ pub fn with_hole<'a>(
                         UnspecializedExpr(symbol) => {
                             match procs.ability_member_aliases.get(symbol).unwrap() {
                                 &self::AbilityMember(member) => {
-                                    let resolved_proc = resolve_ability_specialization(env.subs, env.abilities_store, member, fn_var).expect("Recorded as an ability member, but it doesn't have a specialization");
+                                    let resolved_proc = env.abilities.with_module_store(env.home, |store| 
+                                        resolve_ability_specialization(env.subs, store, member, fn_var)
+                                            .expect("Recorded as an ability member, but it doesn't have a specialization")
+                                    );
 
                                     let resolved_proc = match resolved_proc {
                                         Resolved::Specialization(symbol) => symbol,
@@ -5127,7 +5133,11 @@ fn late_resolve_ability_specialization<'a>(
     specialization_id: SpecializationId,
     specialization_var: Variable,
 ) -> Symbol {
-    if let Some(spec_symbol) = env.abilities_store.get_resolved(specialization_id) {
+    let opt_resolved = env
+        .abilities
+        .with_module_store(env.home, |store| store.get_resolved(specialization_id));
+
+    if let Some(spec_symbol) = opt_resolved {
         // Fast path: specialization is monomorphic, was found during solving.
         spec_symbol
     } else if let Content::Structure(FlatType::Func(_, lambda_set, _)) =
@@ -5148,13 +5158,10 @@ fn late_resolve_ability_specialization<'a>(
         env.subs[spec_symbol_index]
     } else {
         // Otherwise, resolve by checking the able var.
-        let specialization = resolve_ability_specialization(
+        let specialization = env.abilities.with_module_store(env.home, |store| {
-            env.subs,
+            resolve_ability_specialization(env.subs, store, member, specialization_var)
-            env.abilities_store,
+                .expect("Ability specialization is unknown - code generation cannot proceed!")
-            member,
+        });
-            specialization_var,
-        )
-        .expect("Ability specialization is unknown - code generation cannot proceed!");
 
         match specialization {
             Resolved::Specialization(symbol) => symbol,
@@ -6987,7 +6994,11 @@ where
     // 1. Handle references to ability members - we could be aliasing an ability member, or another
     //    alias to an ability member.
     {
-        if env.abilities_store.is_ability_member_name(right) {
+        let is_ability_member = env
+            .abilities
+            .with_module_store(env.home, |store| store.is_ability_member_name(right));
+
+        if is_ability_member {
             procs
                 .ability_member_aliases
                 .insert(left, AbilityMember(right));

compiler/solve/src/ability.rs

@@ -1,8 +1,10 @@
+use std::sync::{Arc, RwLock};
+
 use roc_can::abilities::AbilitiesStore;
 use roc_can::expr::PendingDerives;
-use roc_collections::VecMap;
+use roc_collections::{MutMap, VecMap};
 use roc_error_macros::internal_error;
-use roc_module::symbol::Symbol;
+use roc_module::symbol::{ModuleId, Symbol};
 use roc_region::all::{Loc, Region};
 use roc_types::subs::{Content, FlatType, GetSubsSlice, Rank, Subs, Variable};
 use roc_types::types::{AliasKind, Category, ErrorType, PatternCategory};
@@ -12,6 +14,29 @@ use roc_unify::unify::{MustImplementAbility, Obligated};
 use crate::solve::{instantiate_rigids, type_to_var};
 use crate::solve::{Aliases, Pools, TypeError};
 
+pub type AllModuleAbilities = Arc<RwLock<MutMap<ModuleId, AbilitiesStore>>>;
+
+pub enum WorldAbilities<'a> {
+    BigWorld(AllModuleAbilities),
+    TinyWorld(&'a AbilitiesStore),
+}
+
+impl WorldAbilities<'_> {
+    pub fn with_module_store<T, F>(&self, module: ModuleId, mut f: F) -> T
+    where
+        F: FnMut(&AbilitiesStore) -> T,
+    {
+        match self {
+            WorldAbilities::BigWorld(world) => {
+                let world = world.read().unwrap();
+                let module_store = world.get(&module).unwrap();
+                f(module_store)
+            }
+            WorldAbilities::TinyWorld(store) => f(store),
+        }
+    }
+}
+
 #[derive(Debug, Clone)]
 pub enum AbilityImplError {
     /// Promote this to an error that the type does not fully implement an ability