Add a LayoutBuffer to store layouts aside and thread it through

crates/compiler/load_internal/src/file.rs

@@ -38,6 +38,7 @@ use roc_mono::ir::{
 use roc_mono::layout::{
     CapturesNiche, LambdaName, Layout, LayoutCache, LayoutProblem, STLayoutInterner,
 };
+use roc_mono::LayoutBuffer;
 use roc_packaging::cache::{self, RocCacheDir};
 #[cfg(not(target_family = "wasm"))]
 use roc_packaging::https::PackageMetadata;
@@ -135,9 +136,11 @@ struct ModuleCache<'a> {
     pending_abilities: MutMap<ModuleId, PendingAbilitiesStore>,
     constrained: MutMap<ModuleId, ConstrainedModule>,
     typechecked: MutMap<ModuleId, TypeCheckedModule<'a>>,
+
     found_specializations: MutMap<ModuleId, FoundSpecializationsModule<'a>>,
     late_specializations: MutMap<ModuleId, LateSpecializationsModule<'a>>,
     external_specializations_requested: MutMap<ModuleId, Vec<ExternalSpecializations<'a>>>,
+    layout_buffers: VecMap<ModuleId, LayoutBuffer<'a>>,
     expectations: VecMap<ModuleId, Expectations>,
 
     /// Various information
@@ -209,6 +212,7 @@ impl Default for ModuleCache<'_> {
             found_specializations: Default::default(),
             late_specializations: Default::default(),
             external_specializations_requested: Default::default(),
+            layout_buffers: Default::default(),
             imports: Default::default(),
             top_level_thunks: Default::default(),
             documentation: Default::default(),
@@ -454,8 +458,11 @@ fn start_phase<'a>(
                 let build_expects = matches!(state.exec_mode, ExecutionMode::Test)
                     && state.module_cache.expectations.contains_key(&module_id);
 
+                let layout_buffer = LayoutBuffer::default();
+
                 BuildTask::BuildPendingSpecializations {
                     layout_cache,
+                    layout_buffer,
                     module_id,
                     module_timing,
                     solved_subs,
@@ -495,6 +502,10 @@ fn start_phase<'a>(
                         // This is the first time the derived module is introduced into the load
                         // graph. It has no abilities of its own or anything, just generate fresh
                         // information for it.
+                        state
+                            .module_cache
+                            .layout_buffers
+                            .insert(ModuleId::DERIVED_GEN, Default::default());
                         (
                             IdentIds::default(),
                             Subs::default(),
@@ -550,6 +561,13 @@ fn start_phase<'a>(
                         (ident_ids, subs, procs_base, layout_cache, module_timing)
                     };
 
+                let mut layout_buffer = state
+                    .module_cache
+                    .layout_buffers
+                    .remove(&module_id)
+                    .unwrap()
+                    .1;
+
                 if module_id == ModuleId::DERIVED_GEN {
                     load_derived_partial_procs(
                         module_id,
@@ -562,6 +580,7 @@ fn start_phase<'a>(
                         &state.exposed_types,
                         &mut procs_base,
                         &mut state.world_abilities,
+                        &mut layout_buffer,
                     );
                 }
 
@@ -573,6 +592,7 @@ fn start_phase<'a>(
                     subs,
                     procs_base,
                     layout_cache,
+                    layout_buffer,
                     specializations_we_must_make,
                     module_timing,
                     world_abilities: state.world_abilities.clone_ref(),
@@ -712,6 +732,7 @@ pub struct MonomorphizedModule<'a> {
     pub interns: Interns,
     pub subs: Subs,
     pub layout_interner: SingleThreadedInterner<'a, Layout<'a>>,
+    pub layout_buffers: VecMap<ModuleId, LayoutBuffer<'a>>,
     pub output_path: Box<Path>,
     pub can_problems: MutMap<ModuleId, Vec<roc_problem::can::Problem>>,
     pub type_problems: MutMap<ModuleId, Vec<TypeError>>,
@@ -1138,6 +1159,7 @@ enum BuildTask<'a> {
     BuildPendingSpecializations {
         module_timing: ModuleTiming,
         layout_cache: LayoutCache<'a>,
+        layout_buffer: LayoutBuffer<'a>,
         solved_subs: Solved<Subs>,
         imported_module_thunks: &'a [Symbol],
         module_id: ModuleId,
@@ -1155,6 +1177,7 @@ enum BuildTask<'a> {
         subs: Subs,
         procs_base: ProcsBase<'a>,
         layout_cache: LayoutCache<'a>,
+        layout_buffer: LayoutBuffer<'a>,
         specializations_we_must_make: Vec<ExternalSpecializations<'a>>,
         module_timing: ModuleTiming,
         exposed_by_module: ExposedByModule,
@@ -3158,6 +3181,7 @@ fn finish_specialization<'a>(
         type_problems,
         can_problems,
         sources,
+        layout_buffers,
         ..
     } = module_cache;
 
@@ -3257,6 +3281,7 @@ fn finish_specialization<'a>(
         timings: state.timings,
         toplevel_expects,
         uses_prebuilt_platform,
+        layout_buffers,
     })
 }
 
@@ -5224,6 +5249,7 @@ fn make_specializations<'a>(
     mut subs: Subs,
     procs_base: ProcsBase<'a>,
     mut layout_cache: LayoutCache<'a>,
