Merge pull request #3379 from rtfeldman/wasm-box

Wasm box & unbox

crates/compiler/gen_wasm/src/backend.rs

@@ -17,7 +17,7 @@ use roc_std::RocDec;
 
 use crate::layout::{CallConv, ReturnMethod, WasmLayout};
 use crate::low_level::{call_higher_order_lowlevel, LowLevelCall};
-use crate::storage::{Storage, StoredValue, StoredVarKind};
+use crate::storage::{AddressValue, Storage, StoredValue, StoredVarKind};
 use crate::wasm_module::linking::{DataSymbol, WasmObjectSymbol};
 use crate::wasm_module::sections::{
     ConstExpr, DataMode, DataSegment, Export, Global, GlobalType, Import, ImportDesc, Limits,
@@ -928,7 +928,7 @@ impl<'a> WasmBackend<'a> {
                 index,
                 field_layouts,
                 structure,
-            } => self.expr_struct_at_index(sym, storage, *index, field_layouts, *structure),
+            } => self.expr_struct_at_index(sym, *index, field_layouts, *structure),
 
             Expr::Array { elems, elem_layout } => self.expr_array(sym, storage, elem_layout, elems),
 
@@ -953,9 +953,9 @@ impl<'a> WasmBackend<'a> {
                 index,
             } => self.expr_union_at_index(*structure, *tag_id, union_layout, *index, sym),
 
