diff --git a/compiler/gen_wasm/src/backend.rs b/compiler/gen_wasm/src/backend.rs index a5e8477703..29b9d0f8c3 100644 --- a/compiler/gen_wasm/src/backend.rs +++ b/compiler/gen_wasm/src/backend.rs @@ -1,3 +1,4 @@ +use bitvec::vec::BitVec; use bumpalo::collections::{String, Vec}; use code_builder::Align; @@ -53,7 +54,7 @@ pub struct WasmBackend<'a> { module: WasmModule<'a>, layout_ids: LayoutIds<'a>, pub fn_index_offset: u32, - called_preload_fns: Vec<'a, u32>, + called_preload_fns: BitVec, pub proc_lookup: Vec<'a, ProcLookupData<'a>>, host_lookup: Vec<'a, (&'a str, u32)>, helper_proc_gen: CodeGenHelp<'a>, @@ -125,7 +126,7 @@ impl<'a> WasmBackend<'a> { layout_ids, fn_index_offset, - called_preload_fns: Vec::with_capacity_in(2, env.arena), + called_preload_fns: BitVec::repeat(false, host_lookup.len()), proc_lookup, host_lookup, helper_proc_gen, @@ -273,7 +274,7 @@ impl<'a> WasmBackend<'a> { wasm_fn_index } - pub fn finalize(mut self) -> (WasmModule<'a>, Vec<'a, u32>) { + pub fn finalize(mut self) -> (WasmModule<'a>, BitVec) { 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); @@ -328,7 +329,7 @@ impl<'a> WasmBackend<'a> { self.code_builder.build_fn_header_and_footer(&[], 0, None); self.reset(); - self.called_preload_fns.push(main_fn_index); + self.called_preload_fns.set(main_fn_index as usize, true); } /// Register the debug names of Symbols in a global lookup table @@ -1250,7 +1251,7 @@ impl<'a> WasmBackend<'a> { .find(|(fn_name, _)| *fn_name == name) .unwrap_or_else(|| panic!("The Roc app tries to call `{}` but I can't find it!", name)); - self.called_preload_fns.push(*fn_index); + self.called_preload_fns.set(*fn_index as usize, true); self.code_builder .call(*fn_index, num_wasm_args, has_return_val); } diff --git a/compiler/gen_wasm/src/lib.rs b/compiler/gen_wasm/src/lib.rs index ab44287fc5..34d69ad8e8 100644 --- a/compiler/gen_wasm/src/lib.rs +++ b/compiler/gen_wasm/src/lib.rs @@ -8,6 +8,7 @@ pub mod wasm_module; pub mod wasm32_result; pub mod wasm32_sized; +use bitvec::prelude::BitVec; use bumpalo::collections::Vec; use bumpalo::{self, Bump}; @@ -67,7 +68,7 @@ pub fn build_app_binary<'a>( let (mut wasm_module, called_preload_fns, _) = build_app_module(env, interns, host_module, procedures); - wasm_module.eliminate_dead_code(env.arena, &called_preload_fns); + wasm_module.eliminate_dead_code(env.arena, called_preload_fns); let mut buffer = std::vec::Vec::with_capacity(wasm_module.size()); wasm_module.serialize(&mut buffer); @@ -83,7 +84,7 @@ pub fn build_app_module<'a>( interns: &'a mut Interns, host_module: WasmModule<'a>, procedures: MutMap<(Symbol, ProcLayout<'a>), Proc<'a>>, -) -> (WasmModule<'a>, Vec<'a, u32>, u32) { +) -> (WasmModule<'a>, BitVec, u32) { let layout_ids = LayoutIds::default(); let mut procs = Vec::with_capacity_in(procedures.len(), env.arena); let mut proc_lookup = Vec::with_capacity_in(procedures.len() * 2, env.arena); diff --git a/compiler/gen_wasm/src/wasm_module/mod.rs b/compiler/gen_wasm/src/wasm_module/mod.rs index 51594ac5db..cc9abd590a 100644 --- a/compiler/gen_wasm/src/wasm_module/mod.rs +++ b/compiler/gen_wasm/src/wasm_module/mod.rs @@ -174,7 +174,7 @@ impl<'a> WasmModule<'a> { }) } - pub fn eliminate_dead_code(&mut self, arena: &'a Bump, called_host_fns: &[u32]) { + pub fn eliminate_dead_code(&mut self, arena: &'a Bump, called_host_fns: BitVec) { if DEBUG_SETTINGS.skip_dead_code_elim { return; } @@ -308,7 +308,7 @@ impl<'a> WasmModule<'a> { fn trace_live_host_functions>( &self, arena: &'a Bump, - called_host_fns: &[u32], + called_host_fns: BitVec, exported_fns: I, indirect_callees_and_signatures: Vec<'a, (u32, u32)>, host_fn_min: u32, @@ -349,14 +349,10 @@ impl<'a> WasmModule<'a> { ); // Loop variables for the main loop below - let capacity = host_fn_max as usize + self.code.code_builders.len(); - let mut live_flags = BitVec::repeat(false, capacity); - let mut next_pass_fns = BitVec::repeat(false, capacity); - let mut current_pass_fns = BitVec::::repeat(false, capacity); - - // Start with everything called from Roc and everything exported to JS - // Also include everything that can be indirectly called (crude, but good enough!) - for index in called_host_fns.iter().copied().chain(exported_fns) { + let mut live_flags = BitVec::repeat(false, called_host_fns.len()); + let mut next_pass_fns = BitVec::repeat(false, called_host_fns.len()); + let mut current_pass_fns = called_host_fns; + for index in exported_fns { current_pass_fns.set(index as usize, true); } diff --git a/compiler/test_gen/src/helpers/wasm.rs b/compiler/test_gen/src/helpers/wasm.rs index e3389be87e..588d1b0392 100644 --- a/compiler/test_gen/src/helpers/wasm.rs +++ b/compiler/test_gen/src/helpers/wasm.rs @@ -151,7 +151,7 @@ fn compile_roc_to_wasm_bytes<'a, T: Wasm32Result>( index: init_refcount_idx, }); - module.eliminate_dead_code(env.arena, &called_preload_fns); + module.eliminate_dead_code(env.arena, called_preload_fns); let mut app_module_bytes = std::vec::Vec::with_capacity(module.size()); module.serialize(&mut app_module_bytes); diff --git a/compiler/test_gen/src/wasm_linking.rs b/compiler/test_gen/src/wasm_linking.rs index e2a9dbc213..e16ce3799c 100644 --- a/compiler/test_gen/src/wasm_linking.rs +++ b/compiler/test_gen/src/wasm_linking.rs @@ -288,7 +288,7 @@ fn test_linking_with_dce() { let (mut final_module, called_preload_fns, _roc_main_index) = roc_gen_wasm::build_app_module(&env, &mut interns, host_module, procedures); - final_module.eliminate_dead_code(env.arena, &called_preload_fns); + final_module.eliminate_dead_code(env.arena, called_preload_fns); let mut buffer = Vec::with_capacity(final_module.size()); final_module.serialize(&mut buffer); diff --git a/repl_wasm/src/repl.rs b/repl_wasm/src/repl.rs index 3389e8aa25..979632654f 100644 --- a/repl_wasm/src/repl.rs +++ b/repl_wasm/src/repl.rs @@ -227,7 +227,7 @@ pub async fn entrypoint_from_js(src: String) -> Result { &main_fn_layout.result, ); - module.eliminate_dead_code(env.arena, &called_preload_fns); + module.eliminate_dead_code(env.arena, called_preload_fns); let mut buffer = Vec::with_capacity_in(module.size(), arena); module.serialize(&mut buffer);