JIT implementation

1 files changed, 394 insertions, 0 deletions

diff --git a/weerun_parse.cpp b/weerun_parse.cpp
new file mode 100644
index 0000000..2e2a5df
--- /dev/null
+++ b/weerun_parse.cpp
@@ -0,0 +1,394 @@
+#include <algorithm>
+#include <csignal>
+#include <cstdlib>
+#include <cstring>
+#include <vector>
+
+#include "weerun.hpp"
+#include "weerun_jit.hpp"
+
+template <typename T>
+T decode(buffer_t* buf);
+
+template <typename T>
+std::vector<T> decode_vector(buffer_t* buf);
+
+void parse_custom_section(wasm_module_t* module, buffer_t* buf);
+void parse_type_section(wasm_module_t* module, buffer_t* buf);
+void parse_import_section(wasm_module_t* module, buffer_t* buf);
+void parse_function_section(wasm_module_t* module, buffer_t* buf);
+void parse_table_section(wasm_module_t* module, buffer_t* buf);
+void parse_memory_section(wasm_module_t* module, buffer_t* buf);
+void parse_global_section(wasm_module_t* module, buffer_t* buf);
+void parse_export_section(wasm_module_t* module, buffer_t* buf);
+void parse_start_section(wasm_module_t* module, buffer_t* buf);
+void parse_element_section(wasm_module_t* module, buffer_t* buf);
+void parse_code_section(wasm_module_t* module, buffer_t* buf);
+void parse_data_section(wasm_module_t* module, buffer_t* buf);
+
+extern "C" wasm_module_t* parse(buffer_t* buf) {
+  buf->ptr += 8;  // magic and version
+  wasm_module_t* module = new wasm_module_t();
+  module->num_import_funcs_ = 0;
+  module->have_start_ = false;
+  while (buf->ptr != buf->end) {
+    byte type = read_u8(buf);
+    switch (type) {
+      case 0:
+        parse_custom_section(module, buf);
+        break;
+      case WASM_SECT_TYPE:
+        parse_type_section(module, buf);
+        break;
+      case WASM_SECT_IMPORT:
+        parse_import_section(module, buf);
+        break;
+      case WASM_SECT_FUNCTION:
+        parse_function_section(module, buf);
+        break;
+      case WASM_SECT_TABLE:
+        parse_table_section(module, buf);
+        break;
+      case WASM_SECT_MEMORY:
+        parse_memory_section(module, buf);
+        break;
+      case WASM_SECT_GLOBAL:
+        parse_global_section(module, buf);
+        break;
+      case WASM_SECT_EXPORT:
+        parse_export_section(module, buf);
+        break;
+      case WASM_SECT_START:
+        parse_start_section(module, buf);
+        break;
+      case WASM_SECT_ELEMENT:
+        parse_element_section(module, buf);
+        break;
+      case WASM_SECT_CODE:
+        parse_code_section(module, buf);
+        break;
+      case WASM_SECT_DATA:
+        parse_data_section(module, buf);
+        break;
+      default:
+        assert(0);
+    }
+  }
+  return module;
+}
+
+wasm_instance_t* active_instance;
+uint32_t jit = 1;
+
+extern "C" wasm_typed_value_t run(const byte* start,
+                                  const byte* end,
+                                  uint32_t num_args,
+                                  wasm_value_t* args) {
+  buffer_t onstack_buf = {start, start, end};
+  wasm_module_t* module = parse(&onstack_buf);
+  wasm_instance_t* instance = module->create_instance();
+  wasm_value_t result;
+  if (jit != 0) {
+    result =
+        instance->run_jit(std::vector<wasm_value_t>(args, args + num_args));
+  } else {
+    result = instance->run(std::vector<wasm_value_t>(args, args + num_args));
+  }
+  wasm_typed_value_t tr = {module->funcs_[module->main_index_].sig_->ret_[0],
+                           result};
+  delete instance;
+  delete module;
+  return tr;
+}
+
+void parse_custom_section(wasm_module_t* module, buffer_t* buf) {
+  (void)module;
+  uint32_t size = read_u32leb(buf);
+  const byte* section_end = buf->ptr + size;
+  std::string name = decode<std::string>(buf);
+  buf->ptr = section_end;
+  return;
+}
+
+void parse_type_section(wasm_module_t* module, buffer_t* buf) {
+  read_u32leb(buf);
+  module->sigs_ = decode_vector<wasm_func_decl_t>(buf);
+}
+
+void parse_import_section(wasm_module_t* module, buffer_t* buf) {
+  read_u32leb(buf);
+  std::vector<wasm_import_decl_t> imports =
+      decode_vector<wasm_import_decl_t>(buf);
+  // create imported functions
+  for (wasm_import_decl_t& imp : imports) {
+    switch (imp.type_) {
+      case wasm_import_decl_t::type::FUNCTION: {
+        std::string name = imp.module_ + "." + imp.name_;
+        auto result = g_runtime_funcs.find(name);
+        if (result == g_runtime_funcs.end()) {
+          std::cout << "error: " << name << " cannot be imported";
+          throw std::runtime_error(name + " not found");
+        }
+        wasm_func_t f;
+        f.sig_ = &module->sigs_[imp.funcindex_];
+        f.typeid_ = imp.funcindex_;
+        f.run_ = result->second.run_;
+        f.run_jit_ = result->second.run_jit_;
+        module->funcs_.push_back(f);
+        module->num_import_funcs_++;
+        break;
+      }
+      default:
+        break;
+    }
+  }
+}
+
+void parse_function_section(wasm_module_t* module, buffer_t* buf) {
+  read_u32leb(buf);
+  std::vector<uint32_t> funcidxs = decode_vector<uint32_t>(buf);
+  // create function objects
+  for (uint32_t func : funcidxs) {
+    wasm_func_t f;
+    f.sig_ = &module->sigs_[func];
+    f.typeid_ = func;
+    f.run_ = nullptr;
+    module->funcs_.push_back(f);
+  }
+}
+
+void parse_table_section(wasm_module_t* module, buffer_t* buf) {
+  read_u32leb(buf);
+  std::vector<wasm_table_decl_t> tables = decode_vector<wasm_table_decl_t>(buf);
+  module->table_ = tables[0];
+}
+
+void parse_memory_section(wasm_module_t* module, buffer_t* buf) {
+  read_u32leb(buf);
+  std::vector<wasm_memory_decl_t> memories =
+      decode_vector<wasm_memory_decl_t>(buf);
+  module->mem_limits_ = memories[0];
+}
+
+void parse_global_section(wasm_module_t* module, buffer_t* buf) {
+  read_u32leb(buf);
+  module->globals_ = decode_vector<wasm_global_t>(buf);
+}
+
+void parse_export_section(wasm_module_t* module, buffer_t* buf) {
+  read_u32leb(buf);
+  std::vector<wasm_export_t> exps = decode_vector<wasm_export_t>(buf);
+  // we assume that the module only exports a main function
+  module->main_index_ = exps[0].idx_;
+}
+
+void parse_start_section(wasm_module_t* module, buffer_t* buf) {
+  read_u32leb(buf);
+  module->have_start_ = true;
+  module->start_index_ = read_u32leb(buf);
+}
+
+void parse_element_section(wasm_module_t* module, buffer_t* buf) {
+  read_u32leb(buf);
+  module->elems_ = decode_vector<wasm_elem_decl_t>(buf);
+}
+
+std::vector<wasm_code_decl_t> decode_vector_code_decl(buffer_t* buf,
+                                                      wasm_module_t* module);
+
+void parse_code_section(wasm_module_t* module, buffer_t* buf) {
+  read_u32leb(buf);
+  std::vector<wasm_code_decl_t> codes = decode_vector_code_decl(buf, module);
+  // start from the first non-imported functions
+  uint32_t funcidx = module->num_import_funcs_;
+  for (wasm_code_decl_t& code : codes) {
+    module->funcs_[funcidx].locals_ = code.locals_;
+    module->funcs_[funcidx].code_ = code.code_;
+    funcidx++;
+  }
+}
+
+void parse_data_section(wasm_module_t* module, buffer_t* buf) {
+  read_u32leb(buf);
+  module->datas_ = decode_vector<wasm_data_decl_t>(buf);
+}
+
+wasm_instance_t* wasm_module_t::create_instance() {
+  wasm_instance_t* instance = new wasm_instance_t();
+  instance->module_ = this;
+
+  // initialize memory using data
+  instance->memory_.resize(mem_limits_.init_ * wasm_page_size, 0);
+  for (wasm_data_decl_t& data : datas_) {
+    byte* mem = instance->memory_.data();
+    std::memcpy(mem + data.offset_, data.data_.data(), data.data_.size());
+  }
+
+  // initialize global variables
+  instance->globals_.resize(globals_.size());
+  for (uint32_t i = 0; i < globals_.size(); i++) {
+    switch (globals_[i].type_.type_) {
+      case WASM_TYPE_I32:
+        instance->globals_[i].val.i32 = globals_[i].val_.i32;
+        break;
+      case WASM_TYPE_F64:
+        instance->globals_[i].val.f64 = globals_[i].val_.f64;
+        break;
+      default:
+        instance->globals_[i].val.ref = globals_[i].val_.ref;
+        break;
+    }
+  }
+
+  // initialize table
+  instance->table_.resize(table_.init_, nullptr);
+  for (wasm_elem_decl_t& elem : elems_) {
+    instance->ensure_table(elem.offset_ + elem.funcidx_.size());
+    wasm_func_t** p = &instance->table_[elem.offset_];
+    for (uint32_t funcidx : elem.funcidx_) {
+      *p = &funcs_[funcidx];
+      p++;
+    }
+  }
+  instance->jit_table_.resize(instance->table_.size(), {-1, nullptr});
+  return instance;
+}
+
+void wasm_instance_t::push_value(wasm_value_t v) {
+  value_stack_.push_back(v);
+}
+
+wasm_value_t wasm_instance_t::pop_value() {
+  wasm_value_t result = value_stack_.back();
+  value_stack_.pop_back();
+  return result;
+}
+
+void wasm_instance_t::branch(uint32_t to_off) {
+  cur_frame_->pc_.ptr = cur_frame_->pc_.start + to_off;
+}
+
+void wasm_instance_t::push_frame(wasm_func_t* f) {
+  if (f->run_ != nullptr) {
+    // runtime functions
+    f->run_(this);
+    return;
+  }
+
+  // create a frame for function f
+  wasm_stack_frame_t frame;
+  frame.f_ = f;
+
+  // locals contains arguments and f's own locals
+  uint32_t num_args = f->sig_->args_.size();
+  frame.locals_.resize(num_args);
+  // move argument from stack
+  for (int i = num_args - 1; i >= 0; i--) {
+    frame.locals_[i] = pop_value();
+  }
+  // initialize locals
+  for (wasm_local_decl_t& local : f->locals_) {
+    wasm_value_t v;
+    switch (local.type_) {
+      case WASM_TYPE_I32:
+        v.val.i32 = 0;
+        break;
+      case WASM_TYPE_F64:
+        v.val.f64 = 0;
+        break;
+      case WASM_TYPE_EXTERNREF:
+        v.val.ref = nullptr;
+    }
+    for (uint32_t i = 0; i < local.count_; i++) {
+      frame.locals_.push_back(v);
+    }
+  }
+
+  // set pc to the start of function bytecode
+  frame.pc_.start = frame.pc_.ptr = f->code_.data();
+  frame.pc_.end = frame.pc_.start + f->code_.size();
+  frames_.push_back(frame);
+
+  // update current frame
+  cur_frame_ = &frames_.back();
+}
+
+void wasm_instance_t::pop_frame() {
+  // pop the from when a funtcion reaches the end
+  frames_.pop_back();
+  cur_frame_ = &frames_.back();
+}
+
+wasm_value_t wasm_instance_t::run(std::vector<wasm_value_t> args) {
+  if (module_->have_start_ == true) {
+    run_function(module_->start_index_);
+  }
+  // create the frame for main()
+  for (wasm_value_t& v : args) {
+    value_stack_.push_back(v);
+  }
+  active_instance = this;
+  run_function(module_->main_index_);
+  // return the value on stack;
+  return value_stack_.back();
+}
+
+wasm_value_t wasm_instance_t::run_jit(std::vector<wasm_value_t> args) {
+  wasm_jit_t jit(1 << 20);  // 1M code size
+  jit.compile_instance(this);
+
+  active_instance = this;
+  if (module_->have_start_ == true) {
+    jit.run_fn(module_->start_index_, nullptr, nullptr);
+  }
+  wasm_value_t r;
+  jit.run_fn(module_->main_index_, args.data(), &r);
+  return r;
+}
+
+void wasm_instance_t::run_function(uint32_t funcidx) {
+  trapped = WASM_TRAP_REASON_NOTRAP;
+  push_frame(&module_->funcs_[funcidx]);
+
+  // start interpretation
+  while (trapped == WASM_TRAP_REASON_NOTRAP) {
+    byte opcode = read_u8(&cur_frame_->pc_);
+    const wasm_inst_desc_t* desc = &g_inst_desc[opcode];
+    desc->run(this, opcode);
+  }
+
+  if (trapped != WASM_TRAP_REASON_FUNCEND) {
+    std::cout << "!trap" << std::endl;
+    std::exit(0);
+  }
+}
+
+void wasm_instance_t::trap(uint32_t reason) {
+  trapped = reason;
+}
+
+uint32_t wasm_instance_t::read_u32leb() {
+  return ::read_u32leb(&cur_frame_->pc_);
+}
+
+int32_t wasm_instance_t::read_i32leb() {
+  return ::read_i32leb(&cur_frame_->pc_);
+}
+
+double wasm_instance_t::read_double() {
+  double result = *(double*)(cur_frame_->pc_.ptr);
+  cur_frame_->pc_.ptr += sizeof(double);
+  return result;
+}
+
+byte wasm_instance_t::read_byte() {
+  byte result = *(byte*)(cur_frame_->pc_.ptr);
+  cur_frame_->pc_.ptr += sizeof(byte);
+  return result;
+}
+
+void wasm_instance_t::ensure_table(uint32_t offset) {
+  if (table_.size() < offset) {
+    table_.resize(offset, nullptr);
+  }
+}