wasm: use a BitVec for called host functions

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<usize>,
     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<usize>) {
         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);
     }

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<usize>, 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);

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<usize>) {
         if DEBUG_SETTINGS.skip_dead_code_elim {
             return;
         }
@@ -308,7 +308,7 @@ impl<'a> WasmModule<'a> {
     fn trace_live_host_functions<I: Iterator<Item = u32>>(
         &self,
         arena: &'a Bump,
-        called_host_fns: &[u32],
+        called_host_fns: BitVec<usize>,
         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, called_host_fns.len());
-        let mut live_flags = BitVec::repeat(false, capacity);
+        let mut next_pass_fns = BitVec::repeat(false, called_host_fns.len());
-        let mut next_pass_fns = BitVec::repeat(false, capacity);
+        let mut current_pass_fns = called_host_fns;
-        let mut current_pass_fns = BitVec::<usize>::repeat(false, capacity);
+        for index in exported_fns {
-
-        // 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) {
             current_pass_fns.set(index as usize, true);
         }
 

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);

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);

repl_wasm/src/repl.rs

@@ -227,7 +227,7 @@ pub async fn entrypoint_from_js(src: String) -> Result<String, String> {
             &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);