Wasm: Store DataSection as bytes and segment count

compiler/gen_wasm/src/backend.rs

@@ -40,9 +40,6 @@ use crate::{
 /// Follow Emscripten's example by leaving 1kB unused (though 4 bytes would probably do!)
 const CONST_SEGMENT_BASE_ADDR: u32 = 1024;
 
-/// Index of the data segment where we store constants
-const CONST_SEGMENT_INDEX: usize = 0;
-
 pub struct WasmBackend<'a> {
     env: &'a Env<'a>,
     interns: &'a mut Interns,
@@ -50,7 +47,7 @@ pub struct WasmBackend<'a> {
     // Module-level data
     module: WasmModule<'a>,
     layout_ids: LayoutIds<'a>,
-    constant_sym_index_map: MutMap<&'a str, usize>,
+    next_constant_addr: u32,
     builtin_sym_index_map: MutMap<&'a str, usize>,
     proc_symbols: Vec<'a, (Symbol, u32)>,
     linker_symbols: Vec<'a, SymInfo>,
@@ -107,13 +104,6 @@ impl<'a> WasmBackend<'a> {
             name: STACK_POINTER_NAME.to_string(),
         }));
 
-        let const_segment = DataSegment {
-            mode: DataMode::Active {
-                offset: ConstExpr::I32(CONST_SEGMENT_BASE_ADDR as i32),
-            },
-            init: Vec::with_capacity_in(64, arena),
-        };
-
         let module = WasmModule {
             types: TypeSection::new(arena, num_procs),
             import: ImportSection::new(arena),
@@ -131,9 +121,7 @@ impl<'a> WasmBackend<'a> {
                 preloaded_bytes: Vec::with_capacity_in(0, arena),
                 code_builders: Vec::with_capacity_in(num_procs, arena),
             },
-            data: DataSection {
+            data: DataSection::new(arena),
-                segments: bumpalo::vec![in arena; const_segment],
-            },
             linking: LinkingSection::new(arena),
             relocations: RelocationSection::new(arena, "reloc.CODE"),
         };
