Finalisation and serialisation for byte-level function builder

compiler/gen_wasm/src/function_builder.rs

@@ -3,12 +3,14 @@ use bumpalo::Bump;
 use core::panic;
 use std::fmt::Debug;
 
-use parity_wasm::elements::{Instruction, Instruction::*, Local, Serialize, VarUint32};
 use roc_module::symbol::Symbol;
 
 use crate::code_builder::VirtualMachineSymbolState;
 use crate::opcodes::*;
-use crate::{debug_panic, LocalId, STACK_POINTER_GLOBAL_ID};
+use crate::{
+    encode_u32, encode_u64, round_up_to_alignment, LocalId, FRAME_ALIGNMENT_BYTES,
+    STACK_POINTER_GLOBAL_ID,
+};
 
 const DEBUG_LOG: bool = false;
 
@@ -72,11 +74,15 @@ pub struct FunctionBuilder<'a> {
     /// (Go back and set locals when we realise we need them)
     insertions: Vec<'a, (usize, Vec<'a, u8>)>,
 
-    /// Instruction bytes for locals and stack frame setup code
+    /// Bytes for locals and stack frame setup code
-    locals_and_frame_setup: Vec<'a, u8>,
+    /// We can't write this until we've finished the main code,
+    /// but it goes before the main code in the final output.
+    preamble: Vec<'a, u8>,
 
-    /// Encoded bytes for the inner length of the function
+    /// Encoded bytes for the inner length of the function, locals + code.
-    /// This is the total size of locals + code, as encoded into the module
+    /// ("inner" because it doesn't include the bytes of the length itself!)
+    /// We can't write this until we've finished the code and preamble,
+    /// but it goes before them in the final output.
     inner_length: Vec<'a, u8>,
 
     /// Our simulation model of the Wasm stack machine
