Serialize abilities store and solved implementations with subs

crates/compiler/can/src/abilities.rs

@@ -9,6 +9,14 @@ use roc_types::{
     types::{MemberImpl, Type},
 };
 
+/// During type solving and monomorphization, a module must know how its imported ability
+/// implementations are resolved - are they derived, or have a concrete implementation?
+///
+/// Unfortunately we cannot keep this information opaque, as it's important for properly
+/// restoring specialization lambda sets. As such, we need to export implementation information,
+/// which is the job of this structure.
+pub type ResolvedImplementations = VecMap<ImplKey, ResolvedImpl>;
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct MemberVariables {
     pub able_vars: Vec<Variable>,
@@ -475,11 +483,14 @@ impl IAbilitiesStore<Resolved> {
         serialize::serialize(self, writer)
     }
 
-    pub fn deserialize(bytes: &[u8]) -> Self {
+    pub fn deserialize(bytes: &[u8]) -> (Self, usize) {
         serialize::deserialize(bytes)
     }
 }
 
+pub use serialize::deserialize_solved_implementations;
+pub use serialize::serialize_solved_implementations;
+
 impl IAbilitiesStore<Pending> {
     pub fn import_implementation(&mut self, impl_key: ImplKey, resolved_impl: &ResolvedImpl) {
         let member_impl = match resolved_impl {
@@ -682,7 +693,7 @@ mod serialize {
 
     use super::{
         AbilitiesStore, AbilityMemberData, ImplKey, MemberSpecializationInfo, Resolved,
-        ResolvedMemberType, SpecializationId,
+        ResolvedImpl, ResolvedImplementations, ResolvedMemberType, SpecializationId,
     };
 
     use std::io::{self, Write};
@@ -760,7 +771,7 @@ mod serialize {
         Ok(written)
     }
 
-    pub(super) fn deserialize(bytes: &[u8]) -> AbilitiesStore {
+    pub(super) fn deserialize(bytes: &[u8]) -> (AbilitiesStore, usize) {
         let mut offset = 0;
 
         let header_slice = &bytes[..std::mem::size_of::<Header>()];
@@ -786,17 +797,18 @@ mod serialize {
             offset,
         );
 
-        let _ = offset;
-
-        AbilitiesStore {
-            members_of_ability,
-            specialization_to_root,
-            ability_members,
-            declared_implementations,
-            specializations,
-            next_specialization_id: (header.next_specialization_id as u32).try_into().unwrap(),
-            resolved_specializations,
-        }
+        (
+            AbilitiesStore {
+                members_of_ability,
+                specialization_to_root,
+                ability_members,
+                declared_implementations,
+                specializations,
+                next_specialization_id: (header.next_specialization_id as u32).try_into().unwrap(),
+                resolved_specializations,
+            },
+            offset,
+        )
     }
 
     fn serialize_members_of_ability(
@@ -1113,6 +1125,127 @@ mod serialize {
             offset,
         )
     }
+
+    #[derive(Copy, Clone)]
+    #[repr(C)]
+    enum SerResolvedImpl {
+        Impl(SerMemberSpecInfo),
+        Derived,
+        Error,
+    }
+    impl SerResolvedImpl {
+        fn num_regions(&self) -> usize {
+            match self {
+                SerResolvedImpl::Impl(spec) => spec.1.len(),
+                SerResolvedImpl::Derived => 0,
+                SerResolvedImpl::Error => 0,
+            }
+        }
+    }
+
+    pub fn serialize_solved_implementations(
+        solved_impls: &ResolvedImplementations,
+        writer: &mut impl std::io::Write,
+    ) -> std::io::Result<usize> {
+        let len = solved_impls.len() as u64;
+
+        let written = bytes::serialize_slice(&[len], writer, 0)?;
+
+        bytes::serialize_vec_map(
+            solved_impls,
+            |keys, writer, written| {
+                bytes::serialize_slice(
+                    &keys.iter().map(SerImplKey::from).collect::<Vec<_>>(),
+                    writer,
+                    written,
+                )
+            },
+            |resolved_impls, writer, written| {
+                let mut spec_lambda_sets_regions: Vec<u8> = Vec::new();
+                let mut spec_lambda_sets_vars: Vec<Variable> = Vec::new();
+                let mut ser_resolved_impls: Vec<SerResolvedImpl> = Vec::new();
+
+                for resolved_impl in resolved_impls {
+                    let ser_resolved_impl = match resolved_impl {
+                        ResolvedImpl::Impl(MemberSpecializationInfo {
+                            _phase: _,
+                            symbol,
+                            specialization_lambda_sets,
+                        }) => {
+                            let regions = SubsSlice::extend_new(
+                                &mut spec_lambda_sets_regions,
+                                specialization_lambda_sets.keys().copied(),
+                            );
+                            let vars = SubsSlice::extend_new(
+                                &mut spec_lambda_sets_vars,
+                                specialization_lambda_sets.values().copied(),
+                            );
+                            SerResolvedImpl::Impl(SerMemberSpecInfo(*symbol, regions, vars))
+                        }
+                        ResolvedImpl::Derived => SerResolvedImpl::Derived,
+                        ResolvedImpl::Error => SerResolvedImpl::Error,
+                    };
+
+                    ser_resolved_impls.push(ser_resolved_impl);
+                }
+
+                let written = bytes::serialize_slice(&ser_resolved_impls, writer, written)?;
+                let written = bytes::serialize_slice(&spec_lambda_sets_regions, writer, written)?;
+                let written = bytes::serialize_slice(&spec_lambda_sets_vars, writer, written)?;
+
+                Ok(written)
+            },
+            writer,
+            written,
+        )
+    }
+
+    pub fn deserialize_solved_implementations(bytes: &[u8]) -> (ResolvedImplementations, usize) {
+        let (len_slice, offset) = bytes::deserialize_slice::<u64>(bytes, 1, 0);
+        let length = len_slice[0];
+
+        bytes::deserialize_vec_map(
+            bytes,
+            |bytes, length, offset| {
+                let (slice, offset) = bytes::deserialize_slice::<SerImplKey>(bytes, length, offset);
+                (slice.iter().map(ImplKey::from).collect(), offset)
+            },
+            |bytes, length, offset| {
+                let (serialized_slices, offset) =
+                    bytes::deserialize_slice::<SerResolvedImpl>(bytes, length, offset);
+
+                let total_num_regions = serialized_slices.iter().map(|s| s.num_regions()).sum();
+
+                let (regions_slice, offset) =
+                    bytes::deserialize_slice::<u8>(bytes, total_num_regions, offset);
+
+                let (vars_slice, offset) =
+                    { bytes::deserialize_slice::<Variable>(bytes, total_num_regions, offset) };
+
+                let mut all_resolved: Vec<ResolvedImpl> = Vec::with_capacity(length);
+                for ser_resolved in serialized_slices {
+                    let resolved = match ser_resolved {
+                        SerResolvedImpl::Impl(SerMemberSpecInfo(symbol, regions, vars)) => {
+                            let regions = regions_slice[regions.indices()].to_vec();
+                            let lset_vars = vars_slice[vars.indices()].to_vec();
+                            let spec_info = MemberSpecializationInfo::new(*symbol, unsafe {
+                                VecMap::zip(regions, lset_vars)
+                            });
+                            ResolvedImpl::Impl(spec_info)
+                        }
+                        SerResolvedImpl::Derived => ResolvedImpl::Derived,
+                        SerResolvedImpl::Error => ResolvedImpl::Error,
+                    };
+
+                    all_resolved.push(resolved);
+                }
+
+                (all_resolved, offset)
+            },
+            length as _,
+            offset,
+        )
+    }
 }
 
 #[cfg(test)]
@@ -1231,7 +1364,9 @@ mod test {
             resolved_specializations,
         } = store;
 
-        let de_store = AbilitiesStore::deserialize(&bytes);
+        let (de_store, offset) = AbilitiesStore::deserialize(&bytes);
+
+        assert_eq!(bytes.len(), offset);
 
         assert_eq!(members_of_ability, de_store.members_of_ability);
         assert_eq!(specialization_to_root, de_store.specialization_to_root);

crates/compiler/can/src/module.rs

@@ -1,4 +1,4 @@
-use crate::abilities::{ImplKey, PendingAbilitiesStore, ResolvedImpl};
+use crate::abilities::{AbilitiesStore, ImplKey, PendingAbilitiesStore, ResolvedImpl};
 use crate::annotation::canonicalize_annotation;
 use crate::def::{canonicalize_defs, Def};
 use crate::effect_module::HostedGeneratedFunctions;
@@ -17,7 +17,7 @@ use roc_parse::header::HeaderFor;
 use roc_parse::pattern::PatternType;
 use roc_problem::can::{Problem, RuntimeError};
 use roc_region::all::{Loc, Region};
-use roc_types::subs::{ExposedTypesStorageSubs, VarStore, Variable};
+use roc_types::subs::{ExposedTypesStorageSubs, Subs, VarStore, Variable};
 use roc_types::types::{Alias, AliasKind, AliasVar, Type};
 
 /// The types of all exposed values/functions of a collection of modules
@@ -1179,3 +1179,53 @@ fn fix_values_captured_in_closure_expr(
         OpaqueWrapFunction(_) => {}
     }
 }
+
+/// Type state for a single module.
+#[derive(Debug)]
+pub struct TypeState {
+    pub subs: Subs,
+    pub exposed_vars_by_symbol: Vec<(Symbol, Variable)>,
+    pub abilities: AbilitiesStore,
+    pub solved_implementations: ResolvedImplementations,
+}
+
+impl TypeState {
+    pub fn serialize(&self, writer: &mut impl std::io::Write) -> std::io::Result<usize> {
+        let Self {
+            subs,
+            exposed_vars_by_symbol,
+            abilities,
+            solved_implementations,
+        } = self;
+
+        let written_subs = subs.serialize(exposed_vars_by_symbol, writer)?;
+        let written_ab = abilities.serialize(writer)?;
+        let written_solved_impls =
+            crate::abilities::serialize_solved_implementations(solved_implementations, writer)?;
+
+        Ok(written_subs + written_ab + written_solved_impls)
+    }
+
+    pub fn deserialize(bytes: &[u8]) -> (Self, usize) {
+        let ((subs, exposed_vars_by_symbol), len_subs) = Subs::deserialize(bytes);
+        let bytes = &bytes[len_subs..];
+
+        let (abilities, len_abilities) = AbilitiesStore::deserialize(bytes);
+        let bytes = &bytes[len_abilities..];
+
+        let (solved_implementations, len_solved_impls) =
+            crate::abilities::deserialize_solved_implementations(bytes);
+
+        let total_offset = len_subs + len_abilities + len_solved_impls;
+
+        (
+            Self {
+                subs,
+                exposed_vars_by_symbol: exposed_vars_by_symbol.to_vec(),
+                abilities,
+                solved_implementations,
+            },
+            total_offset,
+        )
+    }
+}