Load derived member impls into derived module for mono

Cargo.lock

@@ -4700,6 +4700,7 @@ dependencies = [
  "roc_debug_flags",
  "roc_derive",
  "roc_derive_key",
+ "roc_late_solve",
  "roc_load_internal",
  "roc_module",
  "roc_region",

crates/compiler/derive/src/encoding.rs

@@ -11,7 +11,7 @@ use roc_derive_key::encoding::FlatEncodableKey;
 use roc_error_macros::internal_error;
 use roc_module::called_via::CalledVia;
 use roc_module::ident::Lowercase;
-use roc_module::symbol::{IdentIds, ModuleId, Symbol};
+use roc_module::symbol::{ModuleId, Symbol};
 use roc_region::all::{Loc, Region};
 use roc_types::subs::{
     instantiate_rigids, Content, ExhaustiveMark, FlatType, GetSubsSlice, LambdaSet, OptVariable,

crates/compiler/derive_key/src/lib.rs

@@ -102,6 +102,18 @@ impl DerivedSymbols {
         *symbol
     }
 
+    pub fn iter_all(&self) -> impl Iterator<Item = (&DeriveKey, &Symbol)> {
+        self.map.iter()
+    }
+
+    /// Generate a unique symbol. This should only be used when generating code inside the Derived
+    /// module; other modules should use [`Self::get_or_insert`] to generate a symbol for a derived
+    /// ability member usage.
+    pub fn gen_unique(&mut self) -> Symbol {
+        let ident_id = self.derived_ident_ids.gen_unique();
+        Symbol::new(ModuleId::DERIVED, ident_id)
+    }
+
     /// Steal all created derived ident Ids.
     /// After this is called, [`Self::get_or_insert`] may no longer be called.
     pub fn steal(&mut self) -> IdentIds {
@@ -110,11 +122,22 @@ impl DerivedSymbols {
 
         #[cfg(debug_assertions)]
         {
+            debug_assert!(!self.stolen);
             self.stolen = true;
         }
 
         ident_ids
     }
+
+    pub fn return_ident_ids(&mut self, ident_ids: IdentIds) {
+        #[cfg(debug_assertions)]
+        {
+            debug_assert!(self.stolen);
+            self.stolen = false;
+        }
+
+        self.derived_ident_ids = ident_ids;
+    }
 }
 
 /// Thread-sharable [`DerivedMethods`].

crates/compiler/late_solve/src/lib.rs

@@ -45,6 +45,28 @@ impl WorldAbilities {
         debug_assert!(old_store.is_none(), "{:?} abilities not new", module);
     }
 
+    #[inline(always)]
+    pub fn with_module_exposed_type<T>(
+        &mut self,
+        module: ModuleId,
+        mut f: impl FnMut(&mut ExposedTypesStorageSubs) -> T,
+    ) -> T {
+        let mut world = self.world.write().unwrap();
+        let (_, exposed_types) = world.get_mut(&module).expect("module not in the world");
+
+        f(exposed_types)
+    }
+
+    #[inline(always)]
+    pub fn with_module_abilities_store<T, F>(&self, module: ModuleId, mut f: F) -> T
+    where
+        F: FnMut(&AbilitiesStore) -> T,
+    {
+        let world = self.world.read().unwrap();
+        let (module_store, _module_types) = world.get(&module).unwrap();
+        f(module_store)
+    }
+
     pub fn clone_ref(&self) -> Self {
         Self {
             world: Arc::clone(&self.world),
@@ -53,7 +75,7 @@ impl WorldAbilities {
 }
 
 pub enum AbilitiesView<'a> {
-    World(WorldAbilities),
+    World(&'a WorldAbilities),
     Module(&'a AbilitiesStore),
 }
 
@@ -64,11 +86,7 @@ impl AbilitiesView<'_> {
         F: FnMut(&AbilitiesStore) -> T,
     {
         match self {
-            AbilitiesView::World(wa) => {
-                let world = wa.world.read().unwrap();
-                let (module_store, _module_types) = world.get(&module).unwrap();
-                f(module_store)
-            }
+            AbilitiesView::World(wa) => wa.with_module_abilities_store(module, f),
             AbilitiesView::Module(store) => f(store),
         }
     }

crates/compiler/load_internal/src/file.rs

@@ -131,6 +131,7 @@ struct ModuleCache<'a> {
     found_specializations: MutMap<ModuleId, FoundSpecializationsModule<'a>>,
     late_specializations: MutMap<ModuleId, LateSpecializationsModule<'a>>,
     external_specializations_requested: MutMap<ModuleId, Vec<ExternalSpecializations<'a>>>,
+    derives_module: Option<DerivesModule<'a>>,
 
     /// Various information
     imports: MutMap<ModuleId, MutSet<ModuleId>>,
@@ -179,6 +180,7 @@ impl Default for ModuleCache<'_> {
             found_specializations: Default::default(),
             late_specializations: Default::default(),
             external_specializations_requested: Default::default(),
+            derives_module: Default::default(),
             imports: Default::default(),
             top_level_thunks: Default::default(),
             documentation: Default::default(),
@@ -436,7 +438,37 @@ fn start_phase<'a>(
                     .unwrap_or_default();
 
                 let (ident_ids, subs, procs_base, layout_cache, module_timing) =
-                    if state.make_specializations_pass.current_pass() == 1 {
+                    if module_id == ModuleId::DERIVED {
+                        // The module for derives is treated specially - we keep it isolated, as it
+                        // is only needed for making specializations.
+                        let DerivesModule {
+                            mut layout_cache,
+                            mut procs_base,
+                            mut subs,
+                            mut module_timing,
+                        } = state
+                            .module_cache
+                            .derives_module
+                            .take()
+                            .unwrap_or_else(|| DerivesModule::new(state.target_info));
+
+                        load_derived_partial_procs(
+                            module_id,
+                            arena,
+                            &mut subs,
+                            &state.derived_symbols,
+                            &mut module_timing,
+                            &mut layout_cache,
+                            state.target_info,
+                            &state.exposed_to_host,
+                            &mut procs_base,
+                            &mut state.world_abilities,
+                        );
+
+                        let ident_ids = state.derived_symbols.lock().unwrap().steal();
+
+                        (ident_ids, subs, procs_base, layout_cache, module_timing)
+                    } else if state.make_specializations_pass.current_pass() == 1 {
                         let found_specializations = state
                             .module_cache
                             .found_specializations
@@ -616,6 +648,25 @@ struct LateSpecializationsModule<'a> {
     procs_base: ProcsBase<'a>,
 }
 
+#[derive(Debug)]
+struct DerivesModule<'a> {
+    layout_cache: LayoutCache<'a>,
+    procs_base: ProcsBase<'a>,
+    subs: Subs,
+    module_timing: ModuleTiming,
+}
+
+impl DerivesModule<'_> {
+    fn new(target_info: TargetInfo) -> Self {
+        Self {
+            layout_cache: LayoutCache::new(target_info),
+            procs_base: ProcsBase::default(),
+            subs: Subs::default(),
+            module_timing: ModuleTiming::new(SystemTime::now()),
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct MonomorphizedModule<'a> {
     pub module_id: ModuleId,
@@ -2387,16 +2438,36 @@ fn update<'a>(
             let _ = layout_cache;
 
             state.procedures.extend(procedures);
-            state.module_cache.late_specializations.insert(
-                module_id,
-                LateSpecializationsModule {
-                    ident_ids,
-                    module_timing,
-                    subs,
+            if module_id == ModuleId::DERIVED {
+                // The derives module is treated specially - put the data back and return the ident
+                // IDs to `derived_symbols` so other modules can use it in their monomorphization
+                // if needed.
+                state
+                    .derived_symbols
+                    .lock()
+                    .unwrap()
+                    .return_ident_ids(ident_ids);
+
+                debug_assert!(state.module_cache.derives_module.is_none());
+
+                state.module_cache.derives_module = Some(DerivesModule {
                     layout_cache,
                     procs_base,
-                },
-            );
+                    subs,
+                    module_timing,
+                });
+            } else {
+                state.module_cache.late_specializations.insert(
+                    module_id,
+                    LateSpecializationsModule {
+                        ident_ids,
+                        module_timing,
+                        subs,
+                        layout_cache,
+                        procs_base,
+                    },
+                );
+            }
 
             let work = state
                 .dependencies
@@ -4509,7 +4580,7 @@ fn make_specializations<'a>(
         update_mode_ids: &mut update_mode_ids,
         // call_specialization_counter=0 is reserved
         call_specialization_counter: 1,
-        abilities: AbilitiesView::World(world_abilities),
+        abilities: AbilitiesView::World(&world_abilities),
         derived_symbols: &derived_symbols,
     };
 

crates/compiler/load_internal/src/work.rs

@@ -377,7 +377,8 @@ impl<'a> Dependencies<'a> {
     pub fn reload_make_specialization_pass(&mut self) -> MutSet<(ModuleId, Phase)> {
         let mut output = MutSet::default();
 
-        let mut make_specializations_dependents = Default::default();
+        let mut make_specializations_dependents = MakeSpecializationsDependents::default();
+        let default_make_specializations_dependents_len = make_specializations_dependents.0.len();
         std::mem::swap(
             &mut self.make_specializations_dependents,
             &mut make_specializations_dependents,
@@ -403,7 +404,7 @@ impl<'a> Dependencies<'a> {
                 self.add_dependency(dependent, module, Phase::MakeSpecializations);
             }
 
-            self.add_to_status_for_all_phases(module, Phase::MakeSpecializations);
+            self.add_to_status_for_phase(module, Phase::MakeSpecializations);
             if !has_pred {
                 output.insert((module, Phase::MakeSpecializations));
             }
@@ -413,9 +414,11 @@ impl<'a> Dependencies<'a> {
             &mut self.make_specializations_dependents,
             &mut make_specializations_dependents,
         );
-        debug_assert!(
-            make_specializations_dependents.0.is_empty(),
-            "more modules were added to the graph"
+        debug_assert_eq!(
+            make_specializations_dependents.0.len(),
+            default_make_specializations_dependents_len,
+            "more modules were added to the graph: {:?}",
+            make_specializations_dependents
         );
 
         output

crates/compiler/mono/src/ir.rs

@@ -918,7 +918,7 @@ impl<'a> SymbolSpecializations<'a> {
     }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Default)]
 pub struct ProcsBase<'a> {
     pub partial_procs: BumpMap<Symbol, PartialProc<'a>>,
     pub module_thunks: &'a [Symbol],
@@ -1325,7 +1325,19 @@ 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(
+        if self.home == ModuleId::DERIVED {
+            // When specializing derives, we steal the Derived module's ident ids from
+            // `derived_symbols` for use in the usual mono pass. But during unification we
+            // temporarily return them, so that derived ability symbols are resolved properly in
+            // lambda sets.
+            let mut derived_ident_ids = IdentIds::default();
+            std::mem::swap(&mut derived_ident_ids, self.ident_ids);
+
+            let mut derived_symbols = self.derived_symbols.lock().unwrap();
+            derived_symbols.return_ident_ids(derived_ident_ids);
+        }
+
+        let result = roc_late_solve::unify(
             self.home,
             self.arena,
             self.subs,
@@ -1333,7 +1345,20 @@ impl<'a, 'i> Env<'a, 'i> {
             self.derived_symbols,
             left,
             right,
-        )
+        );
+
+        if self.home == ModuleId::DERIVED {
+            debug_assert!(
+                self.ident_ids.is_empty(),
+                "no ident ids should have been added while they were returned to derived_symbols"
+            );
+
+            let mut derived_symbols = self.derived_symbols.lock().unwrap();
+            let mut real_derived_ident_ids = derived_symbols.steal();
+            std::mem::swap(&mut real_derived_ident_ids, self.ident_ids);
+        }
+
+        result
     }
 }
 

crates/compiler/test_derive/Cargo.toml

@@ -23,6 +23,7 @@ roc_reporting = { path = "../../reporting" }
 roc_constrain = { path = "../constrain" }
 roc_region = { path = "../region" }
 roc_solve = { path = "../solve" }
+roc_late_solve = { path = "../late_solve" }
 roc_debug_flags = { path = "../debug_flags" }
 bumpalo = { version = "3.8.0", features = ["collections"] }
 lazy_static = "1.4.0"

crates/compiler/test_derive/src/encoding.rs

@@ -9,6 +9,7 @@ use std::path::PathBuf;
 use bumpalo::Bump;
 use insta::assert_snapshot;
 use pretty_assertions::assert_eq;
+use roc_late_solve::WorldAbilities;
 use ven_pretty::DocAllocator;
 
 use crate::pretty_print::{pretty_print_def, Ctx};