@@ -146,7 +134,7 @@ impl<'a> WasmBackend<'a> {
             module,
 
             layout_ids,
-            constant_sym_index_map: MutMap::default(),
+            next_constant_addr: CONST_SEGMENT_BASE_ADDR,
             builtin_sym_index_map: MutMap::default(),
             proc_symbols,
             linker_symbols,
@@ -1445,61 +1433,41 @@ impl<'a> WasmBackend<'a> {
         sym: Symbol,
         layout: &Layout<'a>,
     ) -> (u32, u32) {
-        match self.constant_sym_index_map.get(string) {
+        // Place the segment at a 4-byte aligned offset
-            Some(linker_sym_index) => {
+        let segment_addr = round_up_to_alignment!(self.next_constant_addr, PTR_SIZE);
-                // We've seen this string before. The linker metadata has a reference
+        let elements_addr = segment_addr + PTR_SIZE;
-                // to its offset in the constants data segment.
+        let length_with_refcount = 4 + string.len();
-                let syminfo = &self.linker_symbols[*linker_sym_index];
+        self.next_constant_addr = segment_addr + length_with_refcount as u32;
-                match syminfo {
-                    SymInfo::Data(DataSymbol::Defined { segment_offset, .. }) => {
-                        let elements_addr = *segment_offset + CONST_SEGMENT_BASE_ADDR;
-                        (*linker_sym_index as u32, elements_addr)
-                    }
-                    _ => internal_error!(
-                        "Compiler bug: Invalid linker symbol info for string {:?}:\n{:?}",
-                        string,
-                        syminfo
-                    ),
-                }
-            }
 
-            None => {
+        let mut segment = DataSegment {
-                let const_segment_bytes = &mut self.module.data.segments[CONST_SEGMENT_INDEX].init;
+            mode: DataMode::active_at(segment_addr),
+            init: Vec::with_capacity_in(length_with_refcount, self.env.arena),
+        };
 
-                // Pad the existing data segment to make sure the refcount and string are aligned
+        // Prefix the string bytes with "infinite" refcount
-                let aligned_len = round_up_to_alignment!(const_segment_bytes.len(), 4usize);
+        let refcount_max_bytes: [u8; 4] = (REFCOUNT_MAX as i32).to_le_bytes();
-                const_segment_bytes.resize(aligned_len, 0);
+        segment.init.extend_from_slice(&refcount_max_bytes);
+        segment.init.extend_from_slice(string.as_bytes());
 
-                // Prefix the string with "infinite" refcount
+        let segment_index = self.module.data.append_segment(segment);
-                let refcount_max_bytes: [u8; 4] = (REFCOUNT_MAX as i32).to_le_bytes();
-                const_segment_bytes.extend_from_slice(&refcount_max_bytes);
 
-                // Add the string bytes to the data segment
+        // Generate linker info
-                let elements_offset = const_segment_bytes.len() as u32;
+        let name = self
-                let elements_addr = elements_offset + CONST_SEGMENT_BASE_ADDR;
+            .layout_ids
-                const_segment_bytes.extend_from_slice(string.as_bytes());
+            .get(sym, layout)
+            .to_symbol_string(sym, self.interns);
+        let linker_symbol = SymInfo::Data(DataSymbol::Defined {
+            flags: 0,
+            name,
+            segment_index,
+            segment_offset: 4,
+            size: string.len() as u32,
+        });
 
-                // Generate linker info
+        let linker_sym_index = self.linker_symbols.len();
-                // Just pick the symbol name from the first usage
+        self.linker_symbols.push(linker_symbol);
-                let name = self
-                    .layout_ids
-                    .get(sym, layout)
-                    .to_symbol_string(sym, self.interns);
-                let linker_symbol = SymInfo::Data(DataSymbol::Defined {
-                    flags: 0,
-                    name,
-                    segment_index: CONST_SEGMENT_INDEX as u32,
-                    segment_offset: elements_offset,
-                    size: string.len() as u32,
-                });
 
-                let linker_sym_index = self.linker_symbols.len();
+        (linker_sym_index as u32, elements_addr)
-                self.constant_sym_index_map.insert(string, linker_sym_index);
-                self.linker_symbols.push(linker_symbol);
-
-                (linker_sym_index as u32, elements_addr)
-            }
-        }
     }
 
     fn create_struct(&mut self, sym: &Symbol, layout: &Layout<'a>, fields: &'a [Symbol]) {

compiler/gen_wasm/src/wasm_module/sections.rs

@@ -530,6 +530,14 @@ pub enum DataMode {
     Passive,
 }
 
+impl DataMode {
+    pub fn active_at(offset: u32) -> Self {
+        DataMode::Active {
+            offset: ConstExpr::I32(offset as i32),
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct DataSegment<'a> {
     pub mode: DataMode,
@@ -554,19 +562,30 @@ impl Serialize for DataSegment<'_> {
 
 #[derive(Debug)]
 pub struct DataSection<'a> {
-    pub segments: Vec<'a, DataSegment<'a>>,
+    segment_count: u32,
+    bytes: Vec<'a, u8>,
 }
 
 impl<'a> DataSection<'a> {
-    fn is_empty(&self) -> bool {
+    pub fn new(arena: &'a Bump) -> Self {
-        self.segments.is_empty() || self.segments.iter().all(|seg| seg.init.is_empty())
+        DataSection {
+            segment_count: 0,
+            bytes: bumpalo::vec![in arena],
+        }
+    }
+
+    pub fn append_segment(&mut self, segment: DataSegment<'a>) -> u32 {
+        let index = self.segment_count;
+        self.segment_count += 1;
+        segment.serialize(&mut self.bytes);
+        index
     }
 }
 
 impl Serialize for DataSection<'_> {
     fn serialize<T: SerialBuffer>(&self, buffer: &mut T) {
-        if !self.is_empty() {
+        if !self.bytes.is_empty() {
-            serialize_vector_section(buffer, SectionId::Data, &self.segments);
+            serialize_bytes_section(buffer, SectionId::Data, self.segment_count, &self.bytes);
         }
     }
 }