Wasm: Preload WasmModule from object file bytes

compiler/gen_wasm/src/wasm_module/mod.rs

@@ -4,14 +4,16 @@ pub mod opcodes;
 pub mod sections;
 pub mod serialize;
 
+use bumpalo::Bump;
 pub use code_builder::{Align, CodeBuilder, LocalId, ValueType, VmSymbolState};
 pub use linking::SymInfo;
+use roc_reporting::internal_error;
 pub use sections::{ConstExpr, Export, ExportType, Global, GlobalType, Signature};
 
 use self::linking::{LinkingSection, RelocationSection};
 use self::sections::{
     CodeSection, DataSection, ExportSection, FunctionSection, GlobalSection, ImportSection,
-    MemorySection, OpaqueSection, Section, TypeSection,
+    MemorySection, OpaqueSection, Section, SectionId, TypeSection,
 };
 use self::serialize::{SerialBuffer, Serialize};
 
@@ -112,6 +114,47 @@ impl<'a> WasmModule<'a> {
             + self.code.size()
             + self.data.size()
     }
+
+    pub fn preload(arena: &'a Bump, bytes: &[u8]) -> Self {
+        let is_valid_magic_number = &bytes[0..4] == "\0asm".as_bytes();
+        let is_valid_version = &bytes[4..8] == Self::WASM_VERSION.to_le_bytes();
+        if !is_valid_magic_number || !is_valid_version {
+            internal_error!("Invalid Wasm object file header for platform & builtins");
+        }
+
+        let mut cursor: usize = 8;
+
+        let mut types = TypeSection::preload(arena, bytes, &mut cursor);
+        types.cache_offsets();
+        let import = ImportSection::preload(arena, bytes, &mut cursor);
+        let function = FunctionSection::preload(arena, bytes, &mut cursor);
+        let table = OpaqueSection::preload(SectionId::Table, arena, bytes, &mut cursor);
+        let memory = MemorySection::preload(arena, bytes, &mut cursor);
+        let global = GlobalSection::preload(arena, bytes, &mut cursor);
+        let export = ExportSection::preload(arena, bytes, &mut cursor);
+        let start = OpaqueSection::preload(SectionId::Start, arena, bytes, &mut cursor);
+        let element = OpaqueSection::preload(SectionId::Element, arena, bytes, &mut cursor);
+        let code = CodeSection::preload(arena, bytes, &mut cursor);
+        let data = DataSection::preload(arena, bytes, &mut cursor);
+        let linking = LinkingSection::new(arena);
+        let relocations = RelocationSection::new(arena, "reloc.CODE");
+
+        WasmModule {
+            types,
+            import,
+            function,
+            table,
+            memory,
+            global,
+            export,
+            start,
+            element,
+            code,
+            data,
+            linking,
+            relocations,
+        }
+    }
 }
 
 /// Helper struct to count non-empty sections.

compiler/gen_wasm/src/wasm_module/sections.rs

@@ -100,12 +100,12 @@ fn parse_section<'a>(id: SectionId, module_bytes: &'a [u8], cursor: &mut usize)
     *cursor += 1;
 
     let section_size = parse_u32_or_panic(module_bytes, cursor);
-    let count_offset = *cursor;
+    let count_start = *cursor;
     let count = parse_u32_or_panic(module_bytes, cursor);
-    let body_offset = *cursor;
+    let body_start = *cursor;
 
-    let next_section_start = count_offset + section_size as usize;
-    let body = &module_bytes[body_offset..next_section_start];
+    let next_section_start = count_start + section_size as usize;
+    let body = &module_bytes[body_start..next_section_start];
 
     *cursor = next_section_start;
 