@@ -90,7 +96,7 @@ impl<'a> FunctionBuilder<'a> {
         FunctionBuilder {
             code: Vec::with_capacity_in(1024, arena),
             insertions: Vec::with_capacity_in(1024, arena),
-            locals_and_frame_setup: Vec::with_capacity_in(32, arena),
+            preamble: Vec::with_capacity_in(32, arena),
             inner_length: Vec::with_capacity_in(5, arena),
             vm_stack: Vec::with_capacity_in(32, arena),
         }
@@ -137,35 +143,12 @@ impl<'a> FunctionBuilder<'a> {
         true
     }
 
-    fn push_stack_frame(frame_size: i32, local_frame_pointer: LocalId) -> [Instruction; 5] {
+    /// Generate bytes to declare a function's local variables
-        return [
+    fn build_local_declarations(&mut self, local_types: &[ValueType]) {
-            GetGlobal(STACK_POINTER_GLOBAL_ID),
-            I32Const(frame_size),
-            I32Sub,
-            TeeLocal(local_frame_pointer.0),
-            SetGlobal(STACK_POINTER_GLOBAL_ID),
-        ];
-    }
-
-    fn pop_stack_frame(frame_size: i32, local_frame_pointer: LocalId) -> [Instruction; 4] {
-        return [
-            GetLocal(local_frame_pointer.0),
-            I32Const(frame_size),
-            I32Add,
-            SetGlobal(STACK_POINTER_GLOBAL_ID),
-        ];
-    }
-
-    fn serialize_locals<W>(writer: &mut W, local_types: &[parity_wasm::elements::ValueType])
-    where
-        W: std::io::Write,
-    {
         let num_locals = local_types.len();
-        VarUint32::from(num_locals)
+        encode_u32(&mut self.preamble, num_locals as u32);
-            .serialize(writer)
-            .unwrap_or_else(debug_panic);
 
-        // write bytes for locals, in batches of the same ValueType
+        // write declarations in batches of the same ValueType
         if num_locals > 0 {
             let mut batch_type = local_types[0];
             let mut batch_size = 0;
@@ -173,57 +156,76 @@ impl<'a> FunctionBuilder<'a> {
                 if *t == batch_type {
                     batch_size += 1;
                 } else {
-                    let local = Local::new(batch_size, batch_type);
+                    encode_u32(&mut self.preamble, batch_size);
-                    local.serialize(writer).unwrap_or_else(debug_panic);
+                    self.preamble.push(batch_type as u8);
                     batch_type = *t;
                     batch_size = 1;
                 }
             }
-            let local = Local::new(batch_size, batch_type);
+            encode_u32(&mut self.preamble, batch_size);
-            local.serialize(writer).unwrap_or_else(debug_panic);
+            self.preamble.push(batch_type as u8);
         }
     }
 
+    /// Generate instruction bytes to grab a frame of stack memory on entering the function
+    fn build_stack_frame_push(&mut self, frame_size: i32, frame_pointer: LocalId) {
+        // Can't use the usual instruction methods because they push to self.code.
+        // This is the only case where we push instructions somewhere different.
+        self.preamble.push(GETGLOBAL);
+        encode_u32(&mut self.preamble, STACK_POINTER_GLOBAL_ID);
+        self.preamble.push(I32CONST);
+        encode_u32(&mut self.preamble, frame_size as u32);
+        self.preamble.push(I32SUB);
+        self.preamble.push(TEELOCAL);
+        encode_u32(&mut self.preamble, frame_pointer.0);
+        self.preamble.push(SETGLOBAL);
+        encode_u32(&mut self.preamble, STACK_POINTER_GLOBAL_ID);
+    }
+
+    /// Generate instruction bytes to release a frame of stack memory on leaving the function
+    fn build_stack_frame_pop(&mut self, frame_size: i32, frame_pointer: LocalId) {
+        self.get_local(frame_pointer);
+        self.i32_const(frame_size);
+        self.i32_add();
+        self.set_global(STACK_POINTER_GLOBAL_ID);
+    }
+
     /// Finalize a function body
-    pub fn finalize(&mut self, _final_code: &mut std::vec::Vec<Instruction>) {
+    /// Generate all the "extra" bytes: local declarations, stack frame push/pop code, and function length
-        // sort insertions
+    /// After running this, we will know the final byte length of the function. All bytes will have been
-        // encode local declarations
+    /// _generated_, but not yet _serialized_, as they are still stored in multiple vectors.
-        // encode stack frame push
+    pub fn finalize(&mut self, local_types: &[ValueType], frame_size: i32, frame_pointer: LocalId) {
-        // encode stack frame pop
+        self.insertions.sort_by_key(|pair| pair.0);
-        // calculate inner length
+
-        // encode inner length
+        self.build_local_declarations(local_types);
+
+        if frame_size > 0 {
+            let aligned_size = round_up_to_alignment(frame_size, FRAME_ALIGNMENT_BYTES);
+            self.build_stack_frame_push(aligned_size, frame_pointer);
+            self.build_stack_frame_pop(aligned_size, frame_pointer);
         }
 
-    pub fn serialize<W: std::io::Write>(&self, _writer: W) {
+        let inner_length_val = self.preamble.len() + self.code.len() + self.insertions.len();
-        // write inner length
+        encode_u32(&mut self.inner_length, inner_length_val as u32);
-        // write locals
-        // write stack frame push
-        // write code+insertions
-        // write stack frame pop
     }
 
-    /// Total bytes, including inner length
+    /// Write out the function bytes serially
-    /// (to help calculate overall code section length)
+    pub fn serialize<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<usize> {
+        writer.write(&self.inner_length)?;
+        writer.write(&self.preamble)?;
+        let mut pos: usize = 0;
+        for (next_insert_pos, insert_bytes) in self.insertions.iter() {
+            writer.write(&self.code[pos..*next_insert_pos])?;
+            writer.write(insert_bytes)?;
+            pos = *next_insert_pos;
+        }
+        let code_len = self.code.len();
+        writer.write(&self.code[pos..code_len])
+    }
+
+    /// Total bytes this function will occupy in the code section, including its own inner length
     pub fn outer_len(&self) -> usize {
-        self.code.len()
+        self.inner_length.len() + self.preamble.len() + self.code.len() + self.insertions.len()
-            + self.insertions.len()
-            + self.locals_and_frame_setup.len()
-            + self.inner_length.len()
-    }
-
-    pub fn call(&mut self, function_index: u32, pops: usize, push: bool) {
-        let stack_depth = self.vm_stack.len();
-        if pops > stack_depth {
-            panic!(
-                "Trying to call to call function {:?} with {:?} values but only {:?} on the VM stack\n{:?}",
-                function_index, pops, stack_depth, self
-            );
-        }
-        self.vm_stack.truncate(stack_depth - pops);
-        if push {
-            self.vm_stack.push(Symbol::WASM_ANONYMOUS_STACK_VALUE);
-        }
-        self.code.push(CALL);
     }
 
     /// Base method for generating instructions
@@ -244,23 +246,13 @@ impl<'a> FunctionBuilder<'a> {
 
     fn inst_imm32(&mut self, opcode: u8, pops: usize, push: bool, immediate: u32) {
         self.inst(opcode, pops, push);
-        let mut value = immediate;
+        encode_u32(&mut self.code, immediate);
-        while value >= 0x80 {
-            self.code.push(0x80 | ((value & 0x7f) as u8));
-            value >>= 7;
-        }
-        self.code.push(value as u8);
     }
 
     fn inst_mem(&mut self, opcode: u8, pops: usize, push: bool, align: Align, offset: u32) {
         self.inst(opcode, pops, push);
         self.code.push(align as u8);
-        let mut value = offset;
+        encode_u32(&mut self.code, offset);
-        while value >= 0x80 {
-            self.code.push(0x80 | ((value & 0x7f) as u8));
-            value >>= 7;
-        }
-        self.code.push(value as u8);
     }
 
     // Instruction methods
@@ -273,6 +265,7 @@ impl<'a> FunctionBuilder<'a> {
 
     instruction_no_args!(unreachable_, UNREACHABLE, 0, false);
     instruction_no_args!(nop, NOP, 0, false);
+
     pub fn block(&mut self, ty: BlockType) {
         self.inst_imm8(BLOCK, 0, false, ty.as_byte());
     }
@@ -282,20 +275,43 @@ impl<'a> FunctionBuilder<'a> {
     pub fn if_(&mut self, ty: BlockType) {
         self.inst_imm8(IF, 1, false, ty.as_byte());
     }
+
     instruction_no_args!(else_, ELSE, 0, false);
     instruction_no_args!(end, END, 0, false);
+
     pub fn br(&mut self, levels: u32) {
         self.inst_imm32(BR, 0, false, levels);
     }
     pub fn br_if(&mut self, levels: u32) {
         self.inst_imm32(BRIF, 1, false, levels);
     }
-    // br_table: not implemented
+    fn br_table() {
+        panic!("TODO");
+    }
+
     instruction_no_args!(return_, RETURN, 0, false);
-    // call: see above
+
-    // call_indirect: not implemented
+    pub fn call(&mut self, function_index: u32, n_args: usize, has_return_val: bool) {
+        let stack_depth = self.vm_stack.len();
+        if n_args > stack_depth {
+            panic!(
+                "Trying to call to call function {:?} with {:?} values but only {:?} on the VM stack\n{:?}",
+                function_index, n_args, stack_depth, self
+            );
+        }
+        self.vm_stack.truncate(stack_depth - n_args);
+        if has_return_val {
+            self.vm_stack.push(Symbol::WASM_ANONYMOUS_STACK_VALUE);
+        }
+        self.code.push(CALL);
+    }
+    fn call_indirect() {
+        panic!("Not implemented. Roc doesn't use function pointers");
+    }
+
     instruction_no_args!(drop, DROP, 1, false);
     instruction_no_args!(select, SELECT, 3, true);
+
     pub fn get_local(&mut self, id: LocalId) {
         self.inst_imm32(GETLOCAL, 0, true, id.0);
     }
@@ -303,7 +319,7 @@ impl<'a> FunctionBuilder<'a> {
         self.inst_imm32(SETLOCAL, 1, false, id.0);
     }
     pub fn tee_local(&mut self, id: LocalId) {
-        self.inst_imm32(TEELOCAL, 1, true, id.0);
+        self.inst_imm32(TEELOCAL, 0, false, id.0);
     }
     pub fn get_global(&mut self, id: u32) {
         self.inst_imm32(GETGLOBAL, 0, true, id);
@@ -311,6 +327,7 @@ impl<'a> FunctionBuilder<'a> {
     pub fn set_global(&mut self, id: u32) {
         self.inst_imm32(SETGLOBAL, 1, false, id);
     }
+
     instruction_memargs!(i32_load, I32LOAD, 1, true);
     instruction_memargs!(i64_load, I64LOAD, 1, true);
     instruction_memargs!(f32_load, F32LOAD, 1, true);
@@ -334,6 +351,7 @@ impl<'a> FunctionBuilder<'a> {
     instruction_memargs!(i64_store8, I64STORE8, 2, false);
     instruction_memargs!(i64_store16, I64STORE16, 2, false);
     instruction_memargs!(i64_store32, I64STORE32, 2, false);
+
     pub fn memory_size(&mut self) {
         self.inst_imm8(CURRENTMEMORY, 0, true, 0);
     }
@@ -345,12 +363,7 @@ impl<'a> FunctionBuilder<'a> {
     }
     pub fn i64_const(&mut self, x: i64) {
         self.inst(I64CONST, 0, true);
-        let mut value = x as u64;
+        encode_u64(&mut self.code, x as u64);
-        while value >= 0x80 {
-            self.code.push(0x80 | ((value & 0x7f) as u8));
-            value >>= 7;
-        }
-        self.code.push(value as u8);
     }
     pub fn f32_const(&mut self, x: f32) {
         self.inst(F32CONST, 0, true);
@@ -370,8 +383,9 @@ impl<'a> FunctionBuilder<'a> {
             value >>= 8;
         }
     }
-    // TODO: code gen might for "lowlevel" calls be simplified if the numerical op methods
+
-    // took a ValueType as an argument and picked the opcode. Unify all 'eq', all 'add', etc.
+    // TODO: Consider creating unified methods for numerical ops like 'eq' and 'add',
+    // passing the ValueType as an argument. Could simplify lowlevel code gen.
     instruction_no_args!(i32_eqz, I32EQZ, 1, true);
     instruction_no_args!(i32_eq, I32EQ, 2, true);
     instruction_no_args!(i32_ne, I32NE, 2, true);

compiler/gen_wasm/src/lib.rs

@@ -33,7 +33,7 @@ pub const ALIGN_4: u32 = 2;
 pub const ALIGN_8: u32 = 3;
 
 pub const STACK_POINTER_GLOBAL_ID: u32 = 0;
-pub const STACK_ALIGNMENT_BYTES: i32 = 16;
+pub const FRAME_ALIGNMENT_BYTES: i32 = 16;
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub struct LocalId(pub u32);
@@ -189,7 +189,7 @@ pub fn push_stack_frame(
     size: i32,
     local_frame_pointer: LocalId,
 ) {
-    let aligned_size = round_up_to_alignment(size, STACK_ALIGNMENT_BYTES);
+    let aligned_size = round_up_to_alignment(size, FRAME_ALIGNMENT_BYTES);
     instructions.extend([
         GetGlobal(STACK_POINTER_GLOBAL_ID),
         I32Const(aligned_size),
@@ -204,7 +204,7 @@ pub fn pop_stack_frame(
     size: i32,
     local_frame_pointer: LocalId,
 ) {
-    let aligned_size = round_up_to_alignment(size, STACK_ALIGNMENT_BYTES);
+    let aligned_size = round_up_to_alignment(size, FRAME_ALIGNMENT_BYTES);
     instructions.extend([
         GetLocal(local_frame_pointer.0),
         I32Const(aligned_size),
@@ -216,3 +216,27 @@ pub fn pop_stack_frame(
 pub fn debug_panic<E: std::fmt::Debug>(error: E) {
     panic!("{:?}", error);
 }
+
+/// Write a u32 value as LEB-128 encoded bytes, into the provided buffer
+///
+/// All integers in Wasm are variable-length encoded, which saves space for small values.
+/// The most significant bit indicates "more bytes are coming", and the other 7 are payload.
+pub fn encode_u32<'a>(buffer: &mut Vec<'a, u8>, mut value: u32) {
+    while value >= 0x80 {
+        buffer.push(0x80 | ((value & 0x7f) as u8));
+        value >>= 7;
+    }
+    buffer.push(value as u8);
+}
+
+/// Write a u64 value as LEB-128 encoded bytes, into the provided buffer
+///
+/// All integers in Wasm are variable-length encoded, which saves space for small values.
+/// The most significant bit indicates "more bytes are coming", and the other 7 are payload.
+pub fn encode_u64<'a>(buffer: &mut Vec<'a, u8>, mut value: u64) {
+    while value >= 0x80 {
+        buffer.push(0x80 | ((value & 0x7f) as u8));
+        value >>= 7;
+    }
+    buffer.push(value as u8);
+}