+    mut layout_buffer: LayoutBuffer<'a>,
     specializations_we_must_make: Vec<ExternalSpecializations<'a>>,
     mut module_timing: ModuleTiming,
     target_info: TargetInfo,
@@ -5246,6 +5272,7 @@ fn make_specializations<'a>(
         abilities: AbilitiesView::World(&world_abilities),
         exposed_by_module,
         derived_module: &derived_module,
+        layout_buffer: &mut layout_buffer,
     };
 
     let mut procs = Procs::new_in(arena);
@@ -5302,6 +5329,7 @@ fn build_pending_specializations<'a>(
     declarations: Declarations,
     mut module_timing: ModuleTiming,
     mut layout_cache: LayoutCache<'a>,
+    mut layout_buffer: LayoutBuffer<'a>,
     target_info: TargetInfo,
     exposed_to_host: ExposedToHost,
     exposed_by_module: &ExposedByModule,
@@ -5339,6 +5367,7 @@ fn build_pending_specializations<'a>(
         abilities: AbilitiesView::Module(&abilities_store),
         exposed_by_module,
         derived_module: &derived_module,
+        layout_buffer: &mut layout_buffer,
     };
 
     // Add modules' decls to Procs
@@ -5734,6 +5763,7 @@ fn load_derived_partial_procs<'a>(
     exposed_by_module: &ExposedByModule,
     procs_base: &mut ProcsBase<'a>,
     world_abilities: &mut WorldAbilities,
+    layout_buffer: &mut LayoutBuffer<'a>,
 ) {
     debug_assert_eq!(home, ModuleId::DERIVED_GEN);
 
@@ -5769,6 +5799,7 @@ fn load_derived_partial_procs<'a>(
             abilities: AbilitiesView::World(world_abilities),
             exposed_by_module,
             derived_module,
+            layout_buffer,
         };
 
         let partial_proc = match derived_expr {
@@ -5910,6 +5941,7 @@ fn run_task<'a>(
             decls,
             module_timing,
             layout_cache,
+            layout_buffer,
             solved_subs,
             imported_module_thunks,
             exposed_to_host,
@@ -5926,6 +5958,7 @@ fn run_task<'a>(
             decls,
             module_timing,
             layout_cache,
+            layout_buffer,
             target_info,
             exposed_to_host,
             &exposed_by_module,
@@ -5939,6 +5972,7 @@ fn run_task<'a>(
             subs,
             procs_base,
             layout_cache,
+            layout_buffer,
             specializations_we_must_make,
             module_timing,
             world_abilities,
@@ -5951,6 +5985,7 @@ fn run_task<'a>(
             subs,
             procs_base,
             layout_cache,
+            layout_buffer,
             specializations_we_must_make,
             module_timing,
             target_info,

crates/compiler/mono/src/ir.rs

@@ -5,6 +5,7 @@ use crate::layout::{
     LambdaName, LambdaSet, Layout, LayoutCache, LayoutInterner, LayoutProblem, RawFunctionLayout,
     STLayoutInterner, TagIdIntType, UnionLayout, WrappedVariant,
 };
+use crate::LayoutBuffer;
 use bumpalo::collections::{CollectIn, Vec};
 use bumpalo::Bump;
 use roc_builtins::bitcode::{FloatWidth, IntWidth};
@@ -1474,6 +1475,7 @@ pub struct Env<'a, 'i> {
     pub abilities: AbilitiesView<'i>,
     pub exposed_by_module: &'i ExposedByModule,
     pub derived_module: &'i SharedDerivedModule,
+    pub layout_buffer: &'i mut LayoutBuffer<'a>,
 }
 
 impl<'a, 'i> Env<'a, 'i> {

crates/compiler/mono/src/layout_buffer.rs

@@ -0,0 +1,24 @@
+use roc_collections::soa::Slice;
+
+use crate::layout::Layout;
+
+/// A look-aside buffer to keep layouts in.
+#[derive(Default, Debug)]
+pub struct LayoutBuffer<'a>(Vec<Layout<'a>>);
+
+impl<'a> LayoutBuffer<'a> {
+    pub fn reserve(&mut self, size: usize) -> Slice<Layout<'a>> {
+        Slice::extend_new(&mut self.0, std::iter::repeat(Layout::VOID).take(size))
+    }
+
+    pub fn set_reserved(
+        &mut self,
+        slice: Slice<Layout<'a>>,
+        layouts: impl ExactSizeIterator<Item = Layout<'a>>,
+    ) {
+        debug_assert_eq!(layouts.len(), slice.len());
+        for (index, layout) in slice.indices().zip(layouts) {
+            self.0[index] = layout;
+        }
+    }
+}

crates/compiler/mono/src/lib.rs

@@ -12,6 +12,7 @@ pub mod code_gen_help;
 pub mod inc_dec;
 pub mod ir;
 pub mod layout;
+mod layout_buffer;
 pub mod layout_soa;
 pub mod low_level;
 pub mod reset_reuse;
@@ -23,3 +24,5 @@ pub mod tail_recursion;
 pub mod decision_tree;
 
 pub mod debug;
+
+pub use layout_buffer::LayoutBuffer;