CodeBuilder tracks which local variables are set

crates/compiler/gen_wasm/src/backend.rs

@@ -796,7 +796,7 @@ impl<'a, 'r> WasmBackend<'a, 'r> {
         self.expr(sym, expr, layout, &sym_storage);
 
         if let StoredValue::Local { local_id, .. } = sym_storage {
-            if !self.code_builder.stack_is_empty() {
+            if !self.code_builder.is_set(local_id) {
                 self.code_builder.set_local(local_id);
             }
         }
@@ -1153,7 +1153,11 @@ impl<'a, 'r> WasmBackend<'a, 'r> {
         };
 
         match storage {
-            StoredValue::Local { value_type, .. } => {
+            StoredValue::Local {
+                value_type,
+                local_id,
+                ..
+            } => {
                 match (lit, value_type) {
                     (Literal::Float(x), ValueType::F64) => self.code_builder.f64_const(*x),
                     (Literal::Float(x), ValueType::F32) => self.code_builder.f32_const(*x as f32),
@@ -1167,6 +1171,7 @@ impl<'a, 'r> WasmBackend<'a, 'r> {
                     (Literal::Byte(x), ValueType::I32) => self.code_builder.i32_const(*x as i32),
                     _ => invalid_error(),
                 };
+                self.code_builder.set_local(*local_id);
             }
 
             StoredValue::StackMemory { location, .. } => {

crates/compiler/gen_wasm/src/code_builder.rs

@@ -1,3 +1,4 @@
+use bitvec::vec::BitVec;
 use bumpalo::collections::vec::Vec;
 use bumpalo::Bump;
 use core::panic;
@@ -11,6 +12,7 @@ use roc_wasm_module::{
     round_up_to_alignment, Align, LocalId, RelocationEntry, ValueType, WasmModule,
     FRAME_ALIGNMENT_BYTES, STACK_POINTER_GLOBAL_ID,
 };
+use std::iter::repeat;
 
 use crate::DEBUG_SETTINGS;
 
@@ -106,6 +108,9 @@ pub struct CodeBuilder<'a> {
     /// Relocations for calls to JS imports
     /// When we remove unused imports, the live ones are re-indexed
     import_relocations: Vec<'a, (usize, u32)>,
+
+    /// Keep track of which local variables have been set
+    set_locals: BitVec<u64>,
 }
 
 #[allow(clippy::new_without_default)]
@@ -127,6 +132,7 @@ impl<'a> CodeBuilder<'a> {
             inner_length: Vec::with_capacity_in(5, arena),
             vm_block_stack,
             import_relocations: Vec::with_capacity_in(0, arena),
+            set_locals: BitVec::with_capacity(64),
         }
     }
 
@@ -137,6 +143,7 @@ impl<'a> CodeBuilder<'a> {
         self.preamble.clear();
         self.inner_length.clear();
         self.import_relocations.clear();
+        self.set_locals.clear();
 
         self.vm_block_stack.truncate(1);
         self.vm_block_stack[0].value_stack.clear();
@@ -536,6 +543,18 @@ impl<'a> CodeBuilder<'a> {
     }
     pub fn set_local(&mut self, id: LocalId) {
         self.inst_imm32(SETLOCAL, 1, false, id.0);
+        let index = id.0 as usize;
+        let len = self.set_locals.len();
+        if index >= len {
+            self.set_locals.extend(repeat(false).take(index + 1 - len));
+        }
+        self.set_locals.set(index, true);
+    }
+    /// Check if a local variable has been set
+    /// This is not a Wasm instruction, just a helper method
+    pub fn is_set(&self, id: LocalId) -> bool {
+        let index = id.0 as usize;
+        (index < self.set_locals.len()) && self.set_locals[index]
     }
     pub fn tee_local(&mut self, id: LocalId) {
         self.inst_imm32(TEELOCAL, 0, false, id.0);