-            Expr::ExprBox { .. } | Expr::ExprUnbox { .. } => {
-                todo!("Expression `{}`", expr.to_pretty(100))
-            }
+            Expr::ExprBox { symbol: arg_sym } => self.expr_box(sym, *arg_sym, layout, storage),
+
+            Expr::ExprUnbox { symbol: arg_sym } => self.expr_unbox(sym, *arg_sym),
 
             Expr::Reuse {
                 tag_layout,
@@ -1353,16 +1353,15 @@ impl<'a> WasmBackend<'a> {
     fn expr_struct_at_index(
         &mut self,
         sym: Symbol,
-        storage: &StoredValue,
         index: u64,
         field_layouts: &'a [Layout<'a>],
         structure: Symbol,
     ) {
-        self.storage
-            .ensure_value_has_local(&mut self.code_builder, sym, storage.to_owned());
-        let (local_id, mut offset) = match self.storage.get(&structure) {
+        let (from_addr_val, mut offset) = match self.storage.get(&structure) {
             StoredValue::StackMemory { location, .. } => {
-                location.local_and_offset(self.storage.stack_frame_pointer)
+                let (local_id, offset) =
+                    location.local_and_offset(self.storage.stack_frame_pointer);
+                (AddressValue::NotLoaded(local_id), offset)
             }
 
             StoredValue::Local {
@@ -1371,18 +1370,20 @@ impl<'a> WasmBackend<'a> {
                 ..
             } => {
                 debug_assert!(matches!(value_type, ValueType::I32));
-                (*local_id, 0)
+                (AddressValue::NotLoaded(*local_id), 0)
             }
 
             StoredValue::VirtualMachineStack { .. } => {
-                internal_error!("ensure_value_has_local didn't work")
+                self.storage
+                    .load_symbols(&mut self.code_builder, &[structure]);
+                (AddressValue::Loaded, 0)
             }
         };
         for field in field_layouts.iter().take(index as usize) {
             offset += field.stack_size(TARGET_INFO);
         }
         self.storage
-            .copy_value_from_memory(&mut self.code_builder, sym, local_id, offset);
+            .copy_value_from_memory(&mut self.code_builder, sym, from_addr_val, offset);
     }
 
     /*******************************************************************
@@ -1700,20 +1701,83 @@ impl<'a> WasmBackend<'a> {
 
         let stores_tag_id_in_pointer = union_layout.stores_tag_id_in_pointer(TARGET_INFO);
 
-        let from_ptr = if stores_tag_id_in_pointer {
-            let ptr = self.storage.create_anonymous_local(ValueType::I32);
+        let from_addr_val = if stores_tag_id_in_pointer {
             self.code_builder.get_local(tag_local_id);
             self.code_builder.i32_const(-4); // 11111111...1100
             self.code_builder.i32_and();
-            self.code_builder.set_local(ptr);
-            ptr
+            AddressValue::Loaded
         } else {
-            tag_local_id
+            AddressValue::NotLoaded(tag_local_id)
         };
 
         let from_offset = tag_offset + field_offset;
+        self.storage.copy_value_from_memory(
+            &mut self.code_builder,
+            symbol,
+            from_addr_val,
+            from_offset,
+        );
+    }
+
+    /*******************************************************************
+     * Box
+     *******************************************************************/
+
+    fn expr_box(
+        &mut self,
+        ret_sym: Symbol,
+        arg_sym: Symbol,
+        layout: &Layout<'a>,
+        storage: &StoredValue,
+    ) {
+        // create a local variable for the heap pointer
+        let ptr_local_id = match self.storage.ensure_value_has_local(
+            &mut self.code_builder,
+            ret_sym,
+            storage.clone(),
+        ) {
+            StoredValue::Local { local_id, .. } => local_id,
+            _ => internal_error!("A heap pointer will always be an i32"),
+        };
+
+        // allocate heap memory and load its data address onto the value stack
+        let arg_layout = match layout {
+            Layout::Boxed(arg) => *arg,
+            _ => internal_error!("ExprBox should always produce a Boxed layout"),
+        };
+        let (size, alignment) = arg_layout.stack_size_and_alignment(TARGET_INFO);
+        self.allocate_with_refcount(Some(size), alignment, 1);
+
+        // store the pointer value from the value stack into the local variable
+        self.code_builder.set_local(ptr_local_id);
+
+        // copy the argument to the pointer address
         self.storage
-            .copy_value_from_memory(&mut self.code_builder, symbol, from_ptr, from_offset);
+            .copy_value_to_memory(&mut self.code_builder, ptr_local_id, 0, arg_sym);
+    }
+
+    fn expr_unbox(&mut self, ret_sym: Symbol, arg_sym: Symbol) {
+        let (from_addr_val, from_offset) = match self.storage.get(&arg_sym) {
+            StoredValue::VirtualMachineStack { .. } => {
+                self.storage
+                    .load_symbols(&mut self.code_builder, &[arg_sym]);
+                (AddressValue::Loaded, 0)
+            }
+            StoredValue::Local { local_id, .. } => (AddressValue::NotLoaded(*local_id), 0),
+            StoredValue::StackMemory { location, .. } => {
+                let (local_id, offset) =
+                    location.local_and_offset(self.storage.stack_frame_pointer);
+                (AddressValue::NotLoaded(local_id), offset)
+            }
+        };
+
+        // Copy the value
+        self.storage.copy_value_from_memory(
+            &mut self.code_builder,
+            ret_sym,
+            from_addr_val,
+            from_offset,
+        );
     }
 
     /*******************************************************************

crates/compiler/gen_wasm/src/low_level.rs

@@ -11,7 +11,7 @@ use roc_mono::low_level::HigherOrder;
 
 use crate::backend::{ProcLookupData, ProcSource, WasmBackend};
 use crate::layout::{CallConv, StackMemoryFormat, WasmLayout};
-use crate::storage::{StackMemoryLocation, StoredValue};
+use crate::storage::{AddressValue, StackMemoryLocation, StoredValue};
 use crate::wasm_module::{Align, LocalId, ValueType};
 use crate::TARGET_INFO;
 
@@ -316,14 +316,12 @@ impl<'a> LowLevelCall<'a> {
 
                 // Target element heap pointer
                 backend.code_builder.i32_add(); // base + index*size
-                let elem_heap_ptr = backend.storage.create_anonymous_local(ValueType::I32);
-                backend.code_builder.set_local(elem_heap_ptr);
 
                 // Copy to stack
                 backend.storage.copy_value_from_memory(
                     &mut backend.code_builder,
                     self.ret_symbol,
-                    elem_heap_ptr,
+                    AddressValue::Loaded,
                     0,
                 );
 

crates/compiler/gen_wasm/src/storage.rs

@@ -76,6 +76,13 @@ impl StoredValue {
     }
 }
 
+pub enum AddressValue {
+    /// The address value has been loaded to the VM stack
+    Loaded,
+    /// The address value is in a local variable
+    NotLoaded(LocalId),
+}
+
 /// Helper structure for WasmBackend, to keep track of how values are stored,
 /// including the VM stack, local variables, and linear memory
 #[derive(Debug)]
@@ -610,7 +617,7 @@ impl<'a> Storage<'a> {
         &mut self,
         code_builder: &mut CodeBuilder,
         to_symbol: Symbol,
-        from_ptr: LocalId,
+        from_addr: AddressValue,
         from_offset: u32,
     ) -> u32 {
         let to_storage = self.get(&to_symbol).to_owned();
@@ -625,6 +632,16 @@ impl<'a> Storage<'a> {
                     self.stack_frame_pointer = Some(self.get_next_local_id());
                 }
 
+                let from_ptr = match from_addr {
+                    AddressValue::NotLoaded(ptr) => ptr,
+                    AddressValue::Loaded => {
+                        // The `from` address is on the VM stack but we want it in a local for copying
+                        let tmp_local = self.create_anonymous_local(PTR_TYPE);
+                        code_builder.set_local(tmp_local);
+                        tmp_local
+                    }
+                };
+
                 let (to_ptr, to_offset) = location.local_and_offset(self.stack_frame_pointer);
                 copy_memory(
                     code_builder,
@@ -648,7 +665,10 @@ impl<'a> Storage<'a> {
             } => {
                 use crate::wasm_module::Align::*;
 
-                code_builder.get_local(from_ptr);
+                if let AddressValue::NotLoaded(from_ptr) = from_addr {
+                    code_builder.get_local(from_ptr);
+                }
+
                 match (value_type, size) {
                     (ValueType::I64, 8) => code_builder.i64_load(Bytes8, from_offset),
                     (ValueType::I32, 4) => code_builder.i32_load(Bytes4, from_offset),

crates/compiler/test_gen/src/gen_primitives.rs

@@ -3260,7 +3260,7 @@ fn issue_2322() {
 }
 
 #[test]
-#[cfg(any(feature = "gen-llvm"))]
+#[cfg(any(feature = "gen-llvm", feature = "gen-wasm"))]
 fn box_and_unbox_string() {
     assert_evals_to!(
         indoc!(
@@ -3278,7 +3278,7 @@ fn box_and_unbox_string() {
 }
 
 #[test]
-#[cfg(any(feature = "gen-llvm"))]
+#[cfg(any(feature = "gen-llvm", feature = "gen-wasm"))]
 fn box_and_unbox_num() {
     assert_evals_to!(
         indoc!(
@@ -3292,7 +3292,7 @@ fn box_and_unbox_num() {
 }
 
 #[test]
-#[cfg(any(feature = "gen-llvm"))]
+#[cfg(any(feature = "gen-llvm", feature = "gen-wasm"))]
 fn box_and_unbox_record() {
     assert_evals_to!(
         indoc!(
@@ -3306,7 +3306,7 @@ fn box_and_unbox_record() {
 }
 
 #[test]
-#[cfg(any(feature = "gen-llvm"))]
+#[cfg(any(feature = "gen-llvm", feature = "gen-wasm"))]
 fn box_and_unbox_tag_union() {
     assert_evals_to!(
         indoc!(