trying to gather expects

crates/compiler/can/src/module.rs

@@ -7,7 +7,7 @@ use crate::expr::{ClosureData, Declarations, Expr, Output, PendingDerives};
 use crate::pattern::{BindingsFromPattern, Pattern};
 use crate::scope::Scope;
 use bumpalo::Bump;
-use roc_collections::{MutMap, SendMap, VecSet};
+use roc_collections::{MutMap, SendMap, VecMap, VecSet};
 use roc_error_macros::internal_error;
 use roc_module::ident::Ident;
 use roc_module::ident::Lowercase;
@@ -122,6 +122,7 @@ pub struct Module {
     pub aliases: MutMap<Symbol, (bool, Alias)>,
     pub rigid_variables: RigidVariables,
     pub abilities_store: PendingAbilitiesStore,
+    pub loc_expects: VecMap<Region, Vec<(Symbol, Variable)>>,
 }
 
 #[derive(Debug, Default)]
@@ -144,6 +145,7 @@ pub struct ModuleOutput {
     pub symbols_from_requires: Vec<(Loc<Symbol>, Loc<Type>)>,
     pub pending_derives: PendingDerives,
     pub scope: Scope,
+    pub loc_expects: VecMap<Region, Vec<(Symbol, Variable)>>,
 }
 
 fn validate_generate_with<'a>(
@@ -735,6 +737,8 @@ pub fn canonicalize_module_defs<'a>(
         }
     }
 
+    let loc_expects = Default::default();
+
     ModuleOutput {
         scope,
         aliases,
@@ -747,6 +751,7 @@ pub fn canonicalize_module_defs<'a>(
         symbols_from_requires,
         pending_derives,
         lookups,
+        loc_expects,
     }
 }
 

crates/compiler/load_internal/src/file.rs

@@ -131,6 +131,7 @@ struct ModuleCache<'a> {
     found_specializations: MutMap<ModuleId, FoundSpecializationsModule<'a>>,
     late_specializations: MutMap<ModuleId, LateSpecializationsModule<'a>>,
     external_specializations_requested: MutMap<ModuleId, Vec<ExternalSpecializations<'a>>>,
+    expectations: VecMap<ModuleId, Expectations>,
 
     /// Various information
     imports: MutMap<ModuleId, MutSet<ModuleId>>,
@@ -185,6 +186,7 @@ impl Default for ModuleCache<'_> {
             can_problems: Default::default(),
             type_problems: Default::default(),
             sources: Default::default(),
+            expectations: Default::default(),
         }
     }
 }
@@ -682,6 +684,8 @@ struct ParsedModule<'a> {
     header_for: HeaderFor<'a>,
 }
 
+type LocExpects = VecMap<Region, Vec<(Symbol, Variable)>>;
+
 /// A message sent out _from_ a worker thread,
 /// representing a result of work done, or a request for further work
 #[derive(Debug)]
@@ -699,6 +703,7 @@ enum Msg<'a> {
         dep_idents: IdentIdsByModule,
         module_timing: ModuleTiming,
         abilities_store: AbilitiesStore,
+        loc_expects: LocExpects,
     },
     FinishedAllTypeChecking {
         solved_subs: Solved<Subs>,
@@ -2219,6 +2224,7 @@ fn update<'a>(
             dep_idents,
             mut module_timing,
             abilities_store,
+            loc_expects,
         } => {
             log!("solved types for {:?}", module_id);
             module_timing.end_time = SystemTime::now();
@@ -2228,6 +2234,22 @@ fn update<'a>(
                 .type_problems
                 .insert(module_id, solved_module.problems);
 
+            if !loc_expects.is_empty() {
+                let (path, _) = state.module_cache.sources.get(&module_id).unwrap();
+
+                let expectations = Expectations {
+                    expectations: loc_expects,
+                    subs: solved_subs.clone().into_inner(),
+                    path: path.to_owned(),
+                    ident_ids: ident_ids.clone(),
+                };
+
+                state
+                    .module_cache
+                    .expectations
+                    .insert(module_id, expectations);
+            }
+
             let work = state.dependencies.notify(module_id, Phase::SolveTypes);
 
             // if there is a platform, the `platform` module provides host-exposed,
@@ -4115,6 +4137,8 @@ fn run_solve<'a>(
     // TODO remove when we write builtins in roc
     let aliases = module.aliases.clone();
 
+    let loc_expects = std::mem::take(&mut module.loc_expects);
+
     let (solved_subs, solved_specializations, exposed_vars_by_symbol, problems, abilities_store) = {
         if module_id.is_builtin() {
             match cached_subs.lock().remove(&module_id) {
@@ -4181,6 +4205,7 @@ fn run_solve<'a>(
         solved_module,
         module_timing,
         abilities_store,
+        loc_expects,
     }
 }
 
@@ -4378,6 +4403,7 @@ fn canonicalize_and_constrain<'a>(
         aliases,
         rigid_variables: module_output.rigid_variables,
         abilities_store: module_output.scope.abilities_store,
+        loc_expects,
     };
 
     let constrained_module = ConstrainedModule {