[wasm] Refactor code specialization / patching

src/wasm/wasm-code-specialization.cc

Index: src/wasm/wasm-code-specialization.cc
diff --git a/src/wasm/wasm-code-specialization.cc b/src/wasm/wasm-code-specialization.cc
new file mode 100644
index 0000000000000000000000000000000000000000..9614fac601f55041e960b2b2d891f8cd17c7a6c8
--- /dev/null
+++ b/src/wasm/wasm-code-specialization.cc
@@ -0,0 +1,257 @@
+// Copyright 2017 the V8 project authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "src/wasm/wasm-code-specialization.h"
+
+#include "src/assembler-inl.h"
+#include "src/objects-inl.h"
+#include "src/source-position-table.h"
+#include "src/wasm/decoder.h"
+#include "src/wasm/wasm-module.h"
+#include "src/wasm/wasm-opcodes.h"
+
+using namespace v8::internal;
+using namespace v8::internal::wasm;
+
+namespace {
+
+int ExtractDirectCallIndex(wasm::Decoder& decoder, const byte* pc) {
+  DCHECK_EQ(static_cast<int>(kExprCallFunction), static_cast<int>(*pc));
+  decoder.Reset(pc + 1, pc + 6);
+  uint32_t call_idx = decoder.consume_u32v("call index");
+  DCHECK(decoder.ok());
+  DCHECK_GE(kMaxInt, call_idx);
+  return static_cast<int>(call_idx);
+}
+
+int AdvanceSourcePositionTableIterator(SourcePositionTableIterator& iterator,
+                                       size_t offset_l) {
+  DCHECK_GE(kMaxInt, offset_l);
+  int offset = static_cast<int>(offset_l);
+  DCHECK(!iterator.done());
+  int byte_pos;
+  do {
+    byte_pos = iterator.source_position().ScriptOffset();
+    iterator.Advance();
+  } while (!iterator.done() && iterator.code_offset() <= offset);
+  return byte_pos;
+}
+
+}  // namespace
+
+void CodeSpecialization::RelocateMemoryReferences(Address old_start,
+                                                  uint32_t old_size,
+                                                  Address new_start,
+                                                  uint32_t new_size) {
+  DCHECK(old_mem_start == 0 && new_mem_start == 0);
+  DCHECK(old_start != 0 || new_start != 0);
+  old_mem_start = old_start;
+  old_mem_size = old_size;
+  new_mem_start = new_start;
+  new_mem_size = new_size;
+}
+
+void CodeSpecialization::RelocateGlobals(Address old_start, Address new_start) {
+  DCHECK(old_globals_start == 0 && new_globals_start == 0);
+  DCHECK(old_start != 0 || new_start != 0);
+  old_globals_start = old_start;
+  new_globals_start = new_start;
+}
+
+void CodeSpecialization::PatchTableSize(uint32_t old_size, uint32_t new_size) {
+  DCHECK(old_function_table_size == 0 && new_function_table_size == 0);
+  DCHECK(old_size != 0 || new_size != 0);
+  old_function_table_size = old_size;
+  new_function_table_size = new_size;
+}
+
+void CodeSpecialization::RelocateDirectCalls(WasmInstanceObject* instance) {
+  DCHECK_NULL(relocate_direct_calls_instance);
+  DCHECK_NOT_NULL(instance);
+  relocate_direct_calls_instance = instance;
+}
+
+void CodeSpecialization::RelocateObject(Object* old_obj, Object* new_obj) {
+  DCHECK(old_obj && new_obj);
+  objects_to_relocate.insert(std::make_pair(old_obj, new_obj));
+}
+
+bool CodeSpecialization::ApplyToWholeInstance(
+    WasmInstanceObject* instance, ICacheFlushMode icache_flush_mode) {
+  WasmCompiledModule* compiled_module = instance->compiled_module();
+  FixedArray* code_table = compiled_module->ptr_to_code_table();
+  WasmModule* module = compiled_module->module();
+  std::vector<WasmFunction>* wasm_functions =
+      &compiled_module->module()->functions;
+  DCHECK_EQ(wasm_functions->size() +
+                compiled_module->module()->num_exported_functions,
+            code_table->length());
+
+  bool changed = false;
+  int func_index = module->num_imported_functions;
+
+  // Patch all wasm functions.
+  for (int num_wasm_functions = static_cast<int>(wasm_functions->size());
+       func_index < num_wasm_functions; ++func_index) {
+    Code* wasm_function = Code::cast(code_table->get(func_index));
+    changed |= ApplyToWasmCode(wasm_function, icache_flush_mode);
+  }
+
+  // Patch all exported functions.
+  for (auto exp : module->export_table) {
+    if (exp.kind != kExternalFunction) continue;
+    Code* export_wrapper = Code::cast(code_table->get(func_index));
+    DCHECK_EQ(Code::JS_TO_WASM_FUNCTION, export_wrapper->kind());
+    // There must be exactly one call to WASM_FUNCTION or WASM_TO_JS_FUNCTION.
+    int num_wasm_calls = 0;
+    for (RelocIterator it(export_wrapper,
+                          RelocInfo::ModeMask(RelocInfo::CODE_TARGET));
+         !it.done(); it.next()) {
+      DCHECK(RelocInfo::IsCodeTarget(it.rinfo()->rmode()));
+      Code* code = Code::GetCodeFromTargetAddress(it.rinfo()->target_address());
+      // Ignore calls to other builtins like ToNumber.
+      if (code->kind() != Code::WASM_FUNCTION &&
+          code->kind() != Code::WASM_TO_JS_FUNCTION &&
+          code->builtin_index() != Builtins::kIllegal)
+        continue;
+      ++num_wasm_calls;
+      Code* new_code = Code::cast(code_table->get(exp.index));
+      DCHECK(new_code->kind() == Code::WASM_FUNCTION ||
+             new_code->kind() == Code::WASM_TO_JS_FUNCTION);
+      it.rinfo()->set_target_address(new_code->instruction_start(),
+                                     UPDATE_WRITE_BARRIER, SKIP_ICACHE_FLUSH);
+      changed = true;
+    }
+    DCHECK_EQ(1, num_wasm_calls);
+    func_index++;
+  }
+  DCHECK_EQ(code_table->length(), func_index);
+  return changed;
+}
+
+bool CodeSpecialization::ApplyToWasmCode(Code* code,
+                                         ICacheFlushMode icache_flush_mode) {
+  DCHECK_EQ(Code::WASM_FUNCTION, code->kind());
+
+  bool reloc_mem = old_mem_start || new_mem_start;
+  bool reloc_globals = old_globals_start || new_globals_start;
+  bool patch_table_size = old_function_table_size || new_function_table_size;
+  bool reloc_direct_calls = relocate_direct_calls_instance != nullptr;
+  bool reloc_objects = !objects_to_relocate.empty();
+
+  int reloc_mode = 0;
+  auto add_mode = [&reloc_mode](bool cond, RelocInfo::Mode mode) {
+    if (cond) reloc_mode |= RelocInfo::ModeMask(mode);
+  };
+  add_mode(reloc_mem, RelocInfo::WASM_MEMORY_REFERENCE);
+  add_mode(reloc_mem, RelocInfo::WASM_MEMORY_SIZE_REFERENCE);
+  add_mode(reloc_globals, RelocInfo::WASM_GLOBAL_REFERENCE);
+  add_mode(patch_table_size, RelocInfo::WASM_FUNCTION_TABLE_SIZE_REFERENCE);
+  add_mode(reloc_direct_calls, RelocInfo::CODE_TARGET);
+  add_mode(reloc_objects, RelocInfo::EMBEDDED_OBJECT);
+
+  // This is a poor man's replacement for c++17's std::optional. It's lazily
+  // initialized when needed.
+  union OptionalDirectCallSitePatchingUtils {
+    char __;
+    struct {
+      char initialized;
+      SourcePositionTableIterator source_pos_it;
+      Decoder decoder;
+      const byte* func_bytes;
+    } data;
+    OptionalDirectCallSitePatchingUtils() : __(0) { data.initialized = false; }
+    ~OptionalDirectCallSitePatchingUtils() {
+      if (!data.initialized) return;
+      data.source_pos_it.~SourcePositionTableIterator();
+      data.decoder.~Decoder();
+    }
+    void Init(WasmInstanceObject* instance, Code* code) {
+      DCHECK(!data.initialized);
+      data.initialized = true;
+      new (&data.source_pos_it)
+          SourcePositionTableIterator(code->source_position_table());
+      new (&data.decoder) Decoder(nullptr, nullptr);
+      FixedArray* deopt_data = code->deoptimization_data();
+      DCHECK_EQ(2, deopt_data->length());
+      WasmCompiledModule* comp_mod = instance->compiled_module();
+      int func_index = Smi::cast(deopt_data->get(1))->value();
+      data.func_bytes =
+          comp_mod->module_bytes()->GetChars() +
+          comp_mod->module()->functions[func_index].code_start_offset;
+    }
+  } opt;
+
+  bool changed = false;
+
+  for (RelocIterator it(code, reloc_mode); !it.done(); it.next()) {
+    RelocInfo::Mode mode = it.rinfo()->rmode();
+    switch (mode) {
+      case RelocInfo::WASM_MEMORY_REFERENCE:
+      case RelocInfo::WASM_MEMORY_SIZE_REFERENCE:
+        DCHECK(reloc_mem);
+        it.rinfo()->update_wasm_memory_reference(old_mem_start, new_mem_start,
+                                                 old_mem_size, new_mem_size,
+                                                 icache_flush_mode);
+        changed = true;
+        break;
+      case RelocInfo::WASM_GLOBAL_REFERENCE:
+        DCHECK(reloc_globals);
+        it.rinfo()->update_wasm_global_reference(
+            old_globals_start, new_globals_start, icache_flush_mode);
+        changed = true;
+        break;
+      case RelocInfo::CODE_TARGET: {
+        DCHECK(reloc_direct_calls);
+        Code* old_code =
+            Code::GetCodeFromTargetAddress(it.rinfo()->target_address());
+        // Skip everything which is not a wasm call (stack checks, traps, ...).
+        if (old_code->kind() != Code::WASM_FUNCTION &&
+            old_code->kind() != Code::WASM_TO_JS_FUNCTION &&
+            old_code->builtin_index() != Builtins::kIllegal)
+          continue;
+        // Iterate simultaneously over the relocation information and the source
+        // position table. For each call in the reloc info, move the source
+        // position iterator forward to that position to find the byte offset of
+        // the respective call. Then extract the call index from the module wire
+        // bytes to find the new compiled function.
+        size_t offset = it.rinfo()->pc() - code->instruction_start();
+        if (!opt.data.initialized)
+          opt.Init(relocate_direct_calls_instance, code);
+        int byte_pos =
+            AdvanceSourcePositionTableIterator(opt.data.source_pos_it, offset);
+        int called_func_index = ExtractDirectCallIndex(
+            opt.data.decoder, opt.data.func_bytes + byte_pos);
+        FixedArray* code_table =
+            relocate_direct_calls_instance->compiled_module()
+                ->ptr_to_code_table();
+        Code* new_code = Code::cast(code_table->get(called_func_index));
+        it.rinfo()->set_target_address(new_code->instruction_start(),
+                                       UPDATE_WRITE_BARRIER, icache_flush_mode);
+        changed = true;
+      } break;
+      case RelocInfo::EMBEDDED_OBJECT: {
+        DCHECK(reloc_objects);
+        Object* old = it.rinfo()->target_object();
+        auto rel = objects_to_relocate.find(old);
+        if (rel != objects_to_relocate.end()) {
+          it.rinfo()->set_target_object(rel->second, UPDATE_WRITE_BARRIER,
+                                        icache_flush_mode);
+          changed = true;
+        }
+      } break;
+      case RelocInfo::WASM_FUNCTION_TABLE_SIZE_REFERENCE:
+        DCHECK(patch_table_size);
+        it.rinfo()->update_wasm_function_table_size_reference(
+            old_function_table_size, new_function_table_size,
+            icache_flush_mode);
+        changed = true;
+        break;
+      default:
+        UNREACHABLE();
+    }
+  }
+
+  return changed;
+}