Wasm: refactor the model of the Linking section

compiler/gen_wasm/src/backend.rs

@@ -26,8 +26,8 @@ use crate::wasm_module::sections::{
     Import, ImportDesc, ImportSection, MemorySection, TypeSection, WasmModule,
 };
 use crate::wasm_module::{
-    code_builder, CodeBuilder, ConstExpr, Export, ExportType, Global, GlobalType,
+    code_builder, CodeBuilder, ConstExpr, Export, ExportType, Global, GlobalType, LocalId,
-    LinkingSubSection, LocalId, Signature, SymInfo, ValueType,
+    Signature, SymInfo, ValueType,
 };
 use crate::{
     copy_memory, round_up_to_alignment, CopyMemoryConfig, Env, BUILTINS_IMPORT_MODULE_NAME,
@@ -50,7 +50,6 @@ pub struct WasmBackend<'a> {
     next_constant_addr: u32,
     builtin_sym_index_map: MutMap<&'a str, usize>,
     proc_symbols: Vec<'a, (Symbol, u32)>,
-    linker_symbols: Vec<'a, SymInfo>,
     helper_proc_gen: CodeGenHelp<'a>,
 
     // Function-level data
@@ -103,6 +102,8 @@ impl<'a> WasmBackend<'a> {
             index: STACK_POINTER_GLOBAL_ID,
             name: STACK_POINTER_NAME.to_string(),
         }));
+        let mut linking = LinkingSection::new(arena);
+        linking.symbol_table = linker_symbols;
 
         let module = WasmModule {
             types: TypeSection::new(arena, num_procs),
@@ -122,7 +123,7 @@ impl<'a> WasmBackend<'a> {
                 code_builders: Vec::with_capacity_in(num_procs, arena),
             },
             data: DataSection::new(arena),
-            linking: LinkingSection::new(arena),
+            linking,
             relocations: RelocationSection::new(arena, "reloc.CODE"),
         };
 
@@ -137,7 +138,6 @@ impl<'a> WasmBackend<'a> {
             next_constant_addr: CONST_SEGMENT_BASE_ADDR,
             builtin_sym_index_map: MutMap::default(),
             proc_symbols,
-            linker_symbols,
             helper_proc_gen,
 
             // Function-level data
@@ -157,7 +157,7 @@ impl<'a> WasmBackend<'a> {
     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 linker_sym_index = self.module.linking.symbol_table.len() as u32;
 
         let name = self
             .layout_ids
@@ -165,17 +165,15 @@ impl<'a> WasmBackend<'a> {
             .to_symbol_string(new_proc_sym, self.interns);
 
         self.proc_symbols.push((new_proc_sym, linker_sym_index));
-        self.linker_symbols
+        let linker_symbol = SymInfo::Function(WasmObjectSymbol::Defined {
-            .push(SymInfo::Function(WasmObjectSymbol::Defined {
+            flags: 0,
-                flags: 0,
+            index: wasm_fn_index,
-                index: wasm_fn_index,
+            name,
-                name,
+        });
-            }));
+        self.module.linking.symbol_table.push(linker_symbol);
     }
 
-    pub fn finalize_module(mut self) -> WasmModule<'a> {
+    pub fn into_module(self) -> WasmModule<'a> {
-        let symbol_table = LinkingSubSection::SymbolTable(self.linker_symbols);
-        self.module.linking.subsections.push(symbol_table);
         self.module
     }
 
@@ -1464,8 +1462,8 @@ impl<'a> WasmBackend<'a> {
             size: string.len() as u32,
         });
 
-        let linker_sym_index = self.linker_symbols.len();
+        let linker_sym_index = self.module.linking.symbol_table.len();
-        self.linker_symbols.push(linker_symbol);
+        self.module.linking.symbol_table.push(linker_symbol);
 
         (linker_sym_index as u32, elements_addr)
     }
