Refactor Wasm equality operator

compiler/gen_wasm/src/backend.rs

@@ -16,7 +16,7 @@ use roc_reporting::internal_error;
 
 use crate::layout::{CallConv, ReturnMethod, StackMemoryFormat, WasmLayout};
 use crate::low_level::{dispatch_low_level, LowlevelBuildResult};
-use crate::storage::{Storage, StoredValue, StoredValueKind};
+use crate::storage::{StackMemoryLocation, Storage, StoredValue, StoredValueKind};
 use crate::wasm_module::linking::{
     DataSymbol, LinkingSection, RelocationSection, WasmObjectSymbol, WASM_SYM_BINDING_WEAK,
     WASM_SYM_UNDEFINED,
@@ -1081,24 +1081,17 @@ impl<'a> WasmBackend<'a> {
         return_layout: WasmLayout,
         storage: &StoredValue,
     ) {
-        // Load the arguments using Zig calling convention
-        let (param_types, ret_type) = self.storage.load_symbols_for_call(
-            self.env.arena,
-            &mut self.code_builder,
-            arguments,
-            return_sym,
-            &return_layout,
-            CallConv::Zig,
-        );
-
         use StoredValue::*;
         match self.storage.get(&arguments[0]).to_owned() {
-            VirtualMachineStack { value_type, .. } | Local { value_type, .. } => match value_type {
+            VirtualMachineStack { value_type, .. } | Local { value_type, .. } => {
+                self.storage.load_symbols(&mut self.code_builder, arguments);
+                match value_type {
                     ValueType::I32 => self.code_builder.i32_eq(),
                     ValueType::I64 => self.code_builder.i64_eq(),
                     ValueType::F32 => self.code_builder.f32_eq(),
                     ValueType::F64 => self.code_builder.f64_eq(),
-            },
+                }
+            }
             StackMemory {
                 format,
                 location: location0,
@@ -1109,26 +1102,29 @@ impl<'a> WasmBackend<'a> {
                     ..
                 } = self.storage.get(&arguments[1]).to_owned()
                 {
-                    let stack_frame_pointer = self.storage.stack_frame_pointer;
+                    self.build_eq_memory(
-                    let compare_bytes = |code_builder: &mut CodeBuilder| {
+                        format,
-                        let (local0, offset0) = location0.local_and_offset(stack_frame_pointer);
+                        [location0, location1],
-                        let (local1, offset1) = location1.local_and_offset(stack_frame_pointer);
+                        arguments,
-
+                        return_sym,
-                        code_builder.get_local(local0);
+                        return_layout,
-                        code_builder.i64_load(Align::Bytes8, offset0);
+                        storage,
-                        code_builder.get_local(local1);
+                    )
-                        code_builder.i64_load(Align::Bytes8, offset1);
+                }
-                        code_builder.i64_eq();
+            }
-
-                        code_builder.get_local(local0);
-                        code_builder.i64_load(Align::Bytes8, offset0 + 8);
-                        code_builder.get_local(local1);
-                        code_builder.i64_load(Align::Bytes8, offset1 + 8);
-                        code_builder.i64_eq();
-
-                        code_builder.i32_and();
         };
+    }
 
+    // Equality for values in memory (as opposed to VM stack)
+    fn build_eq_memory(
+        &mut self,
+        format: StackMemoryFormat,
+        locations: [StackMemoryLocation; 2],
+        arguments: &'a [Symbol],
+        return_sym: Symbol,
+        return_layout: WasmLayout,
+        storage: &StoredValue,
+    ) {
         match format {
             StackMemoryFormat::Decimal => {
                 // Both args are finite
@@ -1151,11 +1147,14 @@ impl<'a> WasmBackend<'a> {
                 self.code_builder.i32_and();
 
                 // AND they have the same bytes
-                            compare_bytes(&mut self.code_builder);
+                self.build_eq_help_128bit(locations);
                 self.code_builder.i32_and();
             }
-                        StackMemoryFormat::Int128 => compare_bytes(&mut self.code_builder),
+
+            StackMemoryFormat::Int128 => self.build_eq_help_128bit(locations),
+
             StackMemoryFormat::Float128 => todo!("equality for f128"),
+
             StackMemoryFormat::DataStructure => {
                 let layout = self.symbol_layouts[&arguments[0]];
                 let layout_rhs = self.symbol_layouts[&arguments[1]];
@@ -1165,6 +1164,14 @@ impl<'a> WasmBackend<'a> {
                 );
 
                 if layout == Layout::Builtin(Builtin::Str) {
+                    let (param_types, ret_type) = self.storage.load_symbols_for_call(
+                        self.env.arena,
+                        &mut self.code_builder,
+                        arguments,
+                        return_sym,
+                        &return_layout,
+                        CallConv::Zig,
+                    );
                     self.call_zig_builtin(bitcode::STR_EQUAL, param_types, ret_type);
                 } else if layout.stack_size(PTR_SIZE) == 0 {
                     // Always true: `Unit == Unit`, or `{} == {}`
@@ -1186,18 +1193,30 @@ impl<'a> WasmBackend<'a> {
                     }
 
                     let bool_layout = Layout::Builtin(Builtin::Bool);
-                                self.build_expr(
+                    self.build_expr(&return_sym, replacement_expr, &bool_layout, storage);
-                                    &return_sym,
-                                    replacement_expr,
-                                    &bool_layout,
-                                    storage,
-                                );
-                            }
-                        }
                 }
             }
         }
     }
+
+    /// Equality helper for 128-bit numbers
+    fn build_eq_help_128bit(&mut self, locations: [StackMemoryLocation; 2]) {
+        let (local0, offset0) = locations[0].local_and_offset(self.storage.stack_frame_pointer);
+        let (local1, offset1) = locations[1].local_and_offset(self.storage.stack_frame_pointer);
+
+        self.code_builder.get_local(local0);
+        self.code_builder.i64_load(Align::Bytes8, offset0);
+        self.code_builder.get_local(local1);
+        self.code_builder.i64_load(Align::Bytes8, offset1);
+        self.code_builder.i64_eq();
+
+        self.code_builder.get_local(local0);
+        self.code_builder.i64_load(Align::Bytes8, offset0 + 8);
+        self.code_builder.get_local(local1);
+        self.code_builder.i64_load(Align::Bytes8, offset1 + 8);
+        self.code_builder.i64_eq();
+
+        self.code_builder.i32_and();
     }
 
     fn load_literal(