Properly handle immediate derivations

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::{compact_lambda_sets_of_vars, Phase, Pools};
-use roc_types::subs::{Content, FlatType, LambdaSet};
+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, Mode, Unified};
 
@@ -133,7 +133,7 @@ impl Phase for LatePhase<'_> {
                 let (_module_store, module_types) = world.get_mut(&external_module_id).unwrap();
 
                 let storage_lambda_set_var = *module_types
-                    .stored_ability_lambda_set_vars
+                    .stored_specialization_lambda_set_vars
                     .get(&external_lambda_set_var)
                     .unwrap();
                 let LambdaSet {
@@ -160,6 +160,86 @@ impl Phase for LatePhase<'_> {
             }
         }
     }
+
+    #[inline(always)]
+    fn get_and_copy_ability_member_ambient_function(
+        &self,
+        ability_member: roc_module::symbol::Symbol,
+        region: u8,
+        target_subs: &mut Subs,
+    ) -> Variable {
+        match self.abilities {
+            AbilitiesView::Module(abilities_store) => {
+                // No option other than that the var must be in our module store.
+                // Even if the specialization lambda set comes from another module,
+                // we should have taken care to import it before starting solving in this module.
+                let member_def = abilities_store
+                    .member_def(ability_member)
+                    .unwrap_or_else(|| {
+                        internal_error!(
+                            "{:?} is not resolved, or not an ability member!",
+                            ability_member
+                        )
+                    });
+                let member_var = member_def.signature_var();
+
+                let region_lset = get_member_lambda_sets_at_region(target_subs, member_var, region);
+
+                let LambdaSet {
+                    ambient_function, ..
+                } = target_subs.get_lambda_set(region_lset);
+
+                ambient_function
+            }
+            AbilitiesView::World(wa) => {
+                let member_home = ability_member.module_id();
+                let mut world = wa.world.write().unwrap();
+                let (module_store, module_types) = world.get_mut(&member_home).unwrap();
+
+                let member_def = module_store.member_def(ability_member).unwrap_or_else(|| {
+                    internal_error!(
+                        "{:?} is not resolved, or not an ability member!",
+                        ability_member
+                    )
+                });
+                let member_var = member_def.signature_var();
+                let storage_member_var = module_types
+                    .stored_ability_member_vars
+                    .get(&member_var)
+                    .unwrap();
+
+                let storage_lambda_set_var = get_member_lambda_sets_at_region(
+                    module_types.storage_subs.as_inner(),
+                    *storage_member_var,
+                    region,
+                );
+
+                let LambdaSet {
+                    ambient_function, ..
+                } = module_types
+                    .storage_subs
+                    .as_inner()
+                    .get_lambda_set(storage_lambda_set_var);
+
+                let copied = module_types
+                    .storage_subs
+                    // TODO: I think this is always okay, but revisit later when we're in a more
+                    // stable position to see if we can get rid of the bookkeeping done as a result
+                    // of this.
+                    .export_variable_to_directly_to_use_site(target_subs, ambient_function);
+
+                let our_ambient_function_var = copied.variable;
+                instantiate_rigids(target_subs, our_ambient_function_var);
+
+                debug_assert!(matches!(
+                    target_subs.get_content_without_compacting(our_ambient_function_var),
+                    Content::Structure(FlatType::Func(..))
+                ));
+
+                our_ambient_function_var
+            }
+        }
+    }
 }
 
 /// Unifies two variables and performs lambda set compaction.

crates/compiler/load_internal/src/file.rs

