Merge pull request #3723 from rtfeldman/wasm-stack-size-control-i2490

CLI argument for Wasm dev stack size

crates/cli/src/build.rs

@@ -48,6 +48,7 @@ pub fn build_file<'a>(
     precompiled: bool,
     target_valgrind: bool,
     threading: Threading,
+    wasm_dev_stack_bytes: Option<u32>,
 ) -> Result<BuiltFile, LoadingProblem<'a>> {
     let compilation_start = Instant::now();
     let target_info = TargetInfo::from(target);
@@ -246,6 +247,7 @@ pub fn build_file<'a>(
         opt_level,
         emit_debug_info,
         &preprocessed_host_path,
+        wasm_dev_stack_bytes,
     );
 
     buf.push('\n');

crates/cli/src/lib.rs

@@ -59,6 +59,7 @@ pub const FLAG_LINKER: &str = "linker";
 pub const FLAG_PRECOMPILED: &str = "precompiled-host";
 pub const FLAG_VALGRIND: &str = "valgrind";
 pub const FLAG_CHECK: &str = "check";
+pub const FLAG_WASM_STACK_SIZE_KB: &str = "wasm-stack-size-kb";
 pub const ROC_FILE: &str = "ROC_FILE";
 pub const ROC_DIR: &str = "ROC_DIR";
 pub const GLUE_FILE: &str = "GLUE_FILE";