@@ -183,14 +183,6 @@ pub struct TypeSection<'a> {
 }
 
 impl<'a> TypeSection<'a> {
-    pub fn new(arena: &'a Bump, capacity: usize) -> Self {
-        TypeSection {
-            arena,
-            bytes: Vec::with_capacity_in(capacity * 4, arena),
-            offsets: Vec::with_capacity_in(capacity, arena),
-        }
-    }
-
     /// Find a matching signature or insert a new one. Return the index.
     pub fn insert(&mut self, signature: Signature<'a>) -> u32 {
         let mut sig_bytes = Vec::with_capacity_in(signature.param_types.len() + 4, self.arena);
@@ -216,7 +208,7 @@ impl<'a> TypeSection<'a> {
         sig_id as u32
     }
 
-    fn populate_offsets(&mut self) {
+    pub fn cache_offsets(&mut self) {
         self.offsets.clear();
         let mut i = 0;
         while i < self.bytes.len() {
@@ -330,13 +322,6 @@ pub struct ImportSection<'a> {
 }
 
 impl<'a> ImportSection<'a> {
-    pub fn new(arena: &'a Bump) -> Self {
-        ImportSection {
-            count: 0,
-            bytes: bumpalo::vec![in arena],
-        }
-    }
-
     pub fn append(&mut self, import: Import) {
         import.serialize(&mut self.bytes);
         self.count += 1;
@@ -359,13 +344,6 @@ pub struct FunctionSection<'a> {
 }
 
 impl<'a> FunctionSection<'a> {
-    pub fn new(arena: &'a Bump, capacity: usize) -> Self {
-        FunctionSection {
-            count: 0,
-            bytes: Vec::with_capacity_in(capacity, arena),
-        }
-    }
-
     pub fn add_sig(&mut self, sig_id: u32) {
         self.bytes.encode_u32(sig_id);
         self.count += 1;
@@ -633,7 +611,16 @@ impl<'a> CodeSection<'a> {
         SIZE_SECTION_HEADER + self.preloaded_bytes.len() + builders_size
     }
 
-    // fn preload(arena: &'a Bump, module_bytes: &[u8], cursor: &mut usize) -> Self {}
+    pub fn preload(arena: &'a Bump, module_bytes: &[u8], cursor: &mut usize) -> Self {
+        let (preloaded_count, initial_bytes) = parse_section(SectionId::Code, module_bytes, cursor);
+        let mut preloaded_bytes = Vec::with_capacity_in(initial_bytes.len() * 2, arena);
+        preloaded_bytes.extend_from_slice(initial_bytes);
+        CodeSection {
+            preloaded_count,
+            preloaded_bytes,
+            code_builders: Vec::with_capacity_in(0, arena),
+        }
+    }
 }
 
 impl<'a> Serialize for CodeSection<'a> {
@@ -702,13 +689,6 @@ pub struct DataSection<'a> {
 }
 
 impl<'a> DataSection<'a> {
-    pub fn new(arena: &'a Bump) -> Self {
-        DataSection {
-            count: 0,
-            bytes: bumpalo::vec![in arena],
-        }
-    }
-
     pub fn append_segment(&mut self, segment: DataSegment<'a>) -> u32 {
         let index = self.count;
         self.count += 1;
@@ -735,20 +715,38 @@ pub struct OpaqueSection<'a> {
 }
 
 impl<'a> OpaqueSection<'a> {
-    pub fn new(bytes: &'a [u8]) -> Self {
-        Self { bytes }
-    }
-
     pub fn size(&self) -> usize {
         self.bytes.len()
     }
+
+    pub fn preload(
+        id: SectionId,
+        arena: &'a Bump,
+        module_bytes: &[u8],
+        cursor: &mut usize,
+    ) -> Self {
+        let bytes: &[u8];
+
+        if module_bytes[*cursor] != id as u8 {
+            bytes = &[];
+        } else {
+            let section_start = *cursor;
+            *cursor += 1;
+            let section_size = parse_u32_or_panic(module_bytes, cursor);
+            let next_section_start = *cursor + section_size as usize;
+            bytes = &module_bytes[section_start..next_section_start];
+            *cursor = next_section_start;
+        };
+
+        OpaqueSection {
+            bytes: arena.alloc_slice_clone(bytes),
+        }
+    }
 }
 
 impl Serialize for OpaqueSection<'_> {
     fn serialize<T: SerialBuffer>(&self, buffer: &mut T) {
-        if !self.bytes.is_empty() {
-            buffer.append_slice(self.bytes);
-        }
+        buffer.append_slice(self.bytes);
     }
 }
 
@@ -766,7 +764,7 @@ mod tests {
 
         // Reconstruct a new TypeSection by "pre-loading" the bytes of the original!
         let body = &original_serialized[6..];
-        let preloaded = TypeSection::populate_offsets(arena, body);
+        let preloaded = TypeSection::cache_offsets(arena, body);
 
         debug_assert_eq!(original.offsets, preloaded.offsets);
         debug_assert_eq!(original.bytes, preloaded.bytes);
@@ -790,7 +788,13 @@ mod tests {
                 ret_type: Some(I32),
             },
         ];
-        let mut section = TypeSection::new(arena, signatures.len());
+        let capacity = signatures.len();
+        let mut section = TypeSection {
+            arena,
+            bytes: Vec::with_capacity_in(capacity * 4, arena),
+            offsets: Vec::with_capacity_in(capacity, arena),
+        };
+
         for sig in signatures {
             section.insert(sig);
         }

compiler/gen_wasm/src/wasm_module/serialize.rs

@@ -254,6 +254,13 @@ pub fn decode_u32_or_panic(bytes: &[u8]) -> (u32, usize) {
     decode_u32(bytes).unwrap_or_else(|e| internal_error!("{}", e))
 }
 
+pub fn parse_u32_or_panic(bytes: &[u8], cursor: &mut usize) -> u32 {
+    let (value, new_offset) =
+        decode_u32(&bytes[*cursor..]).unwrap_or_else(|e| internal_error!("{}", e));
+    *cursor = new_offset;
+    value
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;