Compile a simple function (#2)

* Implement basics. * Execute code * Add wasm2wat test cases. * abi_loc_for_arg for stack. * Assert that sp_depth is 0 at the epilogue * Do 32bit add. * Assert that RAX can be used as a scratch register * Reuse assembler. * Align stack slots.
2018-11-08 22:56:27 +01:00
parent d3bc26bc93
commit 08240761d5
9 changed files with 338 additions and 72 deletions
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -15,6 +15,7 @@ wasmparser = "0.21.6"
 capstone = "0.5.0"
 failure = "0.1.3"
 failure_derive = "0.1.3"
+wabt = "0.7"

 [badges]
 maintenance = { status = "experimental" }
--- a/examples/test.rs
+++ b/examples/test.rs
@@ -21,7 +21,10 @@ fn read_to_end<P: AsRef<Path>>(path: P) -> io::Result<Vec<u8>> {

 fn maybe_main() -> Result<(), String> {
    let data = read_to_end("test.wasm").map_err(|e| e.to_string())?;
-    translate(&data).map_err(|e| e.to_string())?;
+    let translated = translate(&data).map_err(|e| e.to_string())?;
+    let result = translated.execute_func(0, 5, 3);
+    println!("f(5, 3) = {}", result);
+
    Ok(())
 }

--- a/src/backend.rs
+++ b/src/backend.rs
@@ -1,7 +1,11 @@
 #![allow(dead_code)] // for now

+use error::Error;
 use dynasmrt::x64::Assembler;
-use dynasmrt::DynasmApi;
+use dynasmrt::{DynasmApi, AssemblyOffset, ExecutableBuffer};
+
+/// Size of a pointer on the target in bytes.
+const WORD_SIZE: u32 = 8;

 type GPR = u8;

@@ -15,22 +19,22 @@ impl GPRs {
    }
 }

-static RAX: u8 = 0;
-static RCX: u8 = 1;
-static RDX: u8 = 2;
-static RBX: u8 = 3;
-static RSP: u8 = 4;
-static RBP: u8 = 5;
-static RSI: u8 = 6;
-static RDI: u8 = 7;
-static R8: u8 = 8;
-static R9: u8 = 9;
-static R10: u8 = 10;
-static R11: u8 = 11;
-static R12: u8 = 12;
-static R13: u8 = 13;
-static R14: u8 = 14;
-static R15: u8 = 15;
+const RAX: u8 = 0;
+const RCX: u8 = 1;
+const RDX: u8 = 2;
+const RBX: u8 = 3;
+const RSP: u8 = 4;
+const RBP: u8 = 5;
+const RSI: u8 = 6;
+const RDI: u8 = 7;
+const R8: u8 = 8;
+const R9: u8 = 9;
+const R10: u8 = 10;
+const R11: u8 = 11;
+const R12: u8 = 12;
+const R13: u8 = 13;
+const R14: u8 = 14;
+const R15: u8 = 15;

 impl GPRs {
    fn take(&mut self) -> GPR {
@@ -41,13 +45,16 @@ impl GPRs {
    }

    fn release(&mut self, gpr: GPR) {
-        assert_eq!(
-            self.bits & (1 << gpr),
-            0,
-            "released register was already free"
+        assert!(
+            !self.is_free(gpr),
+            "released register was already free",
        );
        self.bits |= 1 << gpr;
    }
+
+    fn is_free(&self, gpr: GPR) -> bool {
+        (self.bits & (1 << gpr)) != 0
+    }
 }

 pub struct Registers {
@@ -75,35 +82,210 @@ impl Registers {
    }
 }

-fn push_i32(ops: &mut Assembler, regs: &mut Registers, gpr: GPR) {
-    // For now, do an actual push (and pop below). In the future, we could
-    // do on-the-fly register allocation here.
-    dynasm!(ops
-        ; push Rq(gpr)
-    );
-    regs.release_scratch_gpr(gpr);
+/// Describes location of a argument.
+enum ArgLocation {
+    /// Argument is passed via some register.
+    Reg(GPR),
+    /// Value is passed thru the stack.
+    Stack(i32),
 }

-fn pop_i32(ops: &mut Assembler, regs: &mut Registers) -> GPR {
-    let gpr = regs.take_scratch_gpr();
-    dynasm!(ops
+/// Get a location for an argument at the given position.
+fn abi_loc_for_arg(pos: u32) -> ArgLocation {
+    // TODO: This assumes only system-v calling convention.
+    // In system-v calling convention the first 6 arguments are passed via registers. 
+    // All rest arguments are passed on the stack.
+    const ARGS_IN_GPRS: &'static [GPR] = &[
+        RDI,
+        RSI,
+        RDX,
+        RCX,
+        R8,
+        R9,
+    ];
+
+    if let Some(&reg) = ARGS_IN_GPRS.get(pos as usize) {
+        ArgLocation::Reg(reg)
+    } else {
+        let stack_pos = pos - ARGS_IN_GPRS.len() as u32;
+        // +2 is because the first argument is located right after the saved frame pointer slot 
+        // and the incoming return address.
+        let stack_offset = ((stack_pos + 2) * WORD_SIZE) as i32;
+        ArgLocation::Stack(stack_offset)
+    }
+}
+
+pub struct CodeGenSession {
+    assembler: Assembler,
+    func_starts: Vec<AssemblyOffset>,
+}
+
+impl CodeGenSession {
+    pub fn new() -> Self {
+        CodeGenSession {
+            assembler: Assembler::new().unwrap(),
+            func_starts: Vec::new(),
+        }
+    }
+
+    pub fn new_context(&mut self) -> Context {
+        let start_offset = self.assembler.offset();
+        self.func_starts.push(start_offset);
+        Context {
+            asm: &mut self.assembler,
+            start: start_offset,
+            regs: Registers::new(),
+            sp_depth: 0,
+        }
+    }
+
+    pub fn into_translated_code_section(self) -> Result<TranslatedCodeSection, Error>  {
+        let exec_buf = self.assembler
+            .finalize()
+            .map_err(|_asm| Error::Assembler("assembler error".to_owned()))?;
+        Ok(TranslatedCodeSection { exec_buf, func_starts: self.func_starts })
+    }
+}
+
+pub struct TranslatedCodeSection {
+    exec_buf: ExecutableBuffer,
+    func_starts: Vec<AssemblyOffset>,
+}
+
+impl TranslatedCodeSection {
+    pub fn func_start(&self, idx: usize) -> *const u8 {
+        let offset = self.func_starts[idx];
+        self.exec_buf.ptr(offset)
+    }
+}
+
+pub struct Context<'a> {
+    asm: &'a mut Assembler,
+    start: AssemblyOffset,
+    regs: Registers,
+    /// Offset from starting value of SP counted in words. Each push and pop 
+    /// on the value stack increments or decrements this value by 1 respectively.
+    sp_depth: u32,
+}
+
+impl<'a> Context<'a> {
+    /// Returns the offset of the first instruction.
+    fn start(&self) -> AssemblyOffset {
+        self.start
+    }
+}
+
+fn push_i32(ctx: &mut Context, gpr: GPR) {
+    // For now, do an actual push (and pop below). In the future, we could
+    // do on-the-fly register allocation here.
+    ctx.sp_depth += 1;
+    dynasm!(ctx.asm
+        ; push Rq(gpr)
+    );
+    ctx.regs.release_scratch_gpr(gpr);
+}
+
+fn pop_i32(ctx: &mut Context) -> GPR {
+    ctx.sp_depth -= 1;
+    let gpr = ctx.regs.take_scratch_gpr();
+    dynasm!(ctx.asm
        ; pop Rq(gpr)
    );
    gpr
 }

-pub fn add_i32(ops: &mut Assembler, regs: &mut Registers) {
-    let op0 = pop_i32(ops, regs);
-    let op1 = pop_i32(ops, regs);
-    dynasm!(ops
-        ; add Rq(op0), Rq(op1)
+pub fn add_i32(ctx: &mut Context) {
+    let op0 = pop_i32(ctx);
+    let op1 = pop_i32(ctx);
+    dynasm!(ctx.asm
+        ; add Rd(op0), Rd(op1)
    );
-    push_i32(ops, regs, op0);
-    regs.release_scratch_gpr(op1);
+    push_i32(ctx, op0);
+    ctx.regs.release_scratch_gpr(op1);
 }

-pub fn unsupported_opcode(ops: &mut Assembler) {
-    dynasm!(ops
+fn sp_relative_offset(ctx: &mut Context, slot_idx: u32) -> i32 {
+    ((ctx.sp_depth as i32) + slot_idx as i32) * WORD_SIZE as i32
+}
+
+pub fn get_local_i32(ctx: &mut Context, local_idx: u32) {
+    let gpr = ctx.regs.take_scratch_gpr();
+    let offset = sp_relative_offset(ctx, local_idx);
+    dynasm!(ctx.asm
+        ; mov Rq(gpr), [rsp + offset]
+    );
+    push_i32(ctx, gpr);
+}
+
+pub fn store_i32(ctx: &mut Context, local_idx: u32) {
+    let gpr = pop_i32(ctx);
+    let offset = sp_relative_offset(ctx, local_idx);
+    dynasm!(ctx.asm
+        ; mov [rsp + offset], Rq(gpr)
+    );
+    ctx.regs.release_scratch_gpr(gpr);
+}
+
+pub fn prepare_return_value(ctx: &mut Context) {
+    let ret_gpr = pop_i32(ctx);
+    if ret_gpr != RAX {
+        dynasm!(ctx.asm
+            ; mov Rq(RAX), Rq(ret_gpr)
+        );
+        ctx.regs.release_scratch_gpr(ret_gpr);
+    }
+}
+
+pub fn copy_incoming_arg(ctx: &mut Context, arg_pos: u32) {
+    let loc = abi_loc_for_arg(arg_pos);
+
+    // First, ensure the argument is in a register.
+    let reg = match loc {
+        ArgLocation::Reg(reg) => reg,
+        ArgLocation::Stack(offset) => {
+            assert!(
+                ctx.regs.scratch_gprs.is_free(RAX),
+                "we assume that RAX can be used as a scratch register for now",
+            );
+            dynasm!(ctx.asm
+                ; mov Rq(RAX), [rsp + offset]
+            );
+            RAX
+        }
+    };
+
+    // And then move a value from a register into local variable area on the stack.
+    let offset = sp_relative_offset(ctx, arg_pos);
+    dynasm!(ctx.asm
+        ; mov [rsp + offset], Rq(reg) 
+    );
+}
+
+pub fn prologue(ctx: &mut Context, stack_slots: u32) {
+    // Align stack slots to the nearest even number. This is required
+    // by x86-64 ABI.
+    let aligned_stack_slots = (stack_slots + 1) & !1;
+
+    let framesize: i32 = aligned_stack_slots as i32 * WORD_SIZE as i32;
+    dynasm!(ctx.asm
+        ; push rbp
+        ; mov rbp, rsp
+        ; sub rsp, framesize
+    );
+    ctx.sp_depth += aligned_stack_slots - stack_slots;
+}
+
+pub fn epilogue(ctx: &mut Context) {
+    assert_eq!(ctx.sp_depth, 0, "imbalanced pushes and pops detected");
+    dynasm!(ctx.asm
+        ; mov rsp, rbp
+        ; pop rbp
+        ; ret
+    );
+}
+
+pub fn unsupported_opcode(ctx: &mut Context) {
+    dynasm!(ctx.asm
        ; ud2
    );
 }
--- a/src/disassemble.rs
+++ b/src/disassemble.rs
@@ -2,6 +2,7 @@ use capstone::prelude::*;
 use error::Error;
 use std::fmt::Write;

+#[allow(dead_code)]
 pub fn disassemble(mem: &[u8]) -> Result<(), Error> {
    let mut cs = Capstone::new()
        .x86()
--- a/src/function_body.rs
+++ b/src/function_body.rs
@@ -1,34 +1,48 @@
 use backend::*;
-use disassemble::disassemble;
 use error::Error;
 use wasmparser::{FunctionBody, Operator};

-pub fn translate(body: &FunctionBody) -> Result<(), Error> {
+pub fn translate(session: &mut CodeGenSession, body: &FunctionBody) -> Result<(), Error> {
    let locals = body.get_locals_reader()?;
+
+    // Assume signature is (i32, i32) -> i32 for now.
+    // TODO: Use a real signature
+    const ARG_COUNT: u32 = 2;
+
+    let mut framesize = ARG_COUNT;
    for local in locals {
-        local?; // TODO
+        let (count, _ty) = local?;
+        framesize += count;
    }

-    let mut ops = dynasmrt::x64::Assembler::new().unwrap();
+    let mut ctx = session.new_context();
    let operators = body.get_operators_reader()?;
-    let mut regs = Registers::new();
+
+    prologue(&mut ctx, framesize);
+
+    for arg_pos in 0..ARG_COUNT {
+        copy_incoming_arg(&mut ctx, arg_pos);
+    }
+
    for op in operators {
        match op? {
            Operator::I32Add => {
-                add_i32(&mut ops, &mut regs);
+                add_i32(&mut ctx);
+            }
+            Operator::GetLocal { local_index } => {
+                get_local_i32(&mut ctx, local_index);
+            }
+            Operator::End => {
+                // TODO: This is super naive and makes a lot of unfounded assumptions 
+                // but will do for the start.
+                prepare_return_value(&mut ctx);
            }
            _ => {
-                unsupported_opcode(&mut ops);
+                unsupported_opcode(&mut ctx);
            }
        }
    }
-
-    let output = ops
-        .finalize()
-        .map_err(|_asm| Error::Assembler("assembler error".to_owned()))?;
-
-    // TODO: Do something with the output.
-    disassemble(&output)?;
+    epilogue(&mut ctx);

    Ok(())
 }
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -8,6 +8,8 @@ extern crate wasmparser;
 extern crate failure_derive;
 extern crate dynasmrt;

+extern crate wabt;
+
 mod backend;
 mod disassemble;
 mod error;
@@ -15,4 +17,8 @@ mod function_body;
 mod module;
 mod translate_sections;

+#[cfg(test)]
+mod tests;
+
 pub use module::translate;
+pub use module::TranslatedModule;
--- a/src/module.rs
+++ b/src/module.rs
@@ -1,14 +1,37 @@
+use std::mem;
 use error::Error;
 use translate_sections;
+use backend::TranslatedCodeSection;
 use wasmparser::{ModuleReader, SectionCode};

+#[derive(Default)]
+pub struct TranslatedModule {
+    translated_code_section: Option<TranslatedCodeSection>,
+}
+
+impl TranslatedModule {
+    // For testing only.
+    // Assume signature is (i32, i32) -> i32 for now.
+    // TODO: Handle generic signatures.
+    pub fn execute_func(&self, func_idx: u32, a: usize, b: usize) -> usize {
+        let code_section = self.translated_code_section.as_ref().expect("no code section");
+        let start_buf = code_section.func_start(func_idx as usize);
+
+        unsafe {
+            let func = mem::transmute::<_, extern "sysv64" fn(usize, usize) -> usize>(start_buf);
+            func(a, b)
+        }
+    }
+}
+
 /// Translate from a slice of bytes holding a wasm module.
-pub fn translate(data: &[u8]) -> Result<(), Error> {
+pub fn translate(data: &[u8]) -> Result<TranslatedModule, Error> {
    let mut reader = ModuleReader::new(data)?;
+    let mut output = TranslatedModule::default();

    reader.skip_custom_sections()?;
    if reader.eof() {
-        return Ok(());
+        return Ok(output);
    }
    let mut section = reader.read()?;

@@ -18,7 +41,7 @@ pub fn translate(data: &[u8]) -> Result<(), Error> {

        reader.skip_custom_sections()?;
        if reader.eof() {
-            return Ok(());
+            return Ok(output);
        }
        section = reader.read()?;
    }
@@ -29,7 +52,7 @@ pub fn translate(data: &[u8]) -> Result<(), Error> {

        reader.skip_custom_sections()?;
        if reader.eof() {
-            return Ok(());
+            return Ok(output);
        }
        section = reader.read()?;
    }
@@ -40,7 +63,7 @@ pub fn translate(data: &[u8]) -> Result<(), Error> {

        reader.skip_custom_sections()?;
        if reader.eof() {
-            return Ok(());
+            return Ok(output);
        }
        section = reader.read()?;
    }
@@ -51,7 +74,7 @@ pub fn translate(data: &[u8]) -> Result<(), Error> {

        reader.skip_custom_sections()?;
        if reader.eof() {
-            return Ok(());
+            return Ok(output);
        }
        section = reader.read()?;
    }
@@ -62,7 +85,7 @@ pub fn translate(data: &[u8]) -> Result<(), Error> {

        reader.skip_custom_sections()?;
        if reader.eof() {
-            return Ok(());
+            return Ok(output);
        }
        section = reader.read()?;
    }
@@ -73,7 +96,7 @@ pub fn translate(data: &[u8]) -> Result<(), Error> {

        reader.skip_custom_sections()?;
        if reader.eof() {
-            return Ok(());
+            return Ok(output);
        }
        section = reader.read()?;
    }
@@ -84,7 +107,7 @@ pub fn translate(data: &[u8]) -> Result<(), Error> {

        reader.skip_custom_sections()?;
        if reader.eof() {
-            return Ok(());
+            return Ok(output);
        }
        section = reader.read()?;
    }
@@ -95,7 +118,7 @@ pub fn translate(data: &[u8]) -> Result<(), Error> {

        reader.skip_custom_sections()?;
        if reader.eof() {
-            return Ok(());
+            return Ok(output);
        }
        section = reader.read()?;
    }
@@ -106,18 +129,18 @@ pub fn translate(data: &[u8]) -> Result<(), Error> {

        reader.skip_custom_sections()?;
        if reader.eof() {
-            return Ok(());
+            return Ok(output);
        }
        section = reader.read()?;
    }

    if let SectionCode::Code = section.code {
        let code = section.get_code_section_reader()?;
-        translate_sections::code(code)?;
+        output.translated_code_section = Some(translate_sections::code(code)?);

        reader.skip_custom_sections()?;
        if reader.eof() {
-            return Ok(());
+            return Ok(output);
        }
        section = reader.read()?;
    }
@@ -127,5 +150,5 @@ pub fn translate(data: &[u8]) -> Result<(), Error> {
        translate_sections::data(data)?;
    }

-    Ok(())
+    Ok(output)
 }
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -0,0 +1,34 @@
+use super::{translate, TranslatedModule};
+use wabt;
+
+fn translate_wat(wat: &str) -> TranslatedModule {
+    let wasm = wabt::wat2wasm(wat).unwrap();
+    let compiled = translate(&wasm).unwrap();
+    compiled
+}
+
+/// Execute the first function in the module.
+fn execute_wat(wat: &str, a: usize, b: usize) -> usize {
+    let translated = translate_wat(wat);
+    translated.execute_func(0, a, b)
+}
+
+#[test]
+fn adds() {
+    const CASES: &[(usize, usize, usize)] = &[
+        (5, 3, 8),
+        (0, 228, 228),
+        (usize::max_value(), 1, 0),
+    ];
+
+    let code = r#"
+(module
+  (func (param i32) (param i32) (result i32) (i32.add (get_local 0) (get_local 1)))
+)
+    "#;
+    for (a, b, expected) in CASES {
+        assert_eq!(execute_wat(code, *a, *b), *expected);
+    }
+}
+
+// TODO: Add a test that checks argument passing via the stack.
--- a/src/translate_sections.rs
+++ b/src/translate_sections.rs
@@ -7,6 +7,7 @@ use wasmparser::{
    GlobalSectionReader, GlobalType, Import, ImportSectionEntryType, ImportSectionReader,
    MemorySectionReader, MemoryType, Operator, TableSectionReader, Type, TypeSectionReader,
 };
+use backend::{CodeGenSession, TranslatedCodeSection};

 /// Parses the Type section of the wasm module.
 pub fn type_(types: TypeSectionReader) -> Result<(), Error> {
@@ -79,11 +80,12 @@ pub fn element(elements: ElementSectionReader) -> Result<(), Error> {
 }

 /// Parses the Code section of the wasm module.
-pub fn code(code: CodeSectionReader) -> Result<(), Error> {
+pub fn code(code: CodeSectionReader) -> Result<TranslatedCodeSection, Error> {
+    let mut session = CodeGenSession::new();
    for body in code {
-        function_body::translate(&body?)?;
+        function_body::translate(&mut session, &body?)?;
    }
-    Ok(())
+    Ok(session.into_translated_code_section()?)
 }

 /// Parses the Data section of the wasm module.