[wasm] Refactor call site patching

src/wasm/wasm-module.cc

Index: src/wasm/wasm-module.cc
diff --git a/src/wasm/wasm-module.cc b/src/wasm/wasm-module.cc
index b162704cdeb27356541d91ebd55e43d998aaf490..84e99e8e627c447619841d0a02b483320d766388 100644
--- a/src/wasm/wasm-module.cc
+++ b/src/wasm/wasm-module.cc
@@ -17,6 +17,7 @@
 #include "src/v8.h"
 
 #include "src/wasm/function-body-decoder.h"
+#include "src/wasm/leb-helper.h"
 #include "src/wasm/module-decoder.h"
 #include "src/wasm/wasm-js.h"
 #include "src/wasm/wasm-limits.h"
@@ -445,36 +446,80 @@ void CompileSequentially(Isolate* isolate, ModuleBytesEnv* module_env,
   }
 }
 
-void PatchDirectCalls(Handle<FixedArray> old_functions,
-                      Handle<FixedArray> new_functions, int start) {
-  DCHECK_EQ(new_functions->length(), old_functions->length());
+int ExtractDirectCallIndex(const byte* pc) {

[titzer, 2017/01/11 18:02:48]

Can you use a decoder for this? That way we don't need to introduce another LEB
decoder function in leb-helper.h.

[Clemens Hammacher, 2017/01/11 18:28:43]

Done.

FYI: I have to write it as wasm::Decoder, because arm64 defines another Decoder
somewhere, so it's ambiguous:
https://build.chromium.org/p/tryserver.v8/builders/v8_linux_arm64_rel_ng/buil...

+  DCHECK_EQ(static_cast<int>(kExprCallFunction), static_cast<int>(*pc));
+  uint32_t call_idx = LEBHelper::read_u32v(pc + 1);
+  DCHECK_GE(kMaxInt, call_idx);
+  return static_cast<int>(call_idx);
+}
 
+void PatchDirectCalls(Handle<FixedArray> new_functions,
+                      Handle<WasmCompiledModule> compiled_module,
+                      WasmModule* module, int start) {
   DisallowHeapAllocation no_gc;
-  std::map<Code*, Code*> old_to_new_code;
-  for (int i = 0; i < new_functions->length(); ++i) {
-    old_to_new_code.insert(std::make_pair(Code::cast(old_functions->get(i)),
-                                          Code::cast(new_functions->get(i))));
-  }
-  int mode_mask = RelocInfo::ModeMask(RelocInfo::CODE_TARGET);
   AllowDeferredHandleDereference embedding_raw_address;
-  for (int i = start; i < new_functions->length(); ++i) {
+  SeqOneByteString* module_bytes = compiled_module->module_bytes();
+  std::vector<WasmFunction>* wasm_functions =
+      &compiled_module->module()->functions;
+  DCHECK_EQ(wasm_functions->size() +
+                compiled_module->module()->num_exported_functions,
+            new_functions->length());
+  DCHECK_EQ(start, compiled_module->module()->num_imported_functions);
+
+  int num_wasm_functions = static_cast<int>(wasm_functions->size());
+  for (int i = start; i < num_wasm_functions; ++i) {
     Code* wasm_function = Code::cast(new_functions->get(i));
-    for (RelocIterator it(wasm_function, mode_mask); !it.done(); it.next()) {
-      Code* old_code =
-          Code::GetCodeFromTargetAddress(it.rinfo()->target_address());
-      if (old_code->kind() == Code::WASM_TO_JS_FUNCTION ||
-          old_code->kind() == Code::WASM_FUNCTION) {
-        auto found = old_to_new_code.find(old_code);
-        DCHECK(found != old_to_new_code.end());
-        Code* new_code = found->second;
-        if (new_code != old_code) {
-          it.rinfo()->set_target_address(new_code->instruction_start(),
-                                         UPDATE_WRITE_BARRIER,
-                                         SKIP_ICACHE_FLUSH);
-        }
-      }
+    DCHECK(wasm_function->kind() == Code::WASM_FUNCTION);
+    SourcePositionTableIterator source_pos_iterator(
+        wasm_function->source_position_table());
+    const byte* func_bytes =
+        module_bytes->GetChars() +
+        compiled_module->module()->functions[i].code_start_offset;
+    for (RelocIterator it(wasm_function, RelocInfo::kCodeTargetMask);
+         !it.done(); it.next()) {
+      Code::Kind kind =
+          Code::GetCodeFromTargetAddress(it.rinfo()->target_address())->kind();
+      if (kind != Code::WASM_FUNCTION && kind != Code::WASM_TO_JS_FUNCTION)
+        continue;
+      size_t offset_l = it.rinfo()->pc() - wasm_function->instruction_start();
+      DCHECK_GE(kMaxInt, offset_l);
+      int offset = static_cast<int>(offset_l);
+      DCHECK(!source_pos_iterator.done());
+      int byte_pos;
+      do {
+        byte_pos = source_pos_iterator.source_position().ScriptOffset();
+        source_pos_iterator.Advance();
+      } while (!source_pos_iterator.done() &&
+               source_pos_iterator.code_offset() <= offset);
+      int called_func_index = ExtractDirectCallIndex(func_bytes + byte_pos);
+      Code* new_code = Code::cast(new_functions->get(called_func_index));
+      it.rinfo()->set_target_address(new_code->instruction_start(),
+                                     UPDATE_WRITE_BARRIER, SKIP_ICACHE_FLUSH);
     }
   }
+  int func_index = num_wasm_functions;
+  for (auto exp : module->export_table) {
+    if (exp.kind != kExternalFunction) continue;
+    Code* export_wrapper = Code::cast(new_functions->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::kCodeTargetMask);
+         !it.done(); it.next()) {
+      Code::Kind kind =
+          Code::GetCodeFromTargetAddress(it.rinfo()->target_address())->kind();
+      if (kind != Code::WASM_FUNCTION && kind != Code::WASM_TO_JS_FUNCTION)
+        continue;
+      ++num_wasm_calls;
+      Code* new_code = Code::cast(new_functions->get(exp.index));
+      DCHECK_EQ(kind, new_code->kind());
+      it.rinfo()->set_target_address(new_code->instruction_start(),
+                                     UPDATE_WRITE_BARRIER, SKIP_ICACHE_FLUSH);
+    }
+    DCHECK_EQ(1, num_wasm_calls);
+    func_index++;
+  }
+  DCHECK_EQ(new_functions->length(), func_index);
 }
 
 static void ResetCompiledModule(Isolate* isolate, WasmInstanceObject* owner,
@@ -1299,7 +1344,8 @@ class WasmInstanceBuilder {
 
     if (num_imported_functions > 0 || !owner.is_null()) {
       // If the code was cloned, or new imports were compiled, patch.
-      PatchDirectCalls(old_code_table, code_table, num_imported_functions);
+      PatchDirectCalls(code_table, compiled_module_, module_,
+                       num_imported_functions);
     }
 
     FlushICache(isolate_, code_table);