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 {
-                segments: bumpalo::vec![in arena; const_segment],
-            },
+            data: DataSection::new(arena),
             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,42 +1433,25 @@ impl<'a> WasmBackend<'a> {
         sym: Symbol,
         layout: &Layout<'a>,
     ) -> (u32, u32) {
-        match self.constant_sym_index_map.get(string) {
-            Some(linker_sym_index) => {
-                // We've seen this string before. The linker metadata has a reference
-                // to its offset in the constants data segment.
-                let syminfo = &self.linker_symbols[*linker_sym_index];
-                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
-                    ),
-                }
-            }
+        // Place the segment at a 4-byte aligned offset
+        let segment_addr = round_up_to_alignment!(self.next_constant_addr, PTR_SIZE);
+        let elements_addr = segment_addr + PTR_SIZE;
+        let length_with_refcount = 4 + string.len();
+        self.next_constant_addr = segment_addr + length_with_refcount as u32;
 
-            None => {
-                let const_segment_bytes = &mut self.module.data.segments[CONST_SEGMENT_INDEX].init;
+        let mut segment = DataSegment {
+            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
-                let aligned_len = round_up_to_alignment!(const_segment_bytes.len(), 4usize);
-                const_segment_bytes.resize(aligned_len, 0);
-
-                // Prefix the string with "infinite" refcount
+        // Prefix the string bytes with "infinite" refcount
         let refcount_max_bytes: [u8; 4] = (REFCOUNT_MAX as i32).to_le_bytes();
-                const_segment_bytes.extend_from_slice(&refcount_max_bytes);
+        segment.init.extend_from_slice(&refcount_max_bytes);
+        segment.init.extend_from_slice(string.as_bytes());
 
-                // Add the string bytes to the data segment
-                let elements_offset = const_segment_bytes.len() as u32;
-                let elements_addr = elements_offset + CONST_SEGMENT_BASE_ADDR;
-                const_segment_bytes.extend_from_slice(string.as_bytes());
+        let segment_index = self.module.data.append_segment(segment);
 
         // Generate linker info
-                // Just pick the symbol name from the first usage
         let name = self
             .layout_ids
             .get(sym, layout)
@@ -1488,19 +1459,16 @@ impl<'a> WasmBackend<'a> {
         let linker_symbol = SymInfo::Data(DataSymbol::Defined {
             flags: 0,
             name,
-                    segment_index: CONST_SEGMENT_INDEX as u32,
-                    segment_offset: elements_offset,
+            segment_index,
+            segment_offset: 4,
             size: string.len() as u32,
         });
 
         let linker_sym_index = self.linker_symbols.len();
-                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]) {
         // TODO: we just calculated storage and now we're getting it out of a map

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 {
-        self.segments.is_empty() || self.segments.iter().all(|seg| seg.init.is_empty())
+    pub fn new(arena: &'a Bump) -> Self {
+        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() {
-            serialize_vector_section(buffer, SectionId::Data, &self.segments);
+        if !self.bytes.is_empty() {
+            serialize_bytes_section(buffer, SectionId::Data, self.segment_count, &self.bytes);
         }
     }
 }