Wasm: add preloading logic for Types section

compiler/gen_wasm/src/wasm_module/sections.rs

@@ -5,7 +5,7 @@ use super::linking::{
     IndexRelocType, LinkingSection, RelocationEntry, RelocationSection, SymInfo, WasmObjectSymbol,
 };
 use super::opcodes::OpCode;
-use super::serialize::{SerialBuffer, Serialize};
+use super::serialize::{decode_u32_or_panic, SerialBuffer, Serialize};
 use super::{CodeBuilder, ValueType};
 
 /*******************************************************************
@@ -112,9 +112,13 @@ pub struct Signature<'a> {
     pub ret_type: Option<ValueType>,
 }
 
+impl Signature<'_> {
+    pub const SEPARATOR: u8 = 0x60;
+}
+
 impl<'a> Serialize for Signature<'a> {
     fn serialize<T: SerialBuffer>(&self, buffer: &mut T) {
-        buffer.append_u8(0x60);
+        buffer.append_u8(Self::SEPARATOR);
         self.param_types.serialize(buffer);
         self.ret_type.serialize(buffer);
     }
@@ -161,6 +165,45 @@ impl<'a> TypeSection<'a> {
 
         sig_id as u32
     }
+
+    pub fn preload(arena: &'a Bump, section_body: &[u8]) -> Self {
+        if section_body.is_empty() {
+            return TypeSection {
+                arena,
+                bytes: Vec::new_in(arena),
+                offsets: Vec::new_in(arena),
+            };
+        }
+
+        let (count, content_offset) = decode_u32_or_panic(section_body);
+
+        let mut bytes = Vec::with_capacity_in(section_body.len() * 2, arena);
+        bytes.extend_from_slice(&section_body[content_offset..]);
+
+        let mut offsets = Vec::with_capacity_in((count * 2) as usize, arena);
+
+        let mut i = 0;
+        while i < bytes.len() {
+            offsets.push(i);
+
+            let sep = bytes[i];
+            debug_assert!(sep == Signature::SEPARATOR);
+            i += 1;
+
+            let (n_params, n_params_size) = decode_u32_or_panic(&bytes[i..]);
+            i += n_params_size; // skip over the array length that we just decoded
+            i += n_params as usize; // skip over one byte per param type
+
+            let n_return_values = bytes[i];
+            i += 1 + n_return_values as usize;
+        }
+
+        TypeSection {
+            arena,
+            bytes,
+            offsets,
+        }
+    }
 }
 
 impl<'a> Serialize for TypeSection<'a> {
@@ -781,3 +824,28 @@ impl<'a> WasmModule<'a> {
         }
     }
 }
+
+/// Assertion to run on the generated Wasm module in every test
+#[cfg(debug_assertions)]
+pub fn test_assert_preload<'a>(arena: &'a Bump, wasm_module: &WasmModule<'a>) {
+    test_assert_types_preload(arena, &wasm_module.types);
+}
+
+#[cfg(debug_assertions)]
+fn test_assert_types_preload<'a>(arena: &'a Bump, original: &TypeSection<'a>) {
+    // Serialize the Type section that we built from Roc code
+    let mut original_serialized = Vec::with_capacity_in(original.bytes.len() + 10, arena);
+    original.serialize(&mut original_serialized);
+
+    debug_assert!(original_serialized[0] == SectionId::Type as u8);
+
+    // Reconstruct a new TypeSection by "pre-loading" the bytes
+    let body = &original_serialized[6..];
+
+    let preloaded = TypeSection::preload(arena, body);
+
+    let mut preloaded_serialized = Vec::with_capacity_in(original.bytes.len() + 10, arena);
+    preloaded.serialize(&mut preloaded_serialized);
+
+    debug_assert_eq!(original_serialized, preloaded_serialized);
+}

compiler/gen_wasm/src/wasm_module/serialize.rs

@@ -1,6 +1,7 @@
-use std::fmt::Debug;
+use std::{fmt::Debug, iter::FromIterator};
 
 use bumpalo::collections::vec::Vec;
+use roc_reporting::internal_error;
 
 pub trait Serialize {
     fn serialize<T: SerialBuffer>(&self, buffer: &mut T);
@@ -231,6 +232,28 @@ impl<'a> SerialBuffer for Vec<'a, u8> {
     }
 }
 
+/// Decode an unsigned 32-bit integer from the provided buffer in LEB-128 format
+/// Return the integer itself and the offset after it ends
+pub fn decode_u32(bytes: &[u8]) -> Result<(u32, usize), String> {
+    let mut value = 0;
+    let mut shift = 0;
+    for (i, byte) in bytes.iter().take(5).enumerate() {
+        value += ((byte & 0x7f) as u32) << shift;
+        if (byte & 0x80) == 0 {
+            return Ok((value, i + 1));
+        }
+        shift += 7;
+    }
+    Err(format!(
+        "Failed to decode u32 as LEB-128 from bytes: {:2x?}",
+        std::vec::Vec::from_iter(bytes.iter().take(5))
+    ))
+}
+
+pub fn decode_u32_or_panic(bytes: &[u8]) -> (u32, usize) {
+    decode_u32(bytes).unwrap_or_else(|e| internal_error!("{}", e))
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;