Setup for specializing equality checks

compiler/gen_wasm/src/backend.rs

@@ -6,7 +6,7 @@ use roc_collections::all::MutMap;
 use roc_module::low_level::LowLevel;
 use roc_module::symbol::{Interns, Symbol};
 use roc_mono::code_gen_help::{CodeGenHelp, REFCOUNT_MAX};
-use roc_mono::ir::{CallType, Expr, JoinPointId, Literal, Proc, Stmt};
+use roc_mono::ir::{CallType, Expr, JoinPointId, Literal, Proc, ProcLayout, Stmt};
 use roc_mono::layout::{Builtin, Layout, LayoutIds, TagIdIntType, UnionLayout};
 use roc_reporting::internal_error;
 
@@ -168,6 +168,25 @@ impl<'a> WasmBackend<'a> {
             .generate_procs(self.env.arena, ident_ids)
     }
 
+    fn register_helper_proc(&mut self, new_proc_info: (Symbol, ProcLayout<'a>)) {
+        let (new_proc_sym, new_proc_layout) = new_proc_info;
+        let wasm_fn_index = self.proc_symbols.len() as u32;
+        let linker_sym_index = self.linker_symbols.len() as u32;
+
+        let name = self
+            .layout_ids
+            .get_toplevel(new_proc_sym, &new_proc_layout)
+            .to_symbol_string(new_proc_sym, self.interns);
+
+        self.proc_symbols.push((new_proc_sym, linker_sym_index));
+        self.linker_symbols
+            .push(SymInfo::Function(WasmObjectSymbol::Defined {
+                flags: 0,
+                index: wasm_fn_index,
+                name,
+            }));
+    }
+
     pub fn finalize_module(mut self) -> WasmModule<'a> {
         let symbol_table = LinkingSubSection::SymbolTable(self.linker_symbols);
         self.module.linking.subsections.push(symbol_table);
@@ -500,7 +519,7 @@ impl<'a> WasmBackend<'a> {
                     .get_mut(&self.env.module_id)
                     .unwrap();
 
-                let (rc_stmt, new_proc_info) = self
+                let (rc_stmt, maybe_new_proc_info) = self
                     .helper_proc_gen
                     .expand_refcount_stmt(ident_ids, *layout, modify, *following);
 
@@ -509,25 +528,8 @@ impl<'a> WasmBackend<'a> {
                     println!("## rc_stmt:\n{}\n{:?}", rc_stmt.to_pretty(200), rc_stmt);
                 }
 
-                // If we're creating a new RC procedure, we need to store its symbol data,
+                // If this is the first call to a new helper proc, register its symbol data
-                // so that we can correctly generate calls to it.
+                maybe_new_proc_info.map(|info| self.register_helper_proc(info));
-                if let Some((rc_proc_sym, rc_proc_layout)) = new_proc_info {
-                    let wasm_fn_index = self.proc_symbols.len() as u32;
-                    let linker_sym_index = self.linker_symbols.len() as u32;
-
-                    let name = self
-                        .layout_ids
-                        .get_toplevel(rc_proc_sym, &rc_proc_layout)
-                        .to_symbol_string(rc_proc_sym, self.interns);
-
-                    self.proc_symbols.push((rc_proc_sym, linker_sym_index));
-                    self.linker_symbols
-                        .push(SymInfo::Function(WasmObjectSymbol::Defined {
-                            flags: 0,
-                            index: wasm_fn_index,
-                            name,
-                        }));
-                }
 
                 self.build_stmt(rc_stmt, ret_layout);
             }
@@ -560,7 +562,13 @@ impl<'a> WasmBackend<'a> {
                 CallType::ByName { name: func_sym, .. } => {
                     // If this function is just a lowlevel wrapper, then inline it
                     if let Some(lowlevel) = LowLevel::from_inlined_wrapper(*func_sym) {
-                        return self.build_low_level(lowlevel, arguments, *sym, wasm_layout);
+                        return self.build_low_level(
+                            lowlevel,
+                            arguments,
+                            *sym,
+                            wasm_layout,
+                            storage,
+                        );
                     }
 
                     let (param_types, ret_type) = self.storage.load_symbols_for_call(
@@ -596,7 +604,7 @@ impl<'a> WasmBackend<'a> {
                 }
 
                 CallType::LowLevel { op: lowlevel, .. } => {
-                    self.build_low_level(*lowlevel, arguments, *sym, wasm_layout)
+                    self.build_low_level(*lowlevel, arguments, *sym, wasm_layout, storage)
                 }
 
                 x => todo!("call type {:?}", x),
@@ -925,6 +933,7 @@ impl<'a> WasmBackend<'a> {
         arguments: &'a [Symbol],
         return_sym: Symbol,
         return_layout: WasmLayout,
+        storage: &StoredValue,
     ) {
         let (param_types, ret_type) = self.storage.load_symbols_for_call(
             self.env.arena,
@@ -949,6 +958,30 @@ impl<'a> WasmBackend<'a> {
             BuiltinCall(name) => {
                 self.call_zig_builtin(name, param_types, ret_type);
             }
+            GeneratedHelper => {
+                let layout = self.symbol_layouts[&arguments[0]];
+
+                let ident_ids = self
+                    .interns
+                    .all_ident_ids
+                    .get_mut(&self.env.module_id)
+                    .unwrap();
+
+                let (replacement_expr, maybe_new_proc_info) = self
+                    .helper_proc_gen
+                    .replace_generic_equals(ident_ids, layout);
+
+                // If this is the first call to a new helper proc, register its symbol data
+                maybe_new_proc_info.map(|info| self.register_helper_proc(info));
+
+                self.build_expr(&return_sym, replacement_expr, &layout, storage);
+
+                if lowlevel == LowLevel::NotEq {
+                    self.code_builder.i32_eqz();
+                } else {
+                    debug_assert!(lowlevel == LowLevel::Eq);
+                }
+            }
             NotImplemented => {
                 todo!("Low level operation {:?}", lowlevel)
             }

compiler/gen_wasm/src/low_level.rs

@@ -10,6 +10,7 @@ use crate::wasm_module::{Align, CodeBuilder, ValueType::*};
 pub enum LowlevelBuildResult {
     Done,
     BuiltinCall(&'static str),
+    GeneratedHelper,
     NotImplemented,
 }
 
@@ -563,7 +564,8 @@ pub fn decode_low_level<'a>(
                                 code_builder.i32_and();
                             }
                             Int128 => compare_bytes(code_builder),
-                            Float128 | DataStructure => return NotImplemented,
+                            Float128 => return NotImplemented,
+                            DataStructure => return GeneratedHelper,
                         }
                     }
                 }
@@ -577,9 +579,13 @@ pub fn decode_low_level<'a>(
                 F32 => code_builder.f32_ne(),
                 F64 => code_builder.f64_ne(),
             },
-            StoredValue::StackMemory { .. } => {
+            StoredValue::StackMemory { format, .. } => {
-                decode_low_level(code_builder, storage, LowLevel::Eq, args, ret_layout);
+                if matches!(format, DataStructure) {
-                code_builder.i32_eqz();
+                    return GeneratedHelper;
+                } else {
+                    decode_low_level(code_builder, storage, LowLevel::Eq, args, ret_layout);
+                    code_builder.i32_eqz();
+                }
             }
         },
         And => code_builder.i32_and(),

compiler/mono/src/code_gen_help.rs

@@ -89,7 +89,7 @@ impl<'a> CodeGenHelp<'a> {
         }
     }
 
-    /// Expands a `Refcounting` node to a `Let` node that calls a specialized helper.
+    /// Expand a `Refcounting` node to a `Let` node that calls a specialized helper proc.
     /// The helper procs themselves are to be generated later with `generate_procs`
     pub fn expand_refcount_stmt(
         &mut self,
@@ -224,6 +224,16 @@ impl<'a> CodeGenHelp<'a> {
         }
     }
 
+    /// Replace a generic `Lowlevel::Eq` call with a specialized helper proc.
+    /// The helper procs themselves are to be generated later with `generate_procs`
+    pub fn replace_generic_equals(
+        &mut self,
+        _ident_ids: &mut IdentIds,
+        _layout: Layout<'a>,
+    ) -> (&'a Expr<'a>, Option<(Symbol, ProcLayout<'a>)>) {
+        todo!()
+    }
+
     // Check if the specialization is implemented yet. In the long term, this will be deleted.
     // In the short term, we have to specify in two places what's complete and what's not:
     // Here and in generate_procs. We use assertions to ensure they match.