Move relocations code into the backend and code builder

compiler/gen_wasm/src/backend.rs

@@ -10,6 +10,7 @@ use roc_mono::layout::{Builtin, Layout};
 
 use crate::code_builder::{BlockType, CodeBuilder, ValueType};
 use crate::layout::WasmLayout;
+use crate::module_builder::RelocationEntry;
 use crate::serialize::SerialBuffer;
 use crate::storage::{Storage, StoredValue, StoredValueKind};
 use crate::{copy_memory, CopyMemoryConfig, Env, LocalId, PTR_TYPE};
@@ -27,7 +28,7 @@ pub struct WasmBackend<'a> {
     // Module-level data
     pub module_builder: ModuleBuilder,
     pub code_section_bytes: std::vec::Vec<u8>,
-    pub call_relocations: Vec<'a, (usize, Symbol)>,
+    pub code_relocations: Vec<'a, RelocationEntry>,
     _data_offset_map: MutMap<Literal<'a>, u32>,
     _data_offset_next: u32,
     proc_symbols: &'a Vec<'a, Symbol>,
@@ -58,7 +59,7 @@ impl<'a> WasmBackend<'a> {
             _data_offset_map: MutMap::default(),
             _data_offset_next: UNUSED_DATA_SECTION_BYTES,
             proc_symbols,
-            call_relocations: Vec::with_capacity_in(256, env.arena),
+            code_relocations: Vec::with_capacity_in(256, env.arena),
 
             // Function-level data
             block_depth: 0,
@@ -141,7 +142,7 @@ impl<'a> WasmBackend<'a> {
         );
 
         let relocs = self.code_builder.serialize(&mut self.code_section_bytes);
-        self.call_relocations.extend(relocs);
+        self.code_relocations.extend(relocs);
         Ok(())
     }
 
@@ -394,6 +395,8 @@ impl<'a> WasmBackend<'a> {
 
                     self.storage.load_symbols(&mut self.code_builder, wasm_args);
 
+                    // Index of the called function in the Wasm module code section
+                    // TODO: update when we start inlining functions (remember we emit procs out of order)
                     let func_index = self
                         .proc_symbols
                         .iter()
@@ -401,8 +404,15 @@ impl<'a> WasmBackend<'a> {
                         .map(|i| i as u32)
                         .unwrap_or(u32::MAX); // foreign symbol, updated at link time
 
-                    self.code_builder
-                        .call(func_index, *func_sym, wasm_args.len(), has_return_val);
+                    // Index of the function's name in the symbol table
+                    let symbol_index = func_index; // TODO: update this when we add other things to the symbol table
+
+                    self.code_builder.call(
+                        func_index,
+                        symbol_index,
+                        wasm_args.len(),
+                        has_return_val,
+                    );
 
                     Ok(())
                 }

compiler/gen_wasm/src/code_builder.rs

@@ -5,6 +5,7 @@ use std::fmt::Debug;
 
 use roc_module::symbol::Symbol;
 
+use crate::module_builder::{IndexRelocType, RelocationEntry};
 use crate::opcodes::*;
 use crate::serialize::SerialBuffer;
 use crate::{round_up_to_alignment, LocalId, FRAME_ALIGNMENT_BYTES, STACK_POINTER_GLOBAL_ID};
@@ -141,8 +142,8 @@ pub struct CodeBuilder<'a> {
     vm_stack: Vec<'a, Symbol>,
 
     /// Which byte offsets in the code section correspond to which symbols.
-    /// e.g. Function indices that may change when we link Roc + Zig modules
-    relocations: Vec<'a, (usize, Symbol)>,
+    /// e.g. Function indices may change when we link Roc + builtins + platform
+    relocations: Vec<'a, RelocationEntry>,
 }
 
 #[allow(clippy::new_without_default)]
@@ -389,7 +390,7 @@ impl<'a> CodeBuilder<'a> {
     pub fn serialize<T: SerialBuffer>(
         &mut self,
         code_section_buf: &mut T,
-    ) -> Drain<(usize, Symbol)> {
+    ) -> Drain<RelocationEntry> {
         let function_offset = code_section_buf.size();
         code_section_buf.append_slice(&self.inner_length);
         code_section_buf.append_slice(&self.preamble);
@@ -409,8 +410,11 @@ impl<'a> CodeBuilder<'a> {
         code_section_buf.append_slice(&self.code[pos..len]);
 
         let extra_offset = function_offset + self.inner_length.len() + self.preamble.len();
-        for (offset, _) in self.relocations.iter_mut() {
-            *offset += extra_offset;
+        for reloc in self.relocations.iter_mut() {
+            match reloc {
+                RelocationEntry::Index { offset, .. } => *offset += extra_offset as u32,
+                RelocationEntry::Offset { offset, .. } => *offset += extra_offset as u32,
+            }
         }
         self.relocations.drain(0..)
     }
@@ -490,7 +494,7 @@ impl<'a> CodeBuilder<'a> {
     pub fn call(
         &mut self,
         function_index: u32,
-        function_sym: Symbol,
+        symbol_index: u32,
         n_args: usize,
         has_return_val: bool,
     ) {
@@ -507,9 +511,13 @@ impl<'a> CodeBuilder<'a> {
         }
         self.code.push(CALL);
 
-        let call_offset = self.code.len() + self.insert_bytes.len();
+        let fn_index_offset = self.code.len() + self.insert_bytes.len();
         self.code.encode_padded_u32(function_index);
-        self.relocations.push((call_offset, function_sym));
+        self.relocations.push(RelocationEntry::Index {
+            type_id: IndexRelocType::FunctionIndexLeb,
+            offset: fn_index_offset as u32,
+            symbol_index,
+        });
     }
 
     #[allow(dead_code)]

compiler/gen_wasm/src/lib.rs

@@ -21,7 +21,7 @@ use roc_mono::layout::LayoutIds;
 use crate::backend::WasmBackend;
 use crate::code_builder::{Align, CodeBuilder, ValueType};
 use crate::module_builder::{
-    LinkingSection, LinkingSubSection, RelocationEntry, RelocationSection, SectionId, SymInfo,
+    LinkingSection, LinkingSubSection, RelocationSection, SectionId, SymInfo,
 };
 use crate::serialize::{SerialBuffer, Serialize};
 
@@ -108,17 +108,9 @@ pub fn build_module_help<'a>(
     });
 
     // Code relocations
-    let call_reloc_iter = backend
-        .call_relocations
-        .iter()
-        .map(|(offset, sym)| -> RelocationEntry {
-            let symbol_index = proc_symbols.iter().position(|s| s == sym).unwrap();
-            RelocationEntry::for_function_call(*offset as u32, symbol_index as u32)
-        });
-
     let code_reloc_section = RelocationSection {
         name: "reloc.CODE".to_string(),
-        entries: Vec::from_iter_in(call_reloc_iter, env.arena),
+        entries: &backend.code_relocations,
     };
 
     let mut code_reloc_section_bytes = std::vec::Vec::with_capacity(256);

compiler/gen_wasm/src/module_builder.rs

@@ -132,6 +132,7 @@ pub enum OffsetRelocType {
     MemoryAddrI64 = 16,
 }
 
+#[derive(Debug)]
 pub enum RelocationEntry {
     Index {
         type_id: IndexRelocType,
@@ -185,7 +186,7 @@ impl Serialize for RelocationEntry {
 
 pub struct RelocationSection<'a> {
     pub name: String,
-    pub entries: Vec<'a, RelocationEntry>,
+    pub entries: &'a Vec<'a, RelocationEntry>,
 }
 
 impl<'a> Serialize for RelocationSection<'a> {