Refactor code section to get rid of a copy

compiler/gen_wasm/src/backend.rs

@@ -11,7 +11,6 @@ use roc_mono::layout::{Builtin, Layout};
 use crate::code_builder::{BlockType, CodeBuilder, ValueType};
 use crate::layout::WasmLayout;
 use crate::module_builder::WasmModule;
-use crate::serialize::SerialBuffer;
 use crate::storage::{Storage, StoredValue, StoredValueKind};
 use crate::{copy_memory, CopyMemoryConfig, Env, LocalId, PTR_TYPE};
 
@@ -43,18 +42,11 @@ pub struct WasmBackend<'a> {
 
 impl<'a> WasmBackend<'a> {
     pub fn new(env: &'a Env<'a>, proc_symbols: Vec<'a, Symbol>) -> Self {
-        let mut module = WasmModule::new(env.arena);
-
-        // Code section header
-        module.code.bytes.reserve_padded_u32(); // byte length, to be written at the end
-        let num_procs = proc_symbols.len() as u32;
-        module.code.bytes.encode_padded_u32(num_procs); // modified later in unit tests
-
         WasmBackend {
             env,
 
             // Module-level data
-            module,
+            module: WasmModule::new(env.arena),
             parity_builder: builder::module(),
             _data_offset_map: MutMap::default(),
             _data_offset_next: UNUSED_DATA_SECTION_BYTES,
@@ -68,8 +60,13 @@ impl<'a> WasmBackend<'a> {
         }
     }
 
+    /// Reset function-level data
     fn reset(&mut self) {
-        self.code_builder.clear();
+        // Push the completed CodeBuilder into the module and swap it for a new empty one
+        let mut swap_code_builder = CodeBuilder::new(self.env.arena);
+        std::mem::swap(&mut swap_code_builder, &mut self.code_builder);
+        self.module.code.code_builders.push(swap_code_builder);
+
         self.storage.clear();
         self.joinpoint_label_map.clear();
         assert_eq!(self.block_depth, 0);
@@ -140,8 +137,6 @@ impl<'a> WasmBackend<'a> {
             self.storage.stack_frame_pointer,
         );
 
-        let relocs = self.code_builder.serialize(&mut self.module.code.bytes);
-        self.module.reloc_code.entries.extend(relocs);
         Ok(())
     }
 

compiler/gen_wasm/src/code_builder.rs

@@ -1,4 +1,4 @@
-use bumpalo::collections::vec::{Drain, Vec};
+use bumpalo::collections::vec::Vec;
 use bumpalo::Bump;
 use core::panic;
 use std::fmt::Debug;
@@ -92,8 +92,8 @@ pub enum VirtualMachineSymbolState {
 
 // An instruction (local.set or local.tee) to be inserted into the function code
 #[derive(Debug)]
-struct InsertLocation {
+struct Insertion {
-    insert_at: usize,
+    at: usize,
     start: usize,
     end: usize,
 }
@@ -124,7 +124,7 @@ pub struct CodeBuilder<'a> {
     insert_bytes: Vec<'a, u8>,
 
     /// Code locations where the insert_bytes should go
-    insert_locations: Vec<'a, InsertLocation>,
+    insertions: Vec<'a, Insertion>,
 
     /// Bytes for local variable declarations and stack-frame setup code.
     /// We can't write this until we've finished the main code. But it goes
@@ -151,7 +151,7 @@ impl<'a> CodeBuilder<'a> {
     pub fn new(arena: &'a Bump) -> Self {
         CodeBuilder {
             code: Vec::with_capacity_in(1024, arena),
-            insert_locations: Vec::with_capacity_in(32, arena),
+            insertions: Vec::with_capacity_in(32, arena),
             insert_bytes: Vec::with_capacity_in(64, arena),
             preamble: Vec::with_capacity_in(32, arena),
             inner_length: Vec::with_capacity_in(5, arena),
@@ -160,15 +160,6 @@ impl<'a> CodeBuilder<'a> {
         }
     }
 
-    pub fn clear(&mut self) {
-        self.code.clear();
-        self.insert_locations.clear();
-        self.insert_bytes.clear();
-        self.preamble.clear();
-        self.inner_length.clear();
-        self.vm_stack.clear();
-    }
-
     /**********************************************************
 
         SYMBOLS
@@ -220,8 +211,8 @@ impl<'a> CodeBuilder<'a> {
         self.insert_bytes.push(opcode);
         self.insert_bytes.encode_u32(immediate);
 
-        self.insert_locations.push(InsertLocation {
+        self.insertions.push(Insertion {
-            insert_at,
+            at: insert_at,
             start,
             end: self.insert_bytes.len(),
         });
@@ -348,52 +339,56 @@ impl<'a> CodeBuilder<'a> {
 
         let inner_len = self.preamble.len() + self.code.len() + self.insert_bytes.len();
         self.inner_length.encode_u32(inner_len as u32);
-    }
-
-    /// Write out all the bytes in the right order
-    pub fn serialize<T: SerialBuffer>(
-        &mut self,
-        code_section_buf: &mut T,
-    ) -> Drain<RelocationEntry> {
-        code_section_buf.append_slice(&self.inner_length);
-        code_section_buf.append_slice(&self.preamble);
 
         // Sort insertions. They are not created in order of assignment, but in order of *second* usage.
-        self.insert_locations.sort_by_key(|loc| loc.insert_at);
+        self.insertions.sort_by_key(|ins| ins.at);
+    }
+
+    /// Serialize all byte vectors in the right order
+    /// Also update relocation offsets relative to the provided base offset in the buffer
+    pub fn serialize_with_relocs<T: SerialBuffer>(
+        &self,
+        buffer: &mut T,
+        final_relocs: &mut Vec<'a, RelocationEntry>,
+        reloc_base_offset: usize,
+    ) {
+        buffer.append_slice(&self.inner_length);
+        buffer.append_slice(&self.preamble);
 
         // Do the insertions & update relocation offsets
-        const CODE_SECTION_BODY_OFFSET: usize = 5;
         let mut reloc_index = 0;
-        let mut code_pos: usize = 0;
+        let mut code_pos = 0;
-        for location in self.insert_locations.iter() {
+        let mut insert_iter = self.insertions.iter();
+
+        loop {
+            let next_insert = insert_iter.next();
+            let next_pos = next_insert.map(|i| i.at).unwrap_or(self.code.len());
+
             // Relocation offset needs to be an index into the body of the code section, but
             // at this point it is an index into self.code. Need to adjust for all previous functions
             // in the code section, and for insertions in the current function.
-            let section_body_pos = code_section_buf.size() - CODE_SECTION_BODY_OFFSET;
+            let section_body_pos = buffer.size() - reloc_base_offset;
             while reloc_index < self.relocations.len()
-                && self.relocations[reloc_index].offset() < location.insert_at as u32
+                && self.relocations[reloc_index].offset() < next_pos as u32
             {
-                let offset_ref = self.relocations[reloc_index].offset_mut();
+                let mut reloc_clone = self.relocations[reloc_index].clone();
-                *offset_ref += (section_body_pos - code_pos) as u32;
+                *reloc_clone.offset_mut() += (section_body_pos - code_pos) as u32;
+                final_relocs.push(reloc_clone);
                 reloc_index += 1;
             }
 
-            code_section_buf.append_slice(&self.code[code_pos..location.insert_at]);
+            buffer.append_slice(&self.code[code_pos..next_pos]);
-            code_section_buf.append_slice(&self.insert_bytes[location.start..location.end]);
+
-            code_pos = location.insert_at;
+            match next_insert {
+                Some(Insertion { at, start, end }) => {
+                    buffer.append_slice(&self.insert_bytes[*start..*end]);
+                    code_pos = *at;
+                }
+                None => {
+                    break;
+                }
+            }
         }
-
-        let section_body_pos = code_section_buf.size() - CODE_SECTION_BODY_OFFSET;
-        while reloc_index < self.relocations.len() {
-            let offset_ref = self.relocations[reloc_index].offset_mut();
-            *offset_ref += (section_body_pos - code_pos) as u32;
-            reloc_index += 1;
-        }
-
-        let len = self.code.len();
-        code_section_buf.append_slice(&self.code[code_pos..len]);
-
-        self.relocations.drain(0..)
     }
 
     /**********************************************************

compiler/gen_wasm/src/lib.rs

@@ -83,14 +83,6 @@ pub fn build_module_help<'a>(
         }
     }
 
-    // Update code section length
-    let inner_length = (backend.module.code.bytes.len() - 5) as u32;
-    backend
-        .module
-        .code
-        .bytes
-        .overwrite_padded_u32(0, inner_length);
-
     let symbol_table = LinkingSubSection::SymbolTable(symbol_table_entries);
     backend.module.linking.subsections.push(symbol_table);
 
@@ -178,8 +170,8 @@ pub fn combine_and_serialize<'a>(
     // wasm_module.data_count.serialize(buffer);
     // maybe_increment_section(buffer.size(), &mut prev_size, &mut index);
 
-    wasm_module.code.serialize(buffer);
     wasm_module.reloc_code.target_section_index = Some(index);
+    wasm_module.code.serialize_mut(buffer, &mut wasm_module.reloc_code.entries);
     maybe_increment_section(buffer.size(), &mut prev_size, &mut index);
 
     wasm_module.data.serialize(buffer);

compiler/gen_wasm/src/module_builder.rs

@@ -1,7 +1,7 @@
 use bumpalo::collections::vec::Vec;
 use bumpalo::Bump;
 
-use crate::code_builder::{Align, ValueType};
+use crate::code_builder::{Align, CodeBuilder, ValueType};
 use crate::opcodes;
 use crate::serialize::{SerialBuffer, Serialize};
 
@@ -456,26 +456,30 @@ impl<'a> Serialize for ExportSection<'a> {
 
 #[derive(Debug)]
 pub struct CodeSection<'a> {
-    pub bytes: Vec<'a, u8>,
+    pub code_builders: Vec<'a, CodeBuilder<'a>>,
 }
 
 impl<'a> CodeSection<'a> {
     pub fn new(arena: &'a Bump) -> Self {
         CodeSection {
-            bytes: Vec::with_capacity_in(4096, arena),
+            code_builders: Vec::with_capacity_in(8, arena),
         }
     }
-}
 
-impl<'a> Serialize for CodeSection<'a> {
+    /// Serialize the code builders for all functions, and get code relocations with final offsets
-    fn serialize<T: SerialBuffer>(&self, buffer: &mut T) {
+    pub fn serialize_mut<T: SerialBuffer>(
-        buffer.append_u8(SectionId::Code as u8);
+        &mut self,
+        buffer: &mut T,
+        relocations: &mut Vec<'a, RelocationEntry>,
+    ) {
+        let header_indices = write_section_header(buffer, SectionId::Code);
+        buffer.encode_u32(self.code_builders.len() as u32);
 
-        // TODO
+        for code_builder in self.code_builders.iter_mut() {
-        // We've copied each function into self.bytes, now we're copying again.
+            code_builder.serialize_with_relocs(buffer, relocations, header_indices.body_index);
-        // Can eliminate one of those copies by refactoring to a vector of CodeBuilders
+        }
 
-        buffer.append_slice(&self.bytes);
+        update_section_size(buffer, header_indices);
     }
 }
 
@@ -542,7 +546,7 @@ pub enum OffsetRelocType {
     MemoryAddrI64 = 16,
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub enum RelocationEntry {
     Index {
         type_id: IndexRelocType,
@@ -1044,8 +1048,9 @@ impl<'a> WasmModule<'a> {
         self.data_count.serialize(buffer);
         maybe_increment_section(buffer.size(), &mut prev_size, &mut index);
 
-        self.code.serialize(buffer);
         self.reloc_code.target_section_index = Some(index);
+        self.code
+            .serialize_mut(buffer, &mut self.reloc_code.entries);
         maybe_increment_section(buffer.size(), &mut prev_size, &mut index);
 
         self.data.serialize(buffer);

compiler/gen_wasm/tests/helpers/wasm32_test_result.rs

@@ -3,7 +3,7 @@ use parity_wasm::builder;
 use roc_gen_wasm::code_builder::{Align, CodeBuilder, ValueType};
 use roc_gen_wasm::from_wasm32_memory::FromWasm32Memory;
 use roc_gen_wasm::module_builder::{Export, ExportType, WasmModule};
-use roc_gen_wasm::{serialize::SerialBuffer, LocalId};
+use roc_gen_wasm::LocalId;
 use roc_std::{RocDec, RocList, RocOrder, RocStr};
 
 pub trait Wasm32TestResult {
@@ -28,19 +28,9 @@ pub trait Wasm32TestResult {
             index: location.body,
         });
 
-
         let mut code_builder = CodeBuilder::new(arena);
         Self::build_wrapper_body(&mut code_builder, main_function_index);
-
+        wasm_module.code.code_builders.push(code_builder);
-        code_builder.serialize(&mut wasm_module.code.bytes);
-
-        let mut num_procs = 0;
-        for (i, byte) in wasm_module.code.bytes[5..10].iter().enumerate() {
-            num_procs += ((byte & 0x7f) as u32) << (i * 7);
-        }
-        let inner_length = (wasm_module.code.bytes.len() - 5) as u32;
-        wasm_module.code.bytes.overwrite_padded_u32(0, inner_length);
-        wasm_module.code.bytes.overwrite_padded_u32(5, num_procs + 1);
     }
 
     fn build_wrapper_body(code_builder: &mut CodeBuilder, main_function_index: u32);