@@ -1518,7 +1516,7 @@ impl<'a> WasmBackend<'a> {
         let has_return_val = ret_type.is_some();
 
         let (fn_index, linker_symbol_index) = match self.builtin_sym_index_map.get(name) {
-            Some(sym_idx) => match &self.linker_symbols[*sym_idx] {
+            Some(sym_idx) => match &self.module.linking.symbol_table[*sym_idx] {
                 SymInfo::Function(WasmObjectSymbol::Imported { index, .. }) => {
                     (*index, *sym_idx as u32)
                 }
@@ -1543,12 +1541,12 @@ impl<'a> WasmBackend<'a> {
                 self.module.import.entries.push(import);
 
                 // Provide symbol information for the linker
-                let sym_idx = self.linker_symbols.len();
+                let sym_idx = self.module.linking.symbol_table.len();
                 let sym_info = SymInfo::Function(WasmObjectSymbol::Imported {
                     flags: WASM_SYM_UNDEFINED,
                     index: import_index,
                 });
-                self.linker_symbols.push(sym_info);
+                self.module.linking.symbol_table.push(sym_info);
 
                 // Remember that we have created all of this data, and don't need to do it again
                 self.builtin_sym_index_map.insert(name, sym_idx);
@@ -1568,7 +1566,7 @@ impl<'a> WasmBackend<'a> {
     /// }
     fn _debug_current_proc_is(&self, linker_name: &'static str) -> bool {
         let (_, linker_sym_index) = self.proc_symbols[self.debug_current_proc_index];
-        let sym_info = &self.linker_symbols[linker_sym_index as usize];
+        let sym_info = &self.module.linking.symbol_table[linker_sym_index as usize];
         match sym_info {
             SymInfo::Function(WasmObjectSymbol::Defined { name, .. }) => name == linker_name,
             _ => false,

compiler/gen_wasm/src/lib.rs

@@ -131,7 +131,7 @@ pub fn build_module_help<'a>(
         backend.build_proc(proc);
     }
 
-    let module = backend.finalize_module();
+    let module = backend.into_module();
 
     Ok((module, main_fn_index.unwrap()))
 }

compiler/gen_wasm/src/wasm_module/linking.rs

@@ -1,6 +1,5 @@
 use bumpalo::collections::vec::Vec;
 use bumpalo::Bump;
-use roc_reporting::internal_error;
 
 use super::sections::{update_section_size, write_custom_section_header};
 use super::serialize::{SerialBuffer, Serialize};
@@ -183,8 +182,12 @@ pub struct LinkingSegment {
 }
 
 impl Serialize for LinkingSegment {
-    fn serialize<T: SerialBuffer>(&self, _buffer: &mut T) {
+    fn serialize<T: SerialBuffer>(&self, buffer: &mut T) {
-        todo!();
+        buffer.encode_u32(self.name.len() as u32);
+        buffer.append_slice(self.name.as_bytes());
+        let align_bytes_pow2 = self.alignment as u32;
+        buffer.encode_u32(align_bytes_pow2);
+        buffer.encode_u32(self.flags);
     }
 }
 
@@ -418,35 +421,25 @@ impl Serialize for SymInfo {
 //  Linking subsections
 //----------------------------------------------------------------
 
+#[repr(u8)]
 #[derive(Debug)]
-pub enum LinkingSubSection<'a> {
+enum SubSectionId {
-    /// Extra metadata about the data segments.
+    SegmentInfo = 5,
-    SegmentInfo(Vec<'a, LinkingSegment>),
+    InitFuncs = 6,
-    /// Specifies a list of constructor functions to be called at startup.
+    ComdatInfo = 7,
-    /// These constructors will be called in priority order after memory has been initialized.
+    SymbolTable = 8,
-    InitFuncs(Vec<'a, LinkingInitFunc>),
-    /// Specifies the COMDAT groups of associated linking objects, which are linked only once and all together.
-    ComdatInfo(Vec<'a, LinkingComdat<'a>>),
-    /// Specifies extra information about the symbols present in the module.
-    SymbolTable(Vec<'a, SymInfo>),
 }
 
-impl<'a> Serialize for LinkingSubSection<'a> {
+fn serialize_subsection<'a, I: Serialize, T: SerialBuffer>(
-    fn serialize<T: SerialBuffer>(&self, buffer: &mut T) {
+    buffer: &mut T,
-        buffer.append_u8(match self {
+    id: SubSectionId,
-            Self::SegmentInfo(_) => 5,
+    items: &[I],
-            Self::InitFuncs(_) => 6,
+) {
-            Self::ComdatInfo(_) => 7,
+    if !items.is_empty() {
-            Self::SymbolTable(_) => 8,
+        buffer.append_u8(id as u8);
-        });
         let payload_len_index = buffer.reserve_padded_u32();
         let payload_start_index = buffer.size();
-        match self {
+        items.serialize(buffer);
-            Self::SegmentInfo(items) => items.serialize(buffer),
-            Self::InitFuncs(items) => items.serialize(buffer),
-            Self::ComdatInfo(items) => items.serialize(buffer),
-            Self::SymbolTable(items) => items.serialize(buffer),
-        }
         buffer.overwrite_padded_u32(
             payload_len_index,
             (buffer.size() - payload_start_index) as u32,
@@ -460,35 +453,39 @@ impl<'a> Serialize for LinkingSubSection<'a> {
 
 const LINKING_VERSION: u8 = 2;
 
+/// The spec describes this in very weird way, so we're doing something saner.
+/// They call it an "array" of subsections with different variants, BUT this "array"
+/// has an implicit length, and none of the items can be repeated, so a struct is better.
+/// No point writing code to "find" the symbol table, when we know there's exactly one.
 #[derive(Debug)]
 pub struct LinkingSection<'a> {
-    pub subsections: Vec<'a, LinkingSubSection<'a>>,
+    pub symbol_table: Vec<'a, SymInfo>,
+    pub segment_info: Vec<'a, LinkingSegment>,
+    pub init_funcs: Vec<'a, LinkingInitFunc>,
+    pub comdat_info: Vec<'a, LinkingComdat<'a>>,
 }
 
 impl<'a> LinkingSection<'a> {
     pub fn new(arena: &'a Bump) -> Self {
         LinkingSection {
-            subsections: Vec::with_capacity_in(1, arena),
+            symbol_table: Vec::with_capacity_in(16, arena),
+            segment_info: Vec::with_capacity_in(16, arena),
+            init_funcs: Vec::with_capacity_in(0, arena),
+            comdat_info: Vec::with_capacity_in(0, arena),
         }
     }
-
-    pub fn symbol_table_mut(&mut self) -> &mut Vec<'a, SymInfo> {
-        for sub in self.subsections.iter_mut() {
-            if let LinkingSubSection::SymbolTable(syminfos) = sub {
-                return syminfos;
-            }
-        }
-        internal_error!("Symbol table not found");
-    }
 }
 
 impl<'a> Serialize for LinkingSection<'a> {
     fn serialize<T: SerialBuffer>(&self, buffer: &mut T) {
         let header_indices = write_custom_section_header(buffer, "linking");
         buffer.append_u8(LINKING_VERSION);
-        for subsection in self.subsections.iter() {
+
-            subsection.serialize(buffer);
+        serialize_subsection(buffer, SubSectionId::SymbolTable, &self.symbol_table);
-        }
+        serialize_subsection(buffer, SubSectionId::SegmentInfo, &self.segment_info);
+        serialize_subsection(buffer, SubSectionId::InitFuncs, &self.init_funcs);
+        serialize_subsection(buffer, SubSectionId::ComdatInfo, &self.comdat_info);
+
         update_section_size(buffer, header_indices);
     }
 }

compiler/gen_wasm/src/wasm_module/mod.rs

@@ -5,5 +5,5 @@ pub mod sections;
 pub mod serialize;
 
 pub use code_builder::{Align, CodeBuilder, LocalId, ValueType, VmSymbolState};
-pub use linking::{LinkingSubSection, SymInfo};
+pub use linking::SymInfo;
 pub use sections::{ConstExpr, Export, ExportType, Global, GlobalType, Signature, WasmModule};

compiler/gen_wasm/src/wasm_module/sections.rs

@@ -728,13 +728,11 @@ impl<'a> WasmModule<'a> {
         code_section_body_index: usize,
         n_imported_fns: u32,
     ) {
-        let symbol_table = self.linking.symbol_table_mut();
-
         // Lookup vector of symbol index to new function index
-        let mut new_index_lookup = std::vec::Vec::with_capacity(symbol_table.len());
+        let mut new_index_lookup = std::vec::Vec::with_capacity(self.linking.symbol_table.len());
 
         // Modify symbol table entries and fill the lookup vector
-        for sym_info in symbol_table.iter_mut() {
+        for sym_info in self.linking.symbol_table.iter_mut() {
             match sym_info {
                 SymInfo::Function(WasmObjectSymbol::Defined { index, .. }) => {
                     let new_fn_index = *index + n_imported_fns;

compiler/test_gen/src/helpers/wasm32_test_result.rs

@@ -27,12 +27,12 @@ pub trait Wasm32TestResult {
             index,
         });
 
-        let symbol_table = module.linking.symbol_table_mut();
+        let linker_symbol = SymInfo::Function(WasmObjectSymbol::Defined {
-        symbol_table.push(SymInfo::Function(WasmObjectSymbol::Defined {
             flags: 0,
             index,
             name: wrapper_name.to_string(),
-        }));
+        });
+        module.linking.symbol_table.push(linker_symbol);
 
         let mut code_builder = CodeBuilder::new(arena);
         Self::build_wrapper_body(&mut code_builder, main_function_index);