clone storage subs less

compiler/constrain/src/module.rs

@@ -2,10 +2,11 @@ use roc_builtins::std::StdLib;
 use roc_can::constraint::{Constraint, Constraints};
 use roc_can::def::Declaration;
 use roc_collections::all::{MutMap, MutSet};
+use roc_error_macros::internal_error;
 use roc_module::symbol::{ModuleId, Symbol};
 use roc_region::all::{Loc, Region};
 use roc_types::solved_types::{FreeVars, SolvedType};
-use roc_types::subs::{StorageSubs, VarStore, Variable};
+use roc_types::subs::{VarStore, Variable};
 use roc_types::types::{Alias, Problem};
 
 #[derive(Clone, Debug, Default)]
@@ -21,6 +22,61 @@ impl ExposedByModule {
     pub fn get(&self, module_id: &ModuleId) -> Option<&ExposedModuleTypes> {
         self.exposed.get(module_id)
     }
+
+    pub fn get_mut(&mut self, module_id: &ModuleId) -> Option<&mut ExposedModuleTypes> {
+        self.exposed.get_mut(module_id)
+    }
+
+    pub fn retain_modules<'a>(&self, it: impl Iterator<Item = &'a ModuleId>) -> Self {
+        let mut output = Self::default();
+
+        for module_id in it {
+            match self.exposed.get(module_id) {
+                None => {
+                    internal_error!("Module {:?} did not register its exposed values", module_id)
+                }
+                Some(exposed_types) => {
+                    output.exposed.insert(*module_id, exposed_types.clone());
+                }
+            }
+        }
+
+        output
+    }
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct ExposedForModule {
+    pub exposed_by_module: ExposedByModule,
+    pub imported_symbols: Vec<Symbol>,
+}
+
+impl ExposedForModule {
+    pub fn new<'a>(
+        it: impl Iterator<Item = &'a Symbol>,
+        exposed_by_module: ExposedByModule,
+    ) -> Self {
+        let mut imported_symbols = Vec::new();
+
+        for symbol in it {
+            if symbol.is_builtin() {
+                continue;
+            }
+
+            if let Some(ExposedModuleTypes::Valid { .. }) =
+                exposed_by_module.exposed.get(&symbol.module_id())
+            {
+                imported_symbols.push(*symbol);
+            } else {
+                continue;
+            }
+        }
+
+        Self {
+            imported_symbols,
+            exposed_by_module,
+        }
+    }
 }
 
 #[derive(Clone, Debug)]
