Add the storage module

crates/compiler/late_solve/src/storage.rs

@@ -0,0 +1,75 @@
+use roc_types::subs::StorageSubs;
+use roc_types::subs::{storage_copy_var_to, Subs, Variable, VariableMapCache};
+use std::iter::Iterator;
+
+/// Storage for types to be sent to an external module, and written to only by one module's subs.
+/// Maintains a cache so that independent writes can re-use types previously inserted into the
+/// storage.
+#[derive(Clone, Debug)]
+pub struct ExternalModuleStorage {
+    storage: StorageSubs,
+    /// Variable they expose -> variable we record into storage
+    variable_mapping_cache: VariableMapCache,
+}
+
+pub struct ExternalModuleStorageSnapshot {
+    mapping_cache_len: usize,
+}
+
+impl ExternalModuleStorage {
+    pub fn new(subs: Subs) -> Self {
+        Self {
+            storage: StorageSubs::new(subs),
+            variable_mapping_cache: VariableMapCache::new(),
+        }
+    }
+
+    pub fn extend_with_variable(&mut self, source: &Subs, variable: Variable) -> Variable {
+        storage_copy_var_to(
+            &mut self.variable_mapping_cache,
+            source,
+            &mut self.storage.as_inner_mut(),
+            variable,
+        )
+    }
+
+    pub fn into_storage_subs(self) -> StorageSubs {
+        self.storage
+    }
+
+    /// Invalidates the whole cache given a sequence of variables that should no longer be indexed
+    /// from the cache.
+    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);
+            }
+        }
+    }
+
+    /// Invalidates the whole cache.
+    /// Should only be called if you need to invalidate the cache but don't have a snapshot.
+    /// Generally you should prefer to create a snapshot and invalidate that snapshot, which avoids
+    /// unnecessary cache invalidation.
+    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();
+    }
+
+    /// Creates a snapshot of the cache, making it suitable for new ephemeral entries.
+    /// The cache can be rolled back to the state it was in prior to the snapshot with [rollback_cache].
+    pub fn snapshot_cache(&mut self) -> ExternalModuleStorageSnapshot {
+        self.variable_mapping_cache.0.push(Default::default());
+        ExternalModuleStorageSnapshot {
+            mapping_cache_len: self.variable_mapping_cache.0.len(),
+        }
+    }
+
+    pub fn rollback_cache(&mut self, snapshot: ExternalModuleStorageSnapshot) {
+        debug_assert_eq!(
+            self.variable_mapping_cache.0.len(),
+            snapshot.mapping_cache_len
+        );
+        self.variable_mapping_cache.0.pop();
+    }
+}