Move external specializations storage subs into its own struct

crates/compiler/late_solve/src/lib.rs

@@ -22,6 +22,8 @@ use roc_unify::unify::{Env, Mode, Unified};
 pub use roc_solve::ability::Resolved;
 pub use roc_types::subs::instantiate_rigids;
 
+pub mod storage;
+
 #[derive(Debug)]
 pub struct UnificationFailed;
 

crates/compiler/mono/src/ir.rs

@@ -23,6 +23,7 @@ use roc_derive::SharedDerivedModule;
 use roc_error_macros::{internal_error, todo_abilities};
 use roc_exhaustive::{Ctor, CtorName, RenderAs, TagId};
 use roc_intern::Interner;
+use roc_late_solve::storage::{ExternalModuleStorage, ExternalModuleStorageSnapshot};
 use roc_late_solve::{resolve_ability_specialization, AbilitiesView, Resolved, UnificationFailed};
 use roc_module::ident::{ForeignSymbol, Lowercase, TagName};
 use roc_module::low_level::LowLevel;
@@ -32,8 +33,8 @@ use roc_region::all::{Loc, Region};
 use roc_std::RocDec;
 use roc_target::TargetInfo;
 use roc_types::subs::{
-    instantiate_rigids, Content, ExhaustiveMark, FlatType, RedundantMark, StorageSnapshot,
+    instantiate_rigids, Content, ExhaustiveMark, FlatType, RedundantMark, StorageSubs, Subs,
-    StorageSubs, Subs, Variable, VariableSubsSlice,
+    Variable, VariableSubsSlice,
 };
 use std::collections::HashMap;
 use ven_pretty::{BoxAllocator, DocAllocator, DocBuilder};
@@ -557,14 +558,15 @@ impl<'a> HostSpecializations<'a> {
     }
 }
 
-/// Specializations of this module's symbols that other modules need
+/// Specializations of this module's symbols that other modules need.
+/// One struct represents one pair of modules, e.g. what module A wants of module B.
 #[derive(Clone, Debug)]
 pub struct ExternalSpecializations<'a> {
     /// Not a bumpalo vec because bumpalo is not thread safe
     /// Separate array so we can search for membership quickly
     /// If it's a value and not a lambda, the value is recorded as LambdaName::no_niche.
     pub symbol_or_lambda: std::vec::Vec<LambdaName<'a>>,
-    storage_subs: StorageSubs,
+    storage: ExternalModuleStorage,
     /// For each symbol, what types to specialize it for, points into the storage_subs
     types_to_specialize: std::vec::Vec<std::vec::Vec<Variable>>,
 }
