Merge pull request #1300 from rtfeldman/specialize-lowlevel

Polyvariant defunctionalization

compiler/load/src/file.rs

@@ -20,6 +20,7 @@ use roc_module::symbol::{
 };
 use roc_mono::ir::{
     CapturedSymbols, ExternalSpecializations, PartialProc, PendingSpecialization, Proc, Procs,
+    TopLevelFunctionLayout,
 };
 use roc_mono::layout::{Layout, LayoutCache, LayoutProblem};
 use roc_parse::ast::{self, StrLiteral, TypeAnnotation};
@@ -704,7 +705,7 @@ pub struct MonomorphizedModule<'a> {
     pub can_problems: MutMap<ModuleId, Vec<roc_problem::can::Problem>>,
     pub type_problems: MutMap<ModuleId, Vec<solve::TypeError>>,
     pub mono_problems: MutMap<ModuleId, Vec<roc_mono::ir::MonoProblem>>,
-    pub procedures: MutMap<(Symbol, Layout<'a>), Proc<'a>>,
+    pub procedures: MutMap<(Symbol, TopLevelFunctionLayout<'a>), Proc<'a>>,
     pub exposed_to_host: MutMap<Symbol, Variable>,
     pub header_sources: MutMap<ModuleId, (PathBuf, Box<str>)>,
     pub sources: MutMap<ModuleId, (PathBuf, Box<str>)>,
@@ -775,7 +776,7 @@ enum Msg<'a> {
         ident_ids: IdentIds,
         layout_cache: LayoutCache<'a>,
         external_specializations_requested: BumpMap<ModuleId, ExternalSpecializations<'a>>,
-        procedures: MutMap<(Symbol, Layout<'a>), Proc<'a>>,
+        procedures: MutMap<(Symbol, TopLevelFunctionLayout<'a>), Proc<'a>>,
         problems: Vec<roc_mono::ir::MonoProblem>,
         module_timing: ModuleTiming,
         subs: Subs,
@@ -817,7 +818,7 @@ struct State<'a> {
 
     pub module_cache: ModuleCache<'a>,
     pub dependencies: Dependencies<'a>,
-    pub procedures: MutMap<(Symbol, Layout<'a>), Proc<'a>>,
+    pub procedures: MutMap<(Symbol, TopLevelFunctionLayout<'a>), Proc<'a>>,
     pub exposed_to_host: MutMap<Symbol, Variable>,
 
     /// This is the "final" list of IdentIds, after canonicalization and constraint gen
@@ -847,7 +848,8 @@ struct State<'a> {
     /// pending specializations in the same thread.
     pub needs_specialization: MutSet<ModuleId>,
 
-    pub all_pending_specializations: MutMap<Symbol, MutMap<Layout<'a>, PendingSpecialization<'a>>>,
+    pub all_pending_specializations:
+        MutMap<Symbol, MutMap<TopLevelFunctionLayout<'a>, PendingSpecialization<'a>>>,
 
     pub specializations_in_flight: u32,
 
@@ -2046,6 +2048,19 @@ fn update<'a>(
                 && state.dependencies.solved_all()
                 && state.goal_phase == Phase::MakeSpecializations
             {
+                // display the mono IR of the module, for debug purposes
+                if roc_mono::ir::PRETTY_PRINT_IR_SYMBOLS {
+                    let procs_string = state
+                        .procedures
+                        .values()
+                        .map(|proc| proc.to_pretty(200))
+                        .collect::<Vec<_>>();
+
+                    let result = procs_string.join("\n");
+
+                    println!("{}", result);
+                }
+
                 if false {
                     let it = state.procedures.iter().map(|x| x.1);
 
@@ -2078,19 +2093,6 @@ fn update<'a>(
                     existing.extend(requested);
                 }
 
-                // display the mono IR of the module, for debug purposes
-                if roc_mono::ir::PRETTY_PRINT_IR_SYMBOLS {
-                    let procs_string = state
-                        .procedures
-                        .values()
-                        .map(|proc| proc.to_pretty(200))
-                        .collect::<Vec<_>>();
-
-                    let result = procs_string.join("\n");
-
-                    println!("{}", result);
-                }
-
                 msg_tx
                     .send(Msg::FinishedAllSpecialization {
                         subs,
@@ -4001,7 +4003,7 @@ fn add_def_to_module<'a>(
 
                         procs.insert_exposed(
                             symbol,
-                            layout,
+                            TopLevelFunctionLayout::from_layout(mono_env.arena, layout),
                             mono_env.arena,
                             mono_env.subs,
                             def.annotation,
@@ -4022,6 +4024,9 @@ fn add_def_to_module<'a>(
                     );
                 }
                 body => {
+                    // mark this symbols as a top-level thunk before any other work on the procs
+                    procs.module_thunks.insert(symbol);
+
                     // If this is an exposed symbol, we need to
                     // register it as such. Otherwise, since it
                     // never gets called by Roc code, it will never
@@ -4034,7 +4039,10 @@ fn add_def_to_module<'a>(
                             annotation,
                             mono_env.subs,
                         ) {
-                            Ok(l) => l,
+                            Ok(l) => {
+                                // remember, this is a 0-argument thunk
+                                Layout::FunctionPointer(&[], mono_env.arena.alloc(l))
+                            }
                             Err(LayoutProblem::Erroneous) => {
                                 let message = "top level function has erroneous type";
                                 procs.runtime_errors.insert(symbol, message);
@@ -4054,7 +4062,7 @@ fn add_def_to_module<'a>(
 
                         procs.insert_exposed(
                             symbol,
-                            layout,
+                            TopLevelFunctionLayout::from_layout(mono_env.arena, layout),
                             mono_env.arena,
                             mono_env.subs,
                             def.annotation,
@@ -4074,7 +4082,6 @@ fn add_def_to_module<'a>(
                     };
 
                     procs.partial_procs.insert(symbol, proc);
-                    procs.module_thunks.insert(symbol);
                 }
             };
         }