Compact lambda sets with structural derived ability impls

crates/compiler/derive/src/lib.rs

@@ -11,7 +11,9 @@ use roc_collections::MutMap;
 use roc_derive_key::DeriveKey;
 use roc_module::symbol::{IdentIds, ModuleId, Symbol};
 use roc_region::all::Loc;
-use roc_types::subs::{Content, Descriptor, Mark, OptVariable, Rank, Subs, Variable};
+use roc_types::subs::{
+    copy_import_to, Content, Descriptor, Mark, OptVariable, Rank, Subs, Variable,
+};
 
 mod encoding;
 
@@ -91,6 +93,7 @@ fn build_derived_body(
 impl DerivedModule {
     pub fn get_or_insert(
         &mut self,
+        // TODO: we only need "exposed by builtin modules that expose builtin abilities"
         exposed_by_module: &ExposedByModule,
         key: DeriveKey,
     ) -> &(Symbol, Def, SpecializationLambdaSets) {
@@ -136,6 +139,11 @@ impl DerivedModule {
     pub fn iter_all(
         &self,
     ) -> impl Iterator<Item = (&DeriveKey, &(Symbol, Def, SpecializationLambdaSets))> {
+        #[cfg(debug_assertions)]
+        {
+            debug_assert!(!self.stolen);
+        }
+
         self.map.iter()
     }
 
@@ -143,6 +151,11 @@ impl DerivedModule {
     /// 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 {
+        #[cfg(debug_assertions)]
+        {
+            debug_assert!(!self.stolen);
+        }
+
         let ident_id = self.derived_ident_ids.gen_unique();
         Symbol::new(DERIVED_MODULE, ident_id)
     }
@@ -176,6 +189,30 @@ impl DerivedModule {
         self.subs = subs;
         self.derived_ident_ids = ident_ids;
     }
+
+    pub fn copy_lambda_set_var_to_subs(&self, var: Variable, target: &mut Subs) -> Variable {
+        #[cfg(debug_assertions)]
+        {
+            debug_assert!(!self.stolen);
+        }
+
+        let copied_import = copy_import_to(
+            &self.subs,
+            target,
+            // bookkeep unspecialized lambda sets of var - I think we want this here
+            true,
+            var,
+            // TODO: I think this is okay because the only use of `copy_lambda_set_var_to_subs`
+            // (at least right now) is for lambda set compaction, which will automatically unify
+            // and lower ranks, and never generalize.
+            //
+            // However this is a bad coupling and maybe not a good assumption, we should revisit
+            // this when possible.
+            Rank::toplevel(),
+        );
+
+        copied_import.variable
+    }
 }
 
 /// Thread-sharable [`DerivedMethods`].

crates/compiler/late_solve/src/lib.rs

@@ -5,6 +5,7 @@ use std::sync::{Arc, RwLock};
 
 use bumpalo::Bump;
 use roc_can::abilities::AbilitiesStore;
+use roc_can::module::ExposedByModule;
 use roc_collections::MutMap;
 use roc_derive::SharedDerivedModule;
 use roc_error_macros::internal_error;
@@ -169,6 +170,7 @@ pub fn unify(
     subs: &mut Subs,
     abilities: &AbilitiesView,
     derived_module: &SharedDerivedModule,
+    exposed_by_module: &ExposedByModule,
     left: Variable,
     right: Variable,
 ) -> Result<(), UnificationFailed> {
@@ -190,6 +192,7 @@ pub fn unify(
                 &mut pools,
                 lambda_sets_to_specialize,
                 &late_phase,
+                exposed_by_module,
                 derived_module,
             );
             // At this point we can't do anything with must-implement constraints, since we're no

crates/compiler/load_internal/src/file.rs

@@ -427,6 +427,8 @@ fn start_phase<'a>(
                     ident_ids,
                     exposed_to_host: state.exposed_to_host.clone(),
                     abilities_store,
+                    // TODO: awful, how can we get rid of the clone?
+                    exposed_by_module: state.exposed_types.clone(),
                     derived_module,
                 }
             }
@@ -459,6 +461,7 @@ fn start_phase<'a>(
                             &mut layout_cache,
                             state.target_info,
                             &state.exposed_to_host,
+                            &state.exposed_types,
                             &mut procs_base,
                             &mut state.world_abilities,
                         );
@@ -526,6 +529,8 @@ fn start_phase<'a>(
                     specializations_we_must_make,
                     module_timing,
                     world_abilities: state.world_abilities.clone_ref(),
+                    // TODO: awful, how can we get rid of the clone?
+                    exposed_by_module: state.exposed_types.clone(),
                     derived_module,
                 }
             }
