Fix merge conflicts

crates/compiler/solve/src/specialize.rs

@@ -0,0 +1,786 @@
+//! Module [specialize] is resolves specialization lambda sets.
+
+use std::collections::VecDeque;
+
+use bumpalo::Bump;
+use roc_can::{
+    abilities::{AbilitiesStore, ImplKey},
+    module::ExposedByModule,
+};
+use roc_collections::{VecMap, VecSet};
+use roc_debug_flags::dbg_do;
+#[cfg(debug_assertions)]
+use roc_debug_flags::ROC_TRACE_COMPACTION;
+use roc_derive::SharedDerivedModule;
+use roc_derive_key::{DeriveBuiltin, DeriveError, DeriveKey};
+use roc_error_macros::{internal_error, todo_abilities};
+use roc_module::symbol::{ModuleId, Symbol};
+use roc_types::{
+    subs::{
+        get_member_lambda_sets_at_region, Content, Descriptor, GetSubsSlice, LambdaSet, Mark,
+        OptVariable, Rank, Subs, SubsSlice, UlsOfVar, Variable,
+    },
+    types::{AliasKind, MemberImpl, Uls},
+};
+use roc_unify::unify::{unify, Env as UEnv, Mode, MustImplementConstraints};
+
+use crate::solve::{deep_copy_var_in, introduce, Pools};
+
+/// What phase in the compiler is reaching out to specialize lambda sets?
+/// 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
+    /// that; during late phases, we expect that all information is resolved.
+    const IS_LATE: bool;
+
+    fn with_module_abilities_store<T, F>(&self, module: ModuleId, f: F) -> T
+    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;
+}
+
+pub(crate) struct SolvePhase<'a> {
+    pub abilities_store: &'a AbilitiesStore,
+}
+impl Phase for SolvePhase<'_> {
+    const IS_LATE: bool = false;
+
+    fn with_module_abilities_store<T, F>(&self, _module: ModuleId, mut f: F) -> T
+    where
+        F: FnMut(&AbilitiesStore) -> T,
+    {
+        // During solving we're only aware of our module's abilities store.
+        f(self.abilities_store)
+    }
+
+    fn copy_lambda_set_ambient_function_to_home_subs(
+        &self,
+        external_lambda_set_var: Variable,
+        _external_module_id: ModuleId,
+        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 LambdaSet {
+            ambient_function, ..
+        } = 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
+    }
+}
+
+pub struct DerivedEnv<'a> {
+    pub derived_module: &'a SharedDerivedModule,
+    /// Exposed types needed by the derived module.
+    pub exposed_types: &'a ExposedByModule,
+}
+
+#[derive(Default)]
+pub struct AwaitingSpecializations {
+    // What variables' specialized lambda sets in `uls_of_var` will be unlocked for specialization
+    // when an implementation key's specialization is resolved?
+    waiting: VecMap<ImplKey, VecSet<Variable>>,
+    uls_of_var: UlsOfVar,
+}
+
+impl AwaitingSpecializations {
+    pub fn remove_for_specialized(&mut self, subs: &Subs, impl_key: ImplKey) -> UlsOfVar {
+        let spec_variables = self
+            .waiting
+            .remove(&impl_key)
+            .map(|(_, set)| set)
+            .unwrap_or_default();
+
+        let mut result = UlsOfVar::default();
+        for var in spec_variables {
+            let target_lambda_sets = self
+                .uls_of_var
+                .remove_dependent_unspecialized_lambda_sets(subs, var);
+
+            result.extend(var, target_lambda_sets);
+        }
+        result
+    }
+
+    pub fn add(
+        &mut self,
+        impl_key: ImplKey,
+        var: Variable,
+        lambda_sets: impl IntoIterator<Item = Variable>,
+    ) {
+        self.uls_of_var.extend(var, lambda_sets);
+        let waiting = self.waiting.get_or_insert(impl_key, Default::default);
+        waiting.insert(var);
+    }
+
+    pub fn union(&mut self, other: Self) {
+        for (impl_key, waiting_vars) in other.waiting {
+            let waiting = self.waiting.get_or_insert(impl_key, Default::default);
+            waiting.extend(waiting_vars);
+        }
+        self.uls_of_var.union(other.uls_of_var);
+    }
+
+    pub fn waiting_for(&self, impl_key: ImplKey) -> bool {
+        self.waiting.contains_key(&impl_key)
+    }
+}
+
+pub struct CompactionResult {
+    pub obligations: MustImplementConstraints,
+    pub awaiting_specialization: AwaitingSpecializations,
+}
+
+#[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);
+    let uls_a = uls_a
+        .iter()
+        .map(|v| {
+            format!(
+                "{:?}",
+                roc_types::subs::SubsFmtContent(subs.get_content_without_compacting(*v), subs)
+            )
+        })
+        .collect::<Vec<_>>()
+        .join(",");
+    eprintln!("===lambda set compaction===");
+    eprintln!("  concrete type: {:?}", c_a);
+    eprintln!("  step 1:");
+    eprintln!("    uls_a = {{ {} }}", uls_a);
+}
+
+#[cfg(debug_assertions)]
+fn trace_compaction_step_2(subs: &Subs, uls_a: &[Variable]) {
+    let uls_a = uls_a
+        .iter()
+        .map(|v| {
+            format!(
+                "{:?}",
+                roc_types::subs::SubsFmtContent(subs.get_content_without_compacting(*v), subs)
+            )
+        })
+        .collect::<Vec<_>>()
+        .join(",");
+    eprintln!("  step 2:");
+    eprintln!("    uls_a' = {{ {} }}", uls_a);
+}
+
+#[cfg(debug_assertions)]
+fn trace_compaction_step_3start() {
+    eprintln!("  step 3:");
+}
+
+#[cfg(debug_assertions)]
+fn trace_compaction_step_3iter_start(
+    subs: &Subs,
+    iteration_lambda_set: Variable,
+    t_f1: Variable,
+    t_f2: Variable,
+) {
+    let iteration_lambda_set = roc_types::subs::SubsFmtContent(
+        subs.get_content_without_compacting(iteration_lambda_set),
+        subs,
+    );
+    let t_f1 = roc_types::subs::SubsFmtContent(subs.get_content_without_compacting(t_f1), subs);
+    let t_f2 = roc_types::subs::SubsFmtContent(subs.get_content_without_compacting(t_f2), subs);
+    eprintln!("    - iteration: {:?}", iteration_lambda_set);
+    eprintln!("         {:?}", t_f1);
+    eprintln!("      ~  {:?}", t_f2);
+}
+
+#[cfg(debug_assertions)]
+#[rustfmt::skip]
+fn trace_compaction_step_3iter_end(subs: &Subs, t_f_result: Variable, skipped: bool) {
+    let t_f_result =
+        roc_types::subs::SubsFmtContent(subs.get_content_without_compacting(t_f_result), subs);
+    if skipped {
+    eprintln!("      SKIP");
+    }
+    eprintln!("      =  {:?}\n", t_f_result);
+}
+
+macro_rules! trace_compact {
+    (1. $subs:expr, $c_a:expr, $uls_a:expr) => {{
+        dbg_do!(ROC_TRACE_COMPACTION, {
+            trace_compaction_step_1($subs, $c_a, $uls_a)
+        })
+    }};
+    (2. $subs:expr, $uls_a:expr) => {{
+        dbg_do!(ROC_TRACE_COMPACTION, {
+            trace_compaction_step_2($subs, $uls_a)
+        })
+    }};
+    (3start.) => {{
+        dbg_do!(ROC_TRACE_COMPACTION, { trace_compaction_step_3start() })
+    }};
+    (3iter_start. $subs:expr, $iteration_lset:expr, $t_f1:expr, $t_f2:expr) => {{
+        dbg_do!(ROC_TRACE_COMPACTION, {
+            trace_compaction_step_3iter_start($subs, $iteration_lset, $t_f1, $t_f2)
+        })
+    }};
+    (3iter_end. $subs:expr, $t_f_result:expr) => {{
+        dbg_do!(ROC_TRACE_COMPACTION, {
+            trace_compaction_step_3iter_end($subs, $t_f_result, false)
+        })
+    }};
+    (3iter_end_skipped. $subs:expr, $t_f_result:expr) => {{
+        dbg_do!(ROC_TRACE_COMPACTION, {
+            trace_compaction_step_3iter_end($subs, $t_f_result, true)
+        })
+    }};
+}
+
+#[inline(always)]
+fn iter_concrete_of_unspecialized<'a>(
+    subs: &'a Subs,
+    c_a: Variable,
+    uls: &'a [Uls],
+) -> impl Iterator<Item = &'a Uls> {
+    uls.iter()
+        .filter(move |Uls(var, _, _)| subs.equivalent_without_compacting(*var, c_a))
+}
+
+/// Gets the unique unspecialized lambda resolving to concrete type `c_a` in a list of
+/// unspecialized lambda sets.
+#[inline(always)]
+fn unique_unspecialized_lambda(subs: &Subs, c_a: Variable, uls: &[Uls]) -> Option<Uls> {
+    let mut iter_concrete = iter_concrete_of_unspecialized(subs, c_a, uls);
+    let uls = iter_concrete.next()?;
+    debug_assert!(iter_concrete.next().is_none(), "multiple concrete");
+    Some(*uls)
+}
+
+#[must_use]
+pub fn compact_lambda_sets_of_vars<P: Phase>(
+    subs: &mut Subs,
+    derived_env: &DerivedEnv,
+    arena: &Bump,
+    pools: &mut Pools,
+    uls_of_var: UlsOfVar,
+    phase: &P,
+) -> CompactionResult {
+    let mut must_implement = MustImplementConstraints::default();
+    let mut awaiting_specialization = AwaitingSpecializations::default();
+
+    let mut uls_of_var_queue = VecDeque::with_capacity(uls_of_var.len());
+    uls_of_var_queue.extend(uls_of_var.drain());
+
+    // Suppose a type variable `a` with `uls_of_var` mapping `uls_a = {l1, ... ln}` has been instantiated to a concrete type `C_a`.
+    while let Some((c_a, uls_a)) = uls_of_var_queue.pop_front() {
+        let c_a = subs.get_root_key_without_compacting(c_a);
+        // 1. Let each `l` in `uls_a` be of form `[solved_lambdas + ... + C:f:r + ...]`.
+        //    NB: There may be multiple unspecialized lambdas of form `C:f:r, C:f1:r1, ..., C:fn:rn` in `l`.
+        //    In this case, let `t1, ... tm` be the other unspecialized lambdas not of form `C:_:_`,
+        //    that is, none of which are now specialized to the type `C`. Then, deconstruct
+        //    `l` such that `l' = [solved_lambdas + t1 + ... + tm + C:f:r]` and `l1 = [[] + C:f1:r1], ..., ln = [[] + C:fn:rn]`.
+        //    Replace `l` with `l', l1, ..., ln` in `uls_a`, flattened.
+        //    TODO: the flattening step described above
+        let uls_a = {
+            let mut uls = uls_a.into_vec();
+
+            // De-duplicate lambdas by root key.
+            uls.iter_mut().for_each(|v| *v = subs.get_root_key(*v));
+            uls.sort();
+            uls.dedup();
+            uls
+        };
+
+        trace_compact!(1. subs, c_a, &uls_a);
+
+        // The flattening step - remove lambda sets that don't reference the concrete var, and for
+        // flatten lambda sets that reference it more than once.
+        let mut uls_a: Vec<_> = uls_a
+            .into_iter()
+            .flat_map(|lambda_set| {
+                let LambdaSet {
+                    solved,
+                    recursion_var,
+                    unspecialized,
+                    ambient_function,
+                } = subs.get_lambda_set(lambda_set);
+                let lambda_set_rank = subs.get_rank(lambda_set);
+                let unspecialized = subs.get_subs_slice(unspecialized);
+                // TODO: is it faster to traverse once, see if we only have one concrete lambda, and
+                // bail in that happy-path, rather than always splitting?
+                let (concrete, mut not_concrete): (Vec<_>, Vec<_>) = unspecialized
+                    .iter()
+                    .copied()
+                    .partition(|Uls(var, _, _)| subs.equivalent_without_compacting(*var, c_a));
+                if concrete.len() == 1 {
+                    // No flattening needs to be done, just return the lambda set as-is
+                    return vec![lambda_set];
+                }
+                // Must flatten
+                concrete
+                    .into_iter()
+                    .enumerate()
+                    .map(|(i, concrete_lambda)| {
+                        let (var, unspecialized) = if i == 0 {
+                            // The first lambda set contains one concrete lambda, plus all solved
+                            // lambdas, plus all other unspecialized lambdas.
+                            // l' = [solved_lambdas + t1 + ... + tm + C:f:r]
+                            let unspecialized = SubsSlice::extend_new(
+                                &mut subs.unspecialized_lambda_sets,
+                                not_concrete
+                                    .drain(..)
+                                    .chain(std::iter::once(concrete_lambda)),
+                            );
+                            (lambda_set, unspecialized)
+                        } else {
+                            // All the other lambda sets consists only of their respective concrete
+                            // lambdas.
+                            // ln = [[] + C:fn:rn]
+                            let unspecialized = SubsSlice::extend_new(
+                                &mut subs.unspecialized_lambda_sets,
+                                [concrete_lambda],
+                            );
+                            let var = subs.fresh(Descriptor {
+                                content: Content::Error,
+                                rank: lambda_set_rank,
+                                mark: Mark::NONE,
+                                copy: OptVariable::NONE,
+                            });
+                            (var, unspecialized)
+                        };
+
+                        subs.set_content(
+                            var,
+                            Content::LambdaSet(LambdaSet {
+                                solved,
+                                recursion_var,
+                                unspecialized,
+                                ambient_function,
+                            }),
+                        );
+                        var
+                    })
+                    .collect()
+            })
+            .collect();
+
+        // 2. Now, each `l` in `uls_a` has a unique unspecialized lambda of form `C:f:r`.
+        //    Sort `uls_a` primarily by `f` (arbitrary order), and secondarily by `r` in descending order.
+        uls_a.sort_by(|v1, v2| {
+            let unspec_1 = subs.get_subs_slice(subs.get_lambda_set(*v1).unspecialized);
+            let unspec_2 = subs.get_subs_slice(subs.get_lambda_set(*v2).unspecialized);
+
+            let Uls(_, f1, r1) = unique_unspecialized_lambda(subs, c_a, unspec_1).unwrap();
+            let Uls(_, f2, r2) = unique_unspecialized_lambda(subs, c_a, unspec_2).unwrap();
+
+            match f1.cmp(&f2) {
+                std::cmp::Ordering::Equal => {
+                    // Order by descending order of region.
+                    r2.cmp(&r1)
+                }
+                ord => ord,
+            }
+        });
+
+        trace_compact!(2. subs, &uls_a);
+
+        // 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.
+        //       - For example, `(b' -[[] + Fo:f:2]-> {})` if `C:f:r=Fo:f:2`. Removing `Fo:f:2`, we get `(b' -[[]]-> {})`.
+        //    2. Let `t_f2` be the directly ambient function of the specialization lambda set resolved by `C:f:r`.
+        //       - For example, `(b -[[] + b:g:1]-> {})` if `C:f:r=Fo:f:2`, running on example from above.
+        //    3. Unify `t_f1 ~ t_f2`.
+        trace_compact!(3start.);
+        for l in uls_a {
+            let compaction_result =
+                compact_lambda_set(subs, derived_env, arena, pools, c_a, l, phase);
+
+            match compaction_result {
+                OneCompactionResult::Compacted {
+                    new_obligations,
+                    new_lambda_sets_to_specialize,
+                } => {
+                    must_implement.extend(new_obligations);
+                    uls_of_var_queue.extend(new_lambda_sets_to_specialize.drain());
+                }
+                OneCompactionResult::MustWaitForSpecialization(impl_key) => {
+                    awaiting_specialization.add(impl_key, c_a, [l])
+                }
+            }
+        }
+    }
+
+    CompactionResult {
+        obligations: must_implement,
+        awaiting_specialization,
+    }
+}
+
+enum OneCompactionResult {
+    Compacted {
+        new_obligations: MustImplementConstraints,
+        new_lambda_sets_to_specialize: UlsOfVar,
+    },
+    MustWaitForSpecialization(ImplKey),
+}
+
+#[must_use]
+#[allow(clippy::too_many_arguments)]
+fn compact_lambda_set<P: Phase>(
+    subs: &mut Subs,
+    derived_env: &DerivedEnv,
+    arena: &Bump,
+    pools: &mut Pools,
+    resolved_concrete: Variable,
+    this_lambda_set: Variable,
+    phase: &P,
+) -> OneCompactionResult {
+    // 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.
+    //       - For example, `(b' -[[] + Fo:f:2]-> {})` if `C:f:r=Fo:f:2`. Removing `Fo:f:2`, we get `(b' -[[]]-> {})`.
+    //    2. Let `t_f2` be the directly ambient function of the specialization lambda set resolved by `C:f:r`.
+    //       - For example, `(b -[[] + b:g:1]-> {})` if `C:f:r=Fo:f:2`, from the algorithm's running example.
+    //    3. Unify `t_f1 ~ t_f2`.
+    let LambdaSet {
+        solved,
+        recursion_var,
+        unspecialized,
+        ambient_function: t_f1,
+    } = subs.get_lambda_set(this_lambda_set);
+    let target_rank = subs.get_rank(this_lambda_set);
+
+    debug_assert!(!unspecialized.is_empty());
+
+    let unspecialized = subs.get_subs_slice(unspecialized);
+
+    // 1. Let `t_f1` be the directly ambient function of the lambda set containing `C:f:r`.
+    let Uls(c, f, r) = unique_unspecialized_lambda(subs, resolved_concrete, unspecialized).unwrap();
+
+    debug_assert!(subs.equivalent_without_compacting(c, resolved_concrete));
+
+    // Now decide: do we
+    // - proceed with specialization
+    // - simply drop the specialization lambda set (due to an error)
+    // - or do we need to wait, because we don't know enough information for the specialization yet?
+    let specialization_decision = make_specialization_decision(subs, phase, c, f);
+    let specialization_key_or_drop = match specialization_decision {
+        SpecializeDecision::Specialize(key) => Ok(key),
+        SpecializeDecision::Drop => Err(()),
+        SpecializeDecision::PendingSpecialization(impl_key) => {
+            // Bail, we need to wait for the specialization to be known.
+            return OneCompactionResult::MustWaitForSpecialization(impl_key);
+        }
+    };
+
+    // 1b. Remove `C:f:r` from `t_f1`'s lambda set.
+    let new_unspecialized: Vec<_> = unspecialized
+        .iter()
+        .filter(|Uls(v, _, _)| !subs.equivalent_without_compacting(*v, resolved_concrete))
+        .copied()
+        .collect();
+    debug_assert_eq!(new_unspecialized.len(), unspecialized.len() - 1);
+    let t_f1_lambda_set_without_concrete = LambdaSet {
+        solved,
+        recursion_var,
+        unspecialized: SubsSlice::extend_new(
+            &mut subs.unspecialized_lambda_sets,
+            new_unspecialized,
+        ),
+        ambient_function: t_f1,
+    };
+    subs.set_content(
+        this_lambda_set,
+        Content::LambdaSet(t_f1_lambda_set_without_concrete),
+    );
+
+    let specialization_key = match specialization_key_or_drop {
+        Ok(specialization_key) => specialization_key,
+        Err(()) => {
+            // Do nothing other than to remove the concrete lambda to drop from the lambda set,
+            // which we already did in 1b above.
+            trace_compact!(3iter_end_skipped. subs, t_f1);
+            return OneCompactionResult::Compacted {
+                new_obligations: Default::default(),
+                new_lambda_sets_to_specialize: Default::default(),
+            };
+        }
+    };
+
+    let specialization_ambient_function_var = get_specialization_lambda_set_ambient_function(
+        subs,
+        derived_env,
+        phase,
+        f,
+        r,
+        specialization_key,
+        target_rank,
+    );
+
+    let t_f2 = match specialization_ambient_function_var {
+        Ok(lset) => lset,
+        Err(()) => {
+            // Do nothing other than to remove the concrete lambda to drop from the lambda set,
+            // which we already did in 1b above.
+            trace_compact!(3iter_end_skipped. subs, t_f1);
+            return OneCompactionResult::Compacted {
+                new_obligations: Default::default(),
+                new_lambda_sets_to_specialize: Default::default(),
+            };
+        }
+    };
+
+    // Ensure the specialized ambient function we'll unify with is not a generalized one, but one
+    // at the rank of the lambda set being compacted.
+    let t_f2 = deep_copy_var_in(subs, target_rank, pools, t_f2, arena);
+
+    // 3. Unify `t_f1 ~ t_f2`.
+    trace_compact!(3iter_start. subs, this_lambda_set, t_f1, t_f2);
+    let (vars, new_obligations, new_lambda_sets_to_specialize, _meta) = unify(
+        &mut UEnv::new(subs),
+        t_f1,
+        t_f2,
+        Mode::LAMBDA_SET_SPECIALIZATION,
+    )
+    .expect_success("ambient functions don't unify");
+    trace_compact!(3iter_end. subs, t_f1);
+
+    introduce(subs, target_rank, pools, &vars);
+
+    OneCompactionResult::Compacted {
+        new_obligations,
+        new_lambda_sets_to_specialize,
+    }
+}
+
+#[derive(Debug)]
+enum SpecializationTypeKey {
+    Opaque(Symbol),
+    Derived(DeriveKey),
+    Immediate(Symbol),
+}
+
+enum SpecializeDecision {
+    Specialize(SpecializationTypeKey),
+    Drop,
+
+    /// Only relevant during module solving of recursive defs - we don't yet know the
+    /// specialization type for a declared ability implementation, so we must hold off on
+    /// specialization.
+    PendingSpecialization(ImplKey),
+}
+
+fn make_specialization_decision<P: Phase>(
+    subs: &Subs,
+    phase: &P,
+    var: Variable,
+    ability_member: Symbol,
+) -> SpecializeDecision {
+    use Content::*;
+    use SpecializationTypeKey::*;
+    match subs.get_content_without_compacting(var) {
+        Alias(opaque, _, _, AliasKind::Opaque) if opaque.module_id() != ModuleId::NUM => {
+            if P::IS_LATE {
+                SpecializeDecision::Specialize(Opaque(*opaque))
+            } else {
+                // Solving within a module.
+                phase.with_module_abilities_store(opaque.module_id(), |abilities_store| {
+                    let impl_key = ImplKey {
+                        opaque: *opaque,
+                        ability_member,
+                    };
+                    match abilities_store.get_implementation(impl_key) {
+                        None => {
+                            // Doesn't specialize; an error will already be reported for this.
+                            SpecializeDecision::Drop
+                        }
+                        Some(MemberImpl::Error | MemberImpl::Derived) => {
+                            // TODO: probably not right, we may want to choose a derive decision!
+                            SpecializeDecision::Specialize(Opaque(*opaque))
+                        }
+                        Some(MemberImpl::Impl(specialization_symbol)) => {
+                            match abilities_store.specialization_info(*specialization_symbol) {
+                                Some(_) => SpecializeDecision::Specialize(Opaque(*opaque)),
+
+                                // If we expect a specialization impl but don't yet know it, we must hold off
+                                // compacting the lambda set until the specialization is well-known.
+                                None => SpecializeDecision::PendingSpecialization(impl_key),
+                            }
+                        }
+                    }
+                })
+            }
+        }
+        Structure(_) | Alias(_, _, _, _) => {
+            // This is a structural type, find the name of the derived ability function it
+            // should use.
+            match roc_derive_key::Derived::builtin(DeriveBuiltin::ToEncoder, subs, var) {
+                Ok(derived) => match derived {
+                    roc_derive_key::Derived::Immediate(imm) => {
+                        SpecializeDecision::Specialize(Immediate(imm))
+                        // todo!("deal with lambda set extraction from immediates")
+                    }
+                    roc_derive_key::Derived::Key(derive_key) => {
+                        SpecializeDecision::Specialize(Derived(derive_key))
+                    }
+                },
+                Err(DeriveError::UnboundVar) => {
+                    // not specialized yet, but that also means that it can't possibly be derivable
+                    // at this point?
+                    // TODO: is this right? Revisit if it causes us problems in the future.
+                    SpecializeDecision::Drop
+                }
+                Err(DeriveError::Underivable) => {
+                    // we should have reported an error for this; drop the lambda set.
+                    SpecializeDecision::Drop
+                }
+            }
+        }
+        Error => SpecializeDecision::Drop,
+        FlexAbleVar(_, _)
+        | RigidAbleVar(..)
+        | FlexVar(..)
+        | RigidVar(..)
+        | RecursionVar { .. }
+        | LambdaSet(..)
+        | RangedNumber(..) => {
+            internal_error!("unexpected")
+        }
+    }
+}
+
+#[allow(clippy::too_many_arguments)]
+fn get_specialization_lambda_set_ambient_function<P: Phase>(
+    subs: &mut Subs,
+    derived_env: &DerivedEnv,
+    phase: &P,
+    ability_member: Symbol,
+    lset_region: u8,
+    specialization_key: SpecializationTypeKey,
+    target_rank: Rank,
+) -> Result<Variable, ()> {
+    match specialization_key {
+        SpecializationTypeKey::Opaque(opaque) => {
+            let opaque_home = opaque.module_id();
+            let external_specialized_lset =
+                phase.with_module_abilities_store(opaque_home, |abilities_store| {
+            let impl_key = roc_can::abilities::ImplKey {
+                opaque,
+                ability_member,
+            };
+                    let opt_specialization =
+                        abilities_store.get_implementation(impl_key);
+                    match opt_specialization {
+                        None => {
+                            if P::IS_LATE {
+                                internal_error!(
+                                    "expected to know a specialization for {:?}#{:?}, but it wasn't found",
+                                    opaque,
+                                    ability_member
+                                );
+                            } else {
+                                // doesn't specialize, we'll have reported an error for this
+                                Err(())
+                            }
+                        }
+                        Some(member_impl) => match member_impl {
+                            MemberImpl::Impl(spec_symbol) => {
+                                let specialization =
+                                    abilities_store.specialization_info(*spec_symbol).expect("expected custom implementations to always have complete specialization info by this point");
+
+                                let specialized_lambda_set = *specialization
+                                    .specialization_lambda_sets
+                                    .get(&lset_region)
+                                    .expect("lambda set region not resolved");
+                                Ok(specialized_lambda_set)
+                            }
+                            MemberImpl::Derived => todo_abilities!(),
+                            MemberImpl::Error => todo_abilities!(),
+                        },
+                    }
+                })?;
+
+            let specialized_ambient = phase.copy_lambda_set_ambient_function_to_home_subs(
+                external_specialized_lset,
+                opaque_home,
+                subs,
+            );
+
+            Ok(specialized_ambient)
+        }
+
+        SpecializationTypeKey::Derived(derive_key) => {
+            let mut derived_module = derived_env.derived_module.lock().unwrap();
+
+            let (_, _, specialization_lambda_sets) =
+                derived_module.get_or_insert(derived_env.exposed_types, derive_key);
+
+            let specialized_lambda_set = *specialization_lambda_sets
+                .get(&lset_region)
+                .expect("lambda set region not resolved");
+
+            let specialized_ambient = derived_module.copy_lambda_set_ambient_function_to_subs(
+                specialized_lambda_set,
+                subs,
+                target_rank,
+            );
+
+            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 available 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/solve/tests/solve_expr.rs

@@ -7370,10 +7370,10 @@ mod solve_expr {
         E#k(10) : E -[[k(10)]]-> {}
         a : j | j has J
         b : j | j has J
-        it : k -[[] + j:j(2):2 + a:j(2):2]-> {} | a has J, j has J, k has K
-        J#j(2) : j -[[] + j:j(2):1]-> (k -[[] + j:j(2):2 + a:j(2):2]-> {}) | a has J, j has J, k has K
-        J#j(2) : j -[[] + j:j(2):1]-> (k -[[] + a:j(2):2 + j:j(2):2]-> {}) | a has J, j has J, k has K
-        it : k -[[] + j:j(2):2 + a:j(2):2]-> {} | a has J, j has J, k has K
+        it : k -[[] + j:j(2):2 + j1:j(2):2]-> {} | j has J, j1 has J, k has K
+        J#j(2) : j -[[] + j:j(2):1]-> (k -[[] + j:j(2):2 + j1:j(2):2]-> {}) | j has J, j1 has J, k has K
+        J#j(2) : j -[[] + j:j(2):1]-> (k -[[] + j1:j(2):2 + j:j(2):2]-> {}) | j has J, j1 has J, k has K
+        it : k -[[] + j:j(2):2 + j1:j(2):2]-> {} | j has J, j1 has J, k has K
         f : [A, B], C, D -[[f(11)]]-> (E -[[k(10)]]-> {})
         f A (@C {}) (@D {}) : E -[[k(10)]]-> {}
         main : {}
@@ -7439,10 +7439,10 @@ mod solve_expr {
         kF : F -[[kF(12)]]-> {}
         a : j | j has J
         b : j | j has J
-        it : k -[[] + j:j(2):2 + a:j(2):2]-> {} | a has J, j has J, k has K
-        J#j(2) : j -[[] + j:j(2):1]-> (k -[[] + j:j(2):2 + a:j(2):2]-> {}) | a has J, j has J, k has K
-        J#j(2) : j -[[] + j:j(2):1]-> (k -[[] + a:j(2):2 + j:j(2):2]-> {}) | a has J, j has J, k has K
-        it : k -[[] + j:j(2):2 + a:j(2):2]-> {} | a has J, j has J, k has K
+        it : k -[[] + j:j(2):2 + j1:j(2):2]-> {} | j has J, j1 has J, k has K
+        J#j(2) : j -[[] + j:j(2):1]-> (k -[[] + j:j(2):2 + j1:j(2):2]-> {}) | j has J, j1 has J, k has K
+        J#j(2) : j -[[] + j:j(2):1]-> (k -[[] + j1:j(2):2 + j:j(2):2]-> {}) | j has J, j1 has J, k has K
+        it : k -[[] + j:j(2):2 + j1:j(2):2]-> {} | j has J, j1 has J, k has K
         main : {}
         it : k -[[] + k:k(4):1]-> {} | k has K
         f : [A, B], C, D -[[f(13)]]-> (k -[[] + k:k(4):1]-> {}) | k has K