@@ -117,6 +118,13 @@ pub fn build_app<'a>() -> Command<'a> {
         .possible_values(["true", "false"])
         .required(false);
 
+    let flag_wasm_stack_size_kb = Arg::new(FLAG_WASM_STACK_SIZE_KB)
+        .long(FLAG_WASM_STACK_SIZE_KB)
+        .help("Stack size in kilobytes for wasm32 target. Only applies when --dev also provided.")
+        .takes_value(true)
+        .validator(|s| s.parse::<u32>())
+        .required(false);
+
     let roc_file_to_run = Arg::new(ROC_FILE)
         .help("The .roc file of an app to run")
         .allow_invalid_utf8(true)
@@ -145,6 +153,7 @@ pub fn build_app<'a>() -> Command<'a> {
             .arg(flag_linker.clone())
             .arg(flag_precompiled.clone())
             .arg(flag_valgrind.clone())
+            .arg(flag_wasm_stack_size_kb.clone())
             .arg(
                 Arg::new(FLAG_TARGET)
                     .long(FLAG_TARGET)
@@ -515,6 +524,11 @@ pub fn build(
         process::exit(1);
     }
 
+    let wasm_dev_stack_bytes: Option<u32> = matches
+        .value_of(FLAG_WASM_STACK_SIZE_KB)
+        .and_then(|s| s.parse::<u32>().ok())
+        .map(|x| x * 1024);
+
     let target_valgrind = matches.is_present(FLAG_VALGRIND);
     let res_binary_path = build_file(
         &arena,
@@ -528,6 +542,7 @@ pub fn build(
         precompiled,
         target_valgrind,
         threading,
+        wasm_dev_stack_bytes,
     );
 
     match res_binary_path {

crates/compiler/build/src/program.rs

@@ -166,6 +166,7 @@ pub fn gen_from_mono_module(
     opt_level: OptLevel,
     emit_debug_info: bool,
     preprocessed_host_path: &Path,
+    wasm_dev_stack_bytes: Option<u32>,
 ) -> CodeGenTiming {
     match opt_level {
         OptLevel::Normal | OptLevel::Size | OptLevel::Optimize => gen_from_mono_module_llvm(
@@ -177,9 +178,14 @@ pub fn gen_from_mono_module(
             opt_level,
             emit_debug_info,
         ),
-        OptLevel::Development => {
-            gen_from_mono_module_dev(arena, loaded, target, app_o_file, preprocessed_host_path)
-        }
+        OptLevel::Development => gen_from_mono_module_dev(
+            arena,
+            loaded,
+            target,
+            app_o_file,
+            preprocessed_host_path,
+            wasm_dev_stack_bytes,
+        ),
     }
 }
 
@@ -416,13 +422,18 @@ pub fn gen_from_mono_module_dev(
     target: &target_lexicon::Triple,
     app_o_file: &Path,
     preprocessed_host_path: &Path,
+    wasm_dev_stack_bytes: Option<u32>,
 ) -> CodeGenTiming {
     use target_lexicon::Architecture;
 
     match target.architecture {
-        Architecture::Wasm32 => {
-            gen_from_mono_module_dev_wasm32(arena, loaded, app_o_file, preprocessed_host_path)
-        }
+        Architecture::Wasm32 => gen_from_mono_module_dev_wasm32(
+            arena,
+            loaded,
+            app_o_file,
+            preprocessed_host_path,
+            wasm_dev_stack_bytes,
+        ),
         Architecture::X86_64 | Architecture::Aarch64(_) => {
             gen_from_mono_module_dev_assembly(arena, loaded, target, app_o_file)
         }
@@ -437,6 +448,7 @@ pub fn gen_from_mono_module_dev(
     target: &target_lexicon::Triple,
     app_o_file: &Path,
     _host_input_path: &Path,
+    _wasm_dev_stack_bytes: Option<u32>,
 ) -> CodeGenTiming {
     use target_lexicon::Architecture;
 
@@ -454,6 +466,7 @@ fn gen_from_mono_module_dev_wasm32(
     loaded: MonomorphizedModule,
     app_o_file: &Path,
     preprocessed_host_path: &Path,
+    wasm_dev_stack_bytes: Option<u32>,
 ) -> CodeGenTiming {
     let code_gen_start = Instant::now();
     let MonomorphizedModule {
@@ -474,6 +487,7 @@ fn gen_from_mono_module_dev_wasm32(
         arena,
         module_id,
         exposed_to_host,
+        stack_bytes: wasm_dev_stack_bytes.unwrap_or(roc_gen_wasm::Env::DEFAULT_STACK_BYTES),
     };
 
     let host_bytes = std::fs::read(preprocessed_host_path).unwrap_or_else(|_| {

crates/compiler/gen_wasm/src/backend.rs

@@ -82,11 +82,7 @@ impl<'a> WasmBackend<'a> {
         fn_index_offset: u32,
         helper_proc_gen: CodeGenHelp<'a>,
     ) -> Self {
-        // TODO: get this from a CLI parameter with some default
-        const STACK_SIZE: u32 = 1024 * 1024;
-        let can_relocate_heap = Self::set_memory_layout(env, &mut module, STACK_SIZE);
-
-        Self::export_globals(&mut module);
+        let can_relocate_heap = module.linking.find_internal_symbol("__heap_base").is_ok();
 
         // We don't want to import any Memory or Tables
         module.import.imports.retain(|import| {
@@ -140,8 +136,8 @@ impl<'a> WasmBackend<'a> {
     /// Since they're all in one block, they can't grow independently. Only the highest one can grow.
     /// Also, there's no "invalid region" below the stack, so stack overflow will overwrite constants!
     /// TODO: Detect stack overflow in function prologue... at least in Roc code...
-    fn set_memory_layout(env: &'a Env<'a>, module: &mut WasmModule<'a>, stack_size: u32) -> bool {
-        let mut stack_heap_boundary = module.data.end_addr + stack_size;
+    fn set_memory_layout(&mut self, stack_size: u32) {
+        let mut stack_heap_boundary = self.module.data.end_addr + stack_size;
         stack_heap_boundary = round_up_to_alignment!(stack_heap_boundary, MemorySection::PAGE_SIZE);
 
         // Stack pointer
@@ -155,12 +151,12 @@ impl<'a> WasmBackend<'a> {
             // Check that __stack_pointer is the only imported global
             // If there were more, we'd have to relocate them, and we don't
             let imported_globals = Vec::from_iter_in(
-                module
+                self.module
                     .import
                     .imports
                     .iter()
                     .filter(|import| matches!(import.description, ImportDesc::Global { .. })),
-                env.arena,
+                self.env.arena,
             );
             if imported_globals.len() != 1
                 || imported_globals[0]
@@ -173,49 +169,55 @@ impl<'a> WasmBackend<'a> {
                 panic!("I can't link this host file. I expected it to have one imported Global called env.__stack_pointer")
             }
         }
-        module
+        self.module
             .import
             .imports
             .retain(|import| !matches!(import.description, ImportDesc::Global { .. }));
-        module.global.append(Global {
+
+        self.module.global.append(Global {
             ty: sp_type,
             init: ConstExpr::I32(stack_heap_boundary as i32),
         });
 
         // Set the initial size of the memory
-        module.memory =
-            MemorySection::new(env.arena, stack_heap_boundary + MemorySection::PAGE_SIZE);
+        self.module.memory = MemorySection::new(
+            self.env.arena,
+            stack_heap_boundary + MemorySection::PAGE_SIZE,
+        );
 
         // Export the memory so that JS can interact with it
-        module.export.append(Export {
+        self.module.export.append(Export {
             name: MEMORY_NAME,
             ty: ExportType::Mem,
             index: 0,
         });
 
         // Set the constant that malloc uses to know where the heap begins
-        module
-            .relocate_internal_symbol("__heap_base", stack_heap_boundary)
-            .is_ok()
+        // this should be done after we know how much constant data we have (e.g. string literals)
+        if self.can_relocate_heap {
+            self.module
+                .relocate_internal_symbol("__heap_base", stack_heap_boundary)
+                .unwrap();
+        }
     }
 
     /// If the host has some `extern` global variables, we need to create them in the final binary
     /// and make them visible to JavaScript by exporting them
-    fn export_globals(module: &mut WasmModule<'a>) {
-        for (sym_index, sym) in module.linking.symbol_table.iter().enumerate() {
+    fn export_globals(&mut self) {
+        for (sym_index, sym) in self.module.linking.symbol_table.iter().enumerate() {
             match sym {
                 SymInfo::Data(DataSymbol::Imported { name, .. }) if *name != "__heap_base" => {
-                    let global_value_addr = module.data.end_addr;
-                    module.data.end_addr += PTR_SIZE;
+                    let global_value_addr = self.module.data.end_addr;
+                    self.module.data.end_addr += PTR_SIZE;
 
-                    module.reloc_code.apply_relocs_u32(
-                        &mut module.code.preloaded_bytes,
+                    self.module.reloc_code.apply_relocs_u32(
+                        &mut self.module.code.preloaded_bytes,
                         sym_index as u32,
                         global_value_addr,
                     );
 
-                    let global_index = module.global.count;
-                    module.global.append(Global {
+                    let global_index = self.module.global.count;
+                    self.module.global.append(Global {
                         ty: GlobalType {
                             value_type: ValueType::I32,
                             is_mutable: false,
@@ -223,7 +225,7 @@ impl<'a> WasmBackend<'a> {
                         init: ConstExpr::I32(global_value_addr as i32),
                     });
 
-                    module.export.append(Export {
+                    self.module.export.append(Export {
                         name,
                         ty: ExportType::Global,
                         index: global_index,
@@ -270,6 +272,9 @@ impl<'a> WasmBackend<'a> {
     }
 
     pub fn finalize(mut self) -> (WasmModule<'a>, BitVec<usize>) {
+        self.set_memory_layout(self.env.stack_bytes);
+        self.export_globals();
+
         self.maybe_call_host_main();
         let fn_table_size = 1 + self.module.element.max_table_index();
         self.module.table.function_table.limits = Limits::MinMax(fn_table_size, fn_table_size);

crates/compiler/gen_wasm/src/lib.rs

@@ -45,6 +45,11 @@ pub struct Env<'a> {
     pub arena: &'a Bump,
     pub module_id: ModuleId,
     pub exposed_to_host: MutSet<Symbol>,
+    pub stack_bytes: u32,
+}
+
+impl Env<'_> {
+    pub const DEFAULT_STACK_BYTES: u32 = 1024 * 1024;
 }
 
 /// Parse the preprocessed host binary

crates/compiler/test_gen/src/helpers/wasm.rs

@@ -122,6 +122,7 @@ fn compile_roc_to_wasm_bytes<'a, T: Wasm32Result>(
         arena,
         module_id,
         exposed_to_host,
+        stack_bytes: roc_gen_wasm::Env::DEFAULT_STACK_BYTES,
     };
 
     let host_module = roc_gen_wasm::parse_host(env.arena, host_bytes).unwrap_or_else(|e| {

crates/compiler/test_gen/src/wasm_linking.rs

@@ -161,6 +161,7 @@ impl<'a> BackendInputs<'a> {
             arena,
             module_id,
             exposed_to_host,
+            stack_bytes: Env::DEFAULT_STACK_BYTES,
         };
 
         // Identifier stuff for the backend

crates/repl_wasm/src/repl.rs

@@ -213,6 +213,7 @@ pub async fn entrypoint_from_js(src: String) -> Result<String, String> {
         let env = roc_gen_wasm::Env {
             arena,
             module_id,
+            stack_bytes: roc_gen_wasm::Env::DEFAULT_STACK_BYTES,
             exposed_to_host: exposed_to_host
                 .values
                 .keys()