@@ -579,7 +581,7 @@ impl<'a> ExternalSpecializations<'a> {
     pub fn new() -> Self {
         Self {
             symbol_or_lambda: std::vec::Vec::new(),
-            storage_subs: StorageSubs::new(Subs::default()),
+            storage: ExternalModuleStorage::new(Subs::default()),
             types_to_specialize: std::vec::Vec::new(),
         }
     }
@@ -590,7 +592,7 @@ impl<'a> ExternalSpecializations<'a> {
         env_subs: &mut Subs,
         variable: Variable,
     ) {
-        let stored_variable = self.storage_subs.extend_with_variable(env_subs, variable);
+        let stored_variable = self.storage.extend_with_variable(env_subs, variable);
         roc_tracing::debug!(original = ?variable, stored = ?stored_variable, "stored needed external");
 
         match self
@@ -616,27 +618,27 @@ impl<'a> ExternalSpecializations<'a> {
         impl Iterator<Item = (LambdaName<'a>, std::vec::Vec<Variable>)>,
     ) {
         (
-            self.storage_subs,
+            self.storage.into_storage_subs(),
             self.symbol_or_lambda
                 .into_iter()
                 .zip(self.types_to_specialize.into_iter()),
         )
     }
 
-    fn snapshot_cache(&mut self) -> StorageSnapshot {
+    fn snapshot_cache(&mut self) -> ExternalModuleStorageSnapshot {
-        self.storage_subs.snapshot_cache()
+        self.storage.snapshot_cache()
     }
 
-    fn rollback_cache(&mut self, snapshot: StorageSnapshot) {
+    fn rollback_cache(&mut self, snapshot: ExternalModuleStorageSnapshot) {
-        self.storage_subs.rollback_cache(snapshot)
+        self.storage.rollback_cache(snapshot)
     }
 
     fn invalidate_cache(&mut self, changed_variables: &[Variable]) {
-        self.storage_subs.invalidate_cache(changed_variables)
+        self.storage.invalidate_cache(changed_variables)
     }
 
     fn invalidate_whole_cache(&mut self) {
-        self.storage_subs.invalidate_whole_cache()
+        self.storage.invalidate_whole_cache()
     }
 }
 
@@ -684,7 +686,7 @@ impl<'a> Suspended<'a> {
         self.symbol_or_lambdas.push(symbol_or_lambda);
         self.layouts.push(proc_layout);
 
-        let variable = self.store.extend_with_variable(subs, variable);
+        let variable = self.store.import_variable_from(subs, variable).variable;
 
         self.variables.push(variable);
     }
@@ -3114,7 +3116,7 @@ fn generate_runtime_error_function<'a>(
 struct TypeStateSnapshot {
     subs_snapshot: roc_types::subs::SubsSnapshot,
     layout_snapshot: crate::layout::CacheSnapshot,
-    external_storage_snapshot: VecMap<ModuleId, StorageSnapshot>,
+    external_storage_snapshot: VecMap<ModuleId, ExternalModuleStorageSnapshot>,
 }
 
 /// Takes a snapshot of the type state. Snapshots should be taken before new specializations, and

crates/compiler/types/src/subs.rs

@@ -3996,32 +3996,9 @@ pub struct ExposedTypesStorageSubs {
     pub stored_ability_member_vars: VecMap<Variable, Variable>,
 }
 
-#[derive(Clone, Debug)]
-struct VariableMapCache(Vec<FnvMap<Variable, Variable>>);
-
-impl VariableMapCache {
-    fn new() -> Self {
-        Self(vec![Default::default()])
-    }
-
-    fn get(&self, v: &Variable) -> Option<&Variable> {
-        self.0.iter().rev().find_map(|cache| cache.get(v))
-    }
-
-    fn insert(&mut self, key: Variable, value: Variable) -> Option<Variable> {
-        self.0.last_mut().unwrap().insert(key, value)
-    }
-}
-
 #[derive(Clone, Debug)]
 pub struct StorageSubs {
     subs: Subs,
-    /// Variable they expose -> variable we record into storage
-    variable_mapping_cache: VariableMapCache,
-}
-
-pub struct StorageSnapshot {
-    mapping_cache_len: usize,
 }
 
 #[derive(Copy, Clone, Debug)]
@@ -4037,12 +4014,26 @@ struct StorageSubsOffsets {
     problems: u32,
 }
 
+#[derive(Clone, Debug)]
+pub struct VariableMapCache(pub Vec<FnvMap<Variable, Variable>>);
+
+impl VariableMapCache {
+    pub fn new() -> Self {
+        Self(vec![Default::default()])
+    }
+
+    fn get(&self, v: &Variable) -> Option<&Variable> {
+        self.0.iter().rev().find_map(|cache| cache.get(v))
+    }
+
+    fn insert(&mut self, key: Variable, value: Variable) -> Option<Variable> {
+        self.0.last_mut().unwrap().insert(key, value)
+    }
+}
+
 impl StorageSubs {
     pub fn new(subs: Subs) -> Self {
-        Self {
+        Self { subs }
-            subs,
-            variable_mapping_cache: VariableMapCache::new(),
-        }
     }
 
     pub fn fresh_unnamed_flex_var(&mut self) -> Variable {
@@ -4059,41 +4050,13 @@ impl StorageSubs {
 
     pub fn extend_with_variable(&mut self, source: &Subs, variable: Variable) -> Variable {
         storage_copy_var_to(
-            &mut self.variable_mapping_cache,
+            &mut VariableMapCache::new(),
             source,
             &mut self.subs,
             variable,
         )
     }
 
-    pub fn invalidate_cache(&mut self, changed_variables: &[Variable]) {
-        for var in changed_variables {
-            for cache in self.variable_mapping_cache.0.iter_mut().rev() {
-                cache.remove(var);
-            }
-        }
-    }
-
-    pub fn invalidate_whole_cache(&mut self) {
-        debug_assert_eq!(self.variable_mapping_cache.0.len(), 1);
-        self.variable_mapping_cache.0.last_mut().unwrap().clear();
-    }
-
-    pub fn snapshot_cache(&mut self) -> StorageSnapshot {
-        self.variable_mapping_cache.0.push(Default::default());
-        StorageSnapshot {
-            mapping_cache_len: self.variable_mapping_cache.0.len(),
-        }
-    }
-
-    pub fn rollback_cache(&mut self, snapshot: StorageSnapshot) {
-        debug_assert_eq!(
-            self.variable_mapping_cache.0.len(),
-            snapshot.mapping_cache_len
-        );
-        self.variable_mapping_cache.0.pop();
-    }
-
     pub fn import_variable_from(&mut self, source: &Subs, variable: Variable) -> CopiedImport {
         copy_import_to(source, &mut self.subs, false, variable, Rank::import())
     }
@@ -4384,7 +4347,7 @@ fn put_scratchpad(scratchpad: bumpalo::Bump) {
     });
 }
 
-fn storage_copy_var_to(
+pub fn storage_copy_var_to(
     copy_table: &mut VariableMapCache,
     source: &Subs,
     target: &mut Subs,