@@ -111,18 +167,10 @@ pub fn constrain_imports(
 
 pub struct ConstrainableImports {
     pub imported_symbols: Vec<Import>,
-    pub hacky_symbols: Vec<HackyImport>,
     pub imported_aliases: MutMap<Symbol, Alias>,
     pub unused_imports: MutMap<ModuleId, Region>,
 }
 
-#[derive(Debug, Clone)]
-pub struct HackyImport {
-    pub storage_subs: StorageSubs,
-    pub loc_symbol: Loc<Symbol>,
-    pub variable: Variable,
-}
-
 /// Run this before constraining imports.
 ///
 /// Constraining imports is split into two different functions, because this
@@ -136,7 +184,6 @@ pub fn pre_constrain_imports(
     stdlib: &StdLib,
 ) -> ConstrainableImports {
     let mut imported_symbols = Vec::with_capacity(references.len());
-    let mut hacky_symbols = Vec::with_capacity(references.len());
     let mut imported_aliases = MutMap::default();
     let mut unused_imports = imported_modules; // We'll remove these as we encounter them.
 
@@ -203,23 +250,12 @@ pub fn pre_constrain_imports(
                             imported_aliases.insert(k, v);
                         }
 
+                        continue;
+
                         imported_symbols.push(Import {
                             loc_symbol,
                             solved_type: solved_type.clone(),
                         });
-
-                        let variable = stored_vars_by_symbol
-                            .iter()
-                            .find(|(s, _)| *s == loc_symbol.value)
-                            .unwrap()
-                            .1;
-
-                        hacky_symbols.push(HackyImport {
-                            loc_symbol,
-                            variable,
-                            // TODO very bad, so much cloning!
-                            storage_subs: storage_subs.clone(),
-                        });
                     }
                 }
                 Some(ExposedModuleTypes::Invalid) => {
@@ -244,7 +280,6 @@ pub fn pre_constrain_imports(
 
     ConstrainableImports {
         imported_symbols,
-        hacky_symbols,
         imported_aliases,
         unused_imports,
     }

compiler/load/src/file.rs

@@ -11,7 +11,7 @@ use roc_can::module::{canonicalize_module_defs, Module};
 use roc_collections::all::{default_hasher, BumpMap, MutMap, MutSet};
 use roc_constrain::module::{
     constrain_imports, constrain_module, pre_constrain_imports, ConstrainableImports,
-    ExposedByModule, ExposedModuleTypes, HackyImport, Import,
+    ExposedByModule, ExposedForModule, ExposedModuleTypes, Import,
 };
 use roc_module::ident::{Ident, ModuleName, QualifiedModuleName};
 use roc_module::symbol::{
@@ -730,7 +730,7 @@ enum BuildTask<'a> {
         module: Module,
         ident_ids: IdentIds,
         imported_symbols: Vec<Import>,
-        imported_storage_subs: Vec<HackyImport>,
+        exposed_for_module: ExposedForModule,
         module_timing: ModuleTiming,
         constraints: Constraints,
         constraint: ConstraintSoa,
@@ -3048,6 +3048,10 @@ impl<'a> BuildTask<'a> {
     ) -> Self {
         let home = module.module_id;
 
+        let our_imports = exposed_types.retain_modules(imported_modules.keys());
+
+        let exposed_for_module = ExposedForModule::new(module.references.iter(), our_imports);
+
         // Get the constraints for this module's imports. We do this on the main thread
         // to avoid having to lock the map of exposed types, or to clone it
         // (which would be more expensive for the main thread).
@@ -3055,7 +3059,6 @@ impl<'a> BuildTask<'a> {
             imported_symbols,
             imported_aliases: _, // TODO well then... do we even need those?
             unused_imports,
-            hacky_symbols,
         } = pre_constrain_imports(
             home,
             &module.references,
@@ -3069,7 +3072,7 @@ impl<'a> BuildTask<'a> {
             module,
             ident_ids,
             imported_symbols,
-            imported_storage_subs: hacky_symbols,
+            exposed_for_module,
             constraints,
             constraint,
             var_store,
@@ -3087,7 +3090,7 @@ fn run_solve<'a>(
     ident_ids: IdentIds,
     mut module_timing: ModuleTiming,
     imported_symbols: Vec<Import>,
-    imported_storage_subs: Vec<HackyImport>,
+    mut exposed_for_module: ExposedForModule,
     mut constraints: Constraints,
     constraint: ConstraintSoa,
     mut var_store: VarStore,
@@ -3105,9 +3108,10 @@ fn run_solve<'a>(
 
     if NEW_TYPES {
         imported_symbols.retain(|k| {
-            !imported_storage_subs
+            !exposed_for_module
+                .imported_symbols
                 .iter()
-                .any(|i| k.loc_symbol.value == i.loc_symbol.value)
+                .any(|i| k.loc_symbol.value == *i)
         });
     }
 
@@ -3124,18 +3128,37 @@ fn run_solve<'a>(
         ..
     } = module;
 
-    // TODO
-    // if false { debug_assert!(constraint.validate(), "{:?}", &constraint); }
-
     let mut subs = Subs::new_from_varstore(var_store);
 
     let mut import_variables = Vec::new();
 
     if NEW_TYPES {
-        for mut import in imported_storage_subs {
+        for symbol in exposed_for_module.imported_symbols {
-            let copied_import = import
+            match exposed_for_module
-                .storage_subs
+                .exposed_by_module
-                .export_variable_to(&mut subs, import.variable);
+                .get_mut(&symbol.module_id())
+            {
+                Some(t) => match t {
+                    ExposedModuleTypes::Invalid => {
+                        // idk, just skip?
+                        continue;
+                    }
+                    ExposedModuleTypes::Valid {
+                        solved_types: _,
+                        aliases: _,
+                        stored_vars_by_symbol,
+                        storage_subs,
+                    } => {
+                        let variable =
+                            match stored_vars_by_symbol.iter().find(|(s, _)| *s == symbol) {
+                                None => {
+                                    // TODO happens for imported types
+                                    continue;
+                                }
+                                Some((_, x)) => *x,
+                            };
+
+                        let copied_import = storage_subs.export_variable_to(&mut subs, variable);
 
                         // not a typo; rigids are turned into flex during type inference, but when imported we must
                         // consider them rigid variables
@@ -3145,10 +3168,14 @@ fn run_solve<'a>(
                         import_variables.extend(copied_import.registered);
 
                         def_types.push((
-                import.loc_symbol.value,
+                            symbol,
                             Loc::at_zero(roc_types::types::Type::Variable(copied_import.variable)),
                         ));
                     }
+                },
+                None => todo!(),
+            }
+        }
     }
 
     let actual_constraint =
@@ -3807,7 +3834,7 @@ fn run_task<'a>(
             module,
             module_timing,
             imported_symbols,
-            imported_storage_subs,
+            exposed_for_module,
             constraints,
             constraint,
             var_store,
@@ -3820,7 +3847,7 @@ fn run_task<'a>(
             ident_ids,
             module_timing,
             imported_symbols,
-            imported_storage_subs,
+            exposed_for_module,
             constraints,
             constraint,
             var_store,