@@ -4076,7 +4076,7 @@ pub fn add_imports(
                 resolved_specializations: _,
             }) => {
                 let var = exposed_types
-                    .stored_ability_lambda_set_vars
+                    .stored_specialization_lambda_set_vars
                     .get(&lset_var)
                     .expect("Lambda set var from other module not available");
 

crates/compiler/solve/src/module.rs

@@ -8,10 +8,7 @@ use roc_collections::VecMap;
 use roc_derive::SharedDerivedModule;
 use roc_error_macros::internal_error;
 use roc_module::symbol::{ModuleId, Symbol};
-use roc_types::subs::{
+use roc_types::subs::{Content, ExposedTypesStorageSubs, FlatType, StorageSubs, Subs, Variable};
-    get_member_lambda_sets_at_region, Content, ExposedTypesStorageSubs, FlatType, StorageSubs,
-    Subs, Variable,
-};
 use roc_types::types::Alias;
 
 /// A marker that a given Subs has been solved.
@@ -117,68 +114,60 @@ pub fn exposed_types_storage_subs(
     let subs = solved_subs.inner_mut();
     let mut storage_subs = StorageSubs::new(Subs::new());
     let mut stored_vars_by_symbol = VecMap::with_capacity(exposed_vars_by_symbol.len());
-    let mut stored_ability_lambda_set_vars = VecMap::with_capacity(solved_specializations.len());
 
     for (symbol, var) in exposed_vars_by_symbol.iter() {
         let new_var = storage_subs.import_variable_from(subs, *var).variable;
         stored_vars_by_symbol.insert(*symbol, new_var);
     }
 
-    // Store all specialization lambda sets solved thanks to this module
+    let mut stored_specialization_lambda_set_vars =
-    let solved_specialization_lambda_sets =
+        VecMap::with_capacity(solved_specializations.len());
-        solved_specializations
+
-            .iter()
+    for (_, member_specialization) in solved_specializations.iter() {
-            .flat_map(|(_, member_specialization)| {
+        for (_, &lset_var) in member_specialization.specialization_lambda_sets.iter() {
-                member_specialization
+            let specialization_lset_ambient_function_var =
-                    .specialization_lambda_sets
+                subs.get_lambda_set(lset_var).ambient_function;
-                    .iter()
+
-                    .map(|(_region, var)| *var)
+            // Import the ambient function of this specialization lambda set; that will import the
-            });
+            // lambda set as well. The ambient function is needed for the lambda set compaction
+            // algorithm.
+            let imported_lset_ambient_function_var = storage_subs
+                .import_variable_from(subs, specialization_lset_ambient_function_var)
+                .variable;
+
+            let imported_lset_var = match storage_subs
+                .as_inner()
+                .get_content_without_compacting(imported_lset_ambient_function_var)
+            {
+                Content::Structure(FlatType::Func(_, lambda_set_var, _)) => *lambda_set_var,
+                content => internal_error!(
+                    "ambient lambda set function import is not a function, found: {:?}",
+                    roc_types::subs::SubsFmtContent(content, storage_subs.as_inner())
+                ),
+            };
+            stored_specialization_lambda_set_vars.insert(lset_var, imported_lset_var);
+        }
+    }
 
     // Store the regioned lambda sets of the ability members defined in this module.
-    let defined_member_lambda_sets = abilities_store
+    let stored_ability_member_vars = abilities_store
         .root_ability_members()
         .iter()
         .filter_map(|(member, data)| {
             if member.module_id() == home {
-                Some(get_member_lambda_sets_at_region(
+                let var = data.signature_var();
-                    subs,
+                let improted_var = storage_subs.import_variable_from(subs, var).variable;
-                    data.signature_var(),
+                Some((var, improted_var))
-                    None,
-                ))
             } else {
                 None
             }
         })
-        .flatten();
+        .collect();
-
-    for lset_var in solved_specialization_lambda_sets.chain(defined_member_lambda_sets) {
-        let specialization_lset_ambient_function_var =
-            subs.get_lambda_set(lset_var).ambient_function;
-
-        // Import the ambient function of this specialization lambda set; that will import the
-        // lambda set as well. The ambient function is needed for the lambda set compaction
-        // algorithm.
-        let imported_lset_ambient_function_var = storage_subs
-            .import_variable_from(subs, specialization_lset_ambient_function_var)
-            .variable;
-
-        let imported_lset_var = match storage_subs
-            .as_inner()
-            .get_content_without_compacting(imported_lset_ambient_function_var)
-        {
-            Content::Structure(FlatType::Func(_, lambda_set_var, _)) => *lambda_set_var,
-            content => internal_error!(
-                "ambient lambda set function import is not a function, found: {:?}",
-                roc_types::subs::SubsFmtContent(content, storage_subs.as_inner())
-            ),
-        };
-        stored_ability_lambda_set_vars.insert(lset_var, imported_lset_var);
-    }
 
     ExposedTypesStorageSubs {
         storage_subs,
         stored_vars_by_symbol,
-        stored_ability_lambda_set_vars,
+        stored_specialization_lambda_set_vars,
+        stored_ability_member_vars,
     }
 }

crates/compiler/solve/src/solve.rs

@@ -22,9 +22,9 @@ use roc_module::symbol::{ModuleId, Symbol};
 use roc_problem::can::CycleEntry;
 use roc_region::all::{Loc, Region};
 use roc_types::subs::{
-    self, AliasVariables, Content, Descriptor, FlatType, GetSubsSlice, LambdaSet, Mark,
+    self, get_member_lambda_sets_at_region, AliasVariables, Content, Descriptor, FlatType,
-    OptVariable, Rank, RecordFields, Subs, SubsIndex, SubsSlice, UlsOfVar, UnionLabels,
+    GetSubsSlice, LambdaSet, Mark, OptVariable, Rank, RecordFields, Subs, SubsIndex, SubsSlice,
-    UnionLambdas, UnionTags, Variable, VariableSubsSlice,
+    UlsOfVar, UnionLabels, UnionLambdas, UnionTags, Variable, VariableSubsSlice,
 };
 use roc_types::types::Type::{self, *};
 use roc_types::types::{
@@ -526,6 +526,8 @@ impl Pools {
 
 /// What phase in the compiler is reaching out to solve types.
 /// This is important to distinguish subtle differences in the behavior of the solving algorithm.
+//
+// TODO the APIs of this trait suck, this needs a nice cleanup.
 pub trait Phase {
     /// The regular type-solving phase, or during some later phase of compilation.
     /// During the solving phase we must anticipate that some information is still unknown and react to
@@ -536,12 +538,23 @@ pub trait Phase {
     where
         F: FnMut(&AbilitiesStore) -> T;
 
+    /// Given a known lambda set's ambient function in an external module, copy that ambient
+    /// function into the given subs.
     fn copy_lambda_set_ambient_function_to_home_subs(
         &self,
         external_lambda_set_var: Variable,
         external_module_id: ModuleId,
         home_subs: &mut Subs,
     ) -> Variable;
+
+    /// Find the ambient function var at a given region for an ability member definition (not a
+    /// specialization!), and copy that into the given subs.
+    fn get_and_copy_ability_member_ambient_function(
+        &self,
+        ability_member: Symbol,
+        region: u8,
+        home_subs: &mut Subs,
+    ) -> Variable;
 }
 
 struct SolvePhase<'a> {
@@ -572,6 +585,35 @@ impl Phase for SolvePhase<'_> {
         } = home_subs.get_lambda_set(external_lambda_set_var);
         ambient_function
     }
+
+    fn get_and_copy_ability_member_ambient_function(
+        &self,
+        ability_member: Symbol,
+        region: u8,
+        home_subs: &mut Subs,
+    ) -> Variable {
+        // During solving we're only aware of our module's abilities store, the var must
+        // be in our module store. Even if the specialization lambda set comes from another
+        // module, we should have taken care to import it before starting solving in this module.
+        let member_def = self
+            .abilities_store
+            .member_def(ability_member)
+            .unwrap_or_else(|| {
+                internal_error!(
+                    "{:?} is not resolved, or not an ability member!",
+                    ability_member
+                )
+            });
+        let member_var = member_def.signature_var();
+
+        let region_lset = get_member_lambda_sets_at_region(home_subs, member_var, region);
+
+        let LambdaSet {
+            ambient_function, ..
+        } = home_subs.get_lambda_set(region_lset);
+
+        ambient_function
+    }
 }
 
 #[derive(Clone)]
@@ -2179,6 +2221,7 @@ fn compact_lambda_set<P: Phase>(
 enum SpecializationTypeKey {
     Opaque(Symbol),
     Derived(DeriveKey),
+    Immediate(Symbol),
 }
 
 enum SpecializeDecision {
@@ -2195,8 +2238,9 @@ fn make_specialization_decision(subs: &Subs, var: Variable) -> SpecializeDecisio
             // should use.
             match roc_derive_key::Derived::encoding(subs, var) {
                 Ok(derived) => match derived {
-                    roc_derive_key::Derived::Immediate(_) => {
+                    roc_derive_key::Derived::Immediate(imm) => {
-                        todo!("deal with lambda set extraction from immediates")
+                        SpecializeDecision::Specialize(Immediate(imm))
+                        // todo!("deal with lambda set extraction from immediates")
                     }
                     roc_derive_key::Derived::Key(derive_key) => {
                         SpecializeDecision::Specialize(Derived(derive_key))
@@ -2295,6 +2339,18 @@ fn get_specialization_lambda_set_ambient_function<P: Phase>(
 
             Ok(specialized_ambient)
         }
+
+        SpecializationTypeKey::Immediate(imm) => {
+            // Immediates are like opaques in that we can simply look up their type definition in
+            // the ability store, there is nothing new to synthesize.
+            //
+            // THEORY: if something can become an immediate, it will always be avaialble in the
+            // local ability store, because the transformation is local (?)
+            let immediate_lambda_set_at_region =
+                phase.get_and_copy_ability_member_ambient_function(imm, lset_region, subs);
+
+            Ok(immediate_lambda_set_at_region)
+        }
     }
 }
 

crates/compiler/types/src/subs.rs

@@ -3964,8 +3964,10 @@ fn get_fresh_var_name(state: &mut ErrorTypeState) -> Lowercase {
 pub struct ExposedTypesStorageSubs {
     pub storage_subs: StorageSubs,
     pub stored_vars_by_symbol: VecMap<Symbol, Variable>,
-    /// lambda set var in other module -> var in storage subs
+    /// specialization lambda set var in other module -> var in storage subs
-    pub stored_ability_lambda_set_vars: VecMap<Variable, Variable>,
+    pub stored_specialization_lambda_set_vars: VecMap<Variable, Variable>,
+    /// ability member signature in other module -> var in storage subs
+    pub stored_ability_member_vars: VecMap<Variable, Variable>,
 }
 
 #[derive(Clone, Debug)]
@@ -5295,19 +5297,9 @@ fn instantiate_rigids_help(subs: &mut Subs, max_rank: Rank, initial: Variable) {
 ///
 /// Panics if the given function type does not correspond with what's expected of an ability
 /// member, namely its lambda sets have more than a single unspecialized lambda set.
-pub fn get_member_lambda_sets_at_region(
+pub fn get_member_lambda_sets_at_region(subs: &Subs, var: Variable, target_region: u8) -> Variable {
-    subs: &Subs,
-    var: Variable,
-    target_region: Option<u8>,
-) -> Vec<Variable> {
     let mut stack = vec![var];
 
-    let mut result = if target_region.is_some() {
-        Vec::with_capacity(1)
-    } else {
-        Vec::with_capacity(4)
-    };
-
     while let Some(var) = stack.pop() {
         match subs.get_content_without_compacting(var) {
             Content::LambdaSet(LambdaSet {
@@ -5320,19 +5312,8 @@ pub fn get_member_lambda_sets_at_region(
                 debug_assert!(recursion_var.is_none());
                 debug_assert_eq!(unspecialized.len(), 1);
                 let Uls(_, _, region) = subs.get_subs_slice(*unspecialized)[0];
-                match target_region {
+                if region == target_region {
-                    Some(r) => {
+                    return var;
-                        if region == r {
-                            result.push(var);
-                            return result;
-                        } else {
-                            continue;
-                        }
-                    }
-                    None => {
-                        result.push(var);
-                        continue;
-                    }
                 }
             }
             Content::Structure(flat_type) => match flat_type {
@@ -5386,5 +5367,5 @@ pub fn get_member_lambda_sets_at_region(
         }
     }
 
-    return result;
+    internal_error!("No lambda set at region {} found", target_region);
 }