@@ -1050,6 +1055,7 @@ enum BuildTask<'a> {
         ident_ids: IdentIds,
         decls: Declarations,
         exposed_to_host: ExposedToHost,
+        exposed_by_module: ExposedByModule,
         abilities_store: AbilitiesStore,
         derived_module: SharedDerivedModule,
     },
@@ -1061,6 +1067,7 @@ enum BuildTask<'a> {
         layout_cache: LayoutCache<'a>,
         specializations_we_must_make: Vec<ExternalSpecializations<'a>>,
         module_timing: ModuleTiming,
+        exposed_by_module: ExposedByModule,
         world_abilities: WorldAbilities,
         derived_module: SharedDerivedModule,
     },
@@ -3932,7 +3939,7 @@ pub fn add_imports(
     my_module: ModuleId,
     subs: &mut Subs,
     mut pending_abilities: PendingAbilitiesStore,
-    mut exposed_for_module: ExposedForModule,
+    exposed_for_module: &ExposedForModule,
     def_types: &mut Vec<(Symbol, Loc<roc_types::types::Type>)>,
     rigid_vars: &mut Vec<Variable>,
 ) -> (Vec<Variable>, AbilitiesStore) {
@@ -3945,7 +3952,7 @@ pub fn add_imports(
     macro_rules! import_var_for_symbol  {
         ($subs:expr, $exposed_by_module:expr, $symbol:ident, $break:stmt) => {
             let module_id = $symbol.module_id();
-            match $exposed_by_module.get_mut(&module_id) {
+            match $exposed_by_module.get(&module_id) {
                 Some(ExposedModuleTypes {
                     exposed_types_storage_subs: exposed_types,
                     resolved_specializations: _,
@@ -3988,7 +3995,7 @@ pub fn add_imports(
         }
     }
 
-    for symbol in exposed_for_module.imported_values {
+    for &symbol in &exposed_for_module.imported_values {
         import_var_for_symbol!(subs, exposed_for_module.exposed_by_module, symbol, continue);
     }
 
@@ -4014,14 +4021,14 @@ pub fn add_imports(
 
     struct Ctx<'a> {
         subs: &'a mut Subs,
-        exposed_by_module: &'a mut ExposedByModule,
+        exposed_by_module: &'a ExposedByModule,
     }
 
     let abilities_store = pending_abilities.resolve_for_module(
         my_module,
         &mut Ctx {
             subs,
-            exposed_by_module: &mut exposed_for_module.exposed_by_module,
+            exposed_by_module: &exposed_for_module.exposed_by_module,
         },
         |ctx, symbol| match cached_symbol_vars.get(&symbol).copied() {
             Some(var) => var,
@@ -4035,7 +4042,7 @@ pub fn add_imports(
                 *cached_symbol_vars.get(&symbol).unwrap()
             }
         },
-        |ctx, module, lset_var| match ctx.exposed_by_module.get_mut(&module) {
+        |ctx, module, lset_var| match ctx.exposed_by_module.get(&module) {
             Some(ExposedModuleTypes {
                 exposed_types_storage_subs: exposed_types,
                 resolved_specializations: _,
@@ -4091,7 +4098,7 @@ fn run_solve_solve(
         module.module_id,
         &mut subs,
         pending_abilities,
-        exposed_for_module,
+        &exposed_for_module,
         &mut def_types,
         &mut rigid_vars,
     );
@@ -4114,6 +4121,7 @@ fn run_solve_solve(
             solve_aliases,
             abilities_store,
             pending_derives,
+            &exposed_for_module.exposed_by_module,
             derived_module,
         );
 
@@ -4567,6 +4575,7 @@ fn make_specializations<'a>(
     mut module_timing: ModuleTiming,
     target_info: TargetInfo,
     world_abilities: WorldAbilities,
+    exposed_by_module: &ExposedByModule,
     derived_module: SharedDerivedModule,
 ) -> Msg<'a> {
     let make_specializations_start = SystemTime::now();
@@ -4582,6 +4591,7 @@ fn make_specializations<'a>(
         // call_specialization_counter=0 is reserved
         call_specialization_counter: 1,
         abilities: AbilitiesView::World(&world_abilities),
+        exposed_by_module,
         derived_module: &derived_module,
     };
 
@@ -4675,6 +4685,7 @@ fn build_pending_specializations<'a>(
         // to know the types and abilities in our modules. Only for building *all* specializations
         // do we need a global view.
         abilities: AbilitiesView::Module(&abilities_store),
+        exposed_by_module,
         derived_module: &derived_module,
     };
 
@@ -5057,6 +5068,7 @@ fn run_task<'a>(
             imported_module_thunks,
             exposed_to_host,
             abilities_store,
+            exposed_by_module,
             derived_module,
         } => Ok(build_pending_specializations(
             arena,
@@ -5069,6 +5081,7 @@ fn run_task<'a>(
             layout_cache,
             target_info,
             exposed_to_host,
+            &exposed_by_module,
             abilities_store,
             derived_module,
         )),
@@ -5081,6 +5094,7 @@ fn run_task<'a>(
             specializations_we_must_make,
             module_timing,
             world_abilities,
+            exposed_by_module,
             derived_module,
         } => Ok(make_specializations(
             arena,
@@ -5093,6 +5107,7 @@ fn run_task<'a>(
             module_timing,
             target_info,
             world_abilities,
+            &exposed_by_module,
             derived_module,
         )),
     }?;

crates/compiler/mono/src/ir.rs

@@ -9,6 +9,7 @@ use bumpalo::Bump;
 use roc_builtins::bitcode::{FloatWidth, IntWidth};
 use roc_can::abilities::SpecializationId;
 use roc_can::expr::{AnnotatedMark, ClosureData, IntValue};
+use roc_can::module::ExposedByModule;
 use roc_collections::all::{default_hasher, BumpMap, BumpMapDefault, MutMap};
 use roc_collections::VecMap;
 use roc_debug_flags::dbg_do;
@@ -1293,7 +1294,9 @@ pub struct Env<'a, 'i> {
     pub target_info: TargetInfo,
     pub update_mode_ids: &'i mut UpdateModeIds,
     pub call_specialization_counter: u32,
+    // TODO: WorldAbilities and exposed_by_module share things, think about how to combine them
     pub abilities: AbilitiesView<'i>,
+    pub exposed_by_module: &'i ExposedByModule,
     pub derived_module: &'i SharedDerivedModule,
 }
 
@@ -1356,6 +1359,7 @@ impl<'a, 'i> Env<'a, 'i> {
             self.subs,
             &self.abilities,
             self.derived_module,
+            &self.exposed_by_module,
             left,
             right,
         );

crates/compiler/solve/src/module.rs

@@ -2,7 +2,7 @@ use crate::solve::{self, Aliases};
 use roc_can::abilities::{AbilitiesStore, ResolvedSpecializations};
 use roc_can::constraint::{Constraint as ConstraintSoa, Constraints};
 use roc_can::expr::PendingDerives;
-use roc_can::module::RigidVariables;
+use roc_can::module::{ExposedByModule, RigidVariables};
 use roc_collections::all::MutMap;
 use roc_collections::VecMap;
 use roc_derive_key::GlobalDerivedSymbols;
@@ -61,6 +61,7 @@ pub fn run_solve(
     mut aliases: Aliases,
     mut abilities_store: AbilitiesStore,
     pending_derives: PendingDerives,
+    exposed_by_module: &ExposedByModule,
     derived_module: SharedDerivedModule,
 ) -> (
     Solved<Subs>,
@@ -93,6 +94,7 @@ pub fn run_solve(
         &constraint,
         pending_derives,
         &mut abilities_store,
+        exposed_by_module,
         derived_module,
     );
 

crates/compiler/solve/src/solve.rs

@@ -9,6 +9,7 @@ use roc_can::constraint::Constraint::{self, *};
 use roc_can::constraint::{Constraints, Cycle, LetConstraint, OpportunisticResolve};
 use roc_can::expected::{Expected, PExpected};
 use roc_can::expr::PendingDerives;
+use roc_can::module::ExposedByModule;
 use roc_collections::all::MutMap;
 use roc_debug_flags::dbg_do;
 #[cfg(debug_assertions)]
@@ -551,6 +552,7 @@ pub fn run(
     constraint: &Constraint,
     pending_derives: PendingDerives,
     abilities_store: &mut AbilitiesStore,
+    exposed_by_module: &ExposedByModule,
     derived_module: SharedDerivedModule,
 ) -> (Solved<Subs>, Env) {
     let env = run_in_place(
@@ -561,6 +563,7 @@ pub fn run(
         constraint,
         pending_derives,
         abilities_store,
+        exposed_by_module,
         derived_module,
     );
 
@@ -577,6 +580,7 @@ fn run_in_place(
     constraint: &Constraint,
     pending_derives: PendingDerives,
     abilities_store: &mut AbilitiesStore,
+    exposed_by_module: &ExposedByModule,
     derived_module: SharedDerivedModule,
 ) -> Env {
     let mut pools = Pools::default();
@@ -619,6 +623,7 @@ fn run_in_place(
         &mut pools,
         deferred_uls_to_resolve,
         &SolvePhase { abilities_store },
+        exposed_by_module,
         &derived_module,
     );
 
@@ -1881,7 +1886,8 @@ pub fn compact_lambda_sets_of_vars<P: Phase>(
     pools: &mut Pools,
     uls_of_var: UlsOfVar,
     phase: &P,
-    derived_symbols: &GlobalDerivedSymbols,
+    exposed_by_module: &ExposedByModule,
+    derived_module: &SharedDerivedModule,
 ) -> MustImplementConstraints {
     // let mut seen = VecSet::default();
     let mut must_implement = MustImplementConstraints::default();
@@ -2306,6 +2312,8 @@ fn get_specialization_lambda_set<P: Phase>(
     ability_member: Symbol,
     lset_region: u8,
     specialization_key: SpecializationTypeKey,
+    exposed_by_module: &ExposedByModule,
+    derived_module: &SharedDerivedModule,
 ) -> Result<Variable, ()> {
     match specialization_key {
         SpecializationTypeKey::Opaque(opaque) => {
@@ -2345,7 +2353,21 @@ fn get_specialization_lambda_set<P: Phase>(
             Ok(local_lset)
         }
 
-        SpecializationTypeKey::Derived(_) => todo!(),
+        SpecializationTypeKey::Derived(derive_key) => {
+            let mut derived_module = derived_module.lock().unwrap();
+
+            let (_, _, specialization_lambda_sets) =
+                derived_module.get_or_insert(exposed_by_module, derive_key);
+
+            let &specialized_lambda_set = specialization_lambda_sets
+                .get(&lset_region)
+                .expect("lambda set region not resolved");
+
+            let local_lset =
+                derived_module.copy_lambda_set_var_to_subs(specialized_lambda_set, subs);
+
+            Ok(local_lset)
+        }
     }
 }