Consolidate exposed types and derived module in a derived environment

crates/compiler/late_solve/src/lib.rs

@@ -11,7 +11,7 @@ use roc_derive::SharedDerivedModule;
 use roc_error_macros::internal_error;
 use roc_module::symbol::ModuleId;
 use roc_solve::solve::Pools;
-use roc_solve::specialize::{compact_lambda_sets_of_vars, Phase};
+use roc_solve::specialize::{compact_lambda_sets_of_vars, DerivedEnv, Phase};
 use roc_types::subs::{get_member_lambda_sets_at_region, Content, FlatType, LambdaSet};
 use roc_types::subs::{ExposedTypesStorageSubs, Subs, Variable};
 use roc_unify::unify::{unify as unify_unify, Env, Mode, Unified};
@@ -273,15 +273,18 @@ pub fn unify(
             let mut pools = Pools::default();
 
             let late_phase = LatePhase { home, abilities };
+            let derived_env = DerivedEnv {
+                derived_module,
+                exposed_types: exposed_by_module,
+            };
 
             let must_implement_constraints = compact_lambda_sets_of_vars(
                 subs,
-                derived_module,
+                &derived_env,
                 arena,
                 &mut pools,
                 lambda_sets_to_specialize,
                 &late_phase,
-                exposed_by_module,
             );
             // At this point we can't do anything with must-implement constraints, since we're no
             // longer solving. We must assume that they were totally caught during solving.

crates/compiler/solve/src/solve.rs

@@ -3,7 +3,7 @@ use crate::ability::{
     CheckedDerives, ObligationCache, PendingDerivesTable, Resolved,
 };
 use crate::module::Solved;
-use crate::specialize::{compact_lambda_sets_of_vars, SolvePhase};
+use crate::specialize::{compact_lambda_sets_of_vars, DerivedEnv, SolvePhase};
 use bumpalo::Bump;
 use roc_can::abilities::{AbilitiesStore, MemberSpecializationInfo};
 use roc_can::constraint::Constraint::{self, *};
@@ -588,14 +588,17 @@ fn run_in_place(
     // Now that the module has been solved, we can run through and check all
     // types claimed to implement abilities. This will also tell us what derives
     // are legal, which we need to register.
+    let derived_env = DerivedEnv {
+        derived_module: &derived_module,
+        exposed_types: exposed_by_module,
+    };
     let new_must_implement = compact_lambda_sets_of_vars(
         subs,
-        &derived_module,
+        &derived_env,
         &arena,
         &mut pools,
         deferred_uls_to_resolve,
         &SolvePhase { abilities_store },
-        exposed_by_module,
     );
     problems.extend(obligation_cache.check_obligations(
         subs,

crates/compiler/solve/src/specialize.rs

@@ -112,6 +112,12 @@ impl Phase for SolvePhase<'_> {
     }
 }
 
+pub struct DerivedEnv<'a> {
+    pub derived_module: &'a SharedDerivedModule,
+    /// Exposed types needed by the derived module.
+    pub exposed_types: &'a ExposedByModule,
+}
+
 #[cfg(debug_assertions)]
 fn trace_compaction_step_1(subs: &Subs, c_a: Variable, uls_a: &[Variable]) {
     let c_a = roc_types::subs::SubsFmtContent(subs.get_content_without_compacting(c_a), subs);
@@ -235,12 +241,11 @@ fn unique_unspecialized_lambda(subs: &Subs, c_a: Variable, uls: &[Uls]) -> Optio
 #[must_use]
 pub fn compact_lambda_sets_of_vars<P: Phase>(
     subs: &mut Subs,
-    derived_module: &SharedDerivedModule,
+    derived_env: &DerivedEnv,
     arena: &Bump,
     pools: &mut Pools,
     uls_of_var: UlsOfVar,
     phase: &P,
-    exposed_by_module: &ExposedByModule,
 ) -> MustImplementConstraints {
     // let mut seen = VecSet::default();
     let mut must_implement = MustImplementConstraints::default();
@@ -365,16 +370,8 @@ pub fn compact_lambda_sets_of_vars<P: Phase>(
             //     continue;
             // }
 
-            let (new_must_implement, new_uls_of_var) = compact_lambda_set(
-                subs,
-                derived_module,
-                arena,
-                pools,
-                c_a,
-                l,
-                phase,
-                exposed_by_module,
-            );
+            let (new_must_implement, new_uls_of_var) =
+                compact_lambda_set(subs, derived_env, arena, pools, c_a, l, phase);
 
             must_implement.extend(new_must_implement);
             uls_of_var_queue.extend(new_uls_of_var.drain());
@@ -390,13 +387,12 @@ pub fn compact_lambda_sets_of_vars<P: Phase>(
 #[allow(clippy::too_many_arguments)]
 fn compact_lambda_set<P: Phase>(
     subs: &mut Subs,
-    derived_module: &SharedDerivedModule,
+    derived_env: &DerivedEnv,
     arena: &Bump,
     pools: &mut Pools,
     resolved_concrete: Variable,
     this_lambda_set: Variable,
     phase: &P,
-    exposed_by_module: &ExposedByModule,
 ) -> (MustImplementConstraints, UlsOfVar) {
     // 3. For each `l` in `uls_a` with unique unspecialized lambda `C:f:r`:
     //    1. Let `t_f1` be the directly ambient function of the lambda set containing `C:f:r`. Remove `C:f:r` from `t_f1`'s lambda set.
@@ -456,12 +452,11 @@ fn compact_lambda_set<P: Phase>(
 
     let specialization_ambient_function_var = get_specialization_lambda_set_ambient_function(
         subs,
-        derived_module,
+        derived_env,
         phase,
         f,
         r,
         specialization_key,
-        exposed_by_module,
         target_rank,
     );
 
@@ -550,12 +545,11 @@ fn make_specialization_decision(subs: &Subs, var: Variable) -> SpecializeDecisio
 #[allow(clippy::too_many_arguments)]
 fn get_specialization_lambda_set_ambient_function<P: Phase>(
     subs: &mut Subs,
-    derived_module: &SharedDerivedModule,
+    derived_env: &DerivedEnv,
     phase: &P,
     ability_member: Symbol,
     lset_region: u8,
     specialization_key: SpecializationTypeKey,
-    exposed_by_module: &ExposedByModule,
     target_rank: Rank,
 ) -> Result<Variable, ()> {
     match specialization_key {
@@ -608,10 +602,10 @@ fn get_specialization_lambda_set_ambient_function<P: Phase>(
         }
 
         SpecializationTypeKey::Derived(derive_key) => {
-            let mut derived_module = derived_module.lock().unwrap();
+            let mut derived_module = derived_env.derived_module.lock().unwrap();
 
             let (_, _, specialization_lambda_sets) =
-                derived_module.get_or_insert(exposed_by_module, derive_key);
+                derived_module.get_or_insert(derived_env.exposed_types, derive_key);
 
             let specialized_lambda_set = *specialization_lambda_sets
                 .get(&lset_region)