[wasm] Refactor code specialization / patching

src/wasm/wasm-module.cc

 // Copyright 2015 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 <memory>
 
 #include "src/assembler-inl.h"
 #include "src/base/adapters.h"
 #include "src/base/atomic-utils.h"
 #include "src/code-stubs.h"
 #include "src/compiler/wasm-compiler.h"
 #include "src/debug/interface-types.h"
 #include "src/objects.h"
 #include "src/property-descriptor.h"
 #include "src/simulator.h"
 #include "src/snapshot/snapshot.h"
 #include "src/v8.h"
 
 #include "src/asmjs/asm-wasm-builder.h"
 #include "src/wasm/function-body-decoder.h"
 #include "src/wasm/module-decoder.h"
+#include "src/wasm/wasm-code-specialization.h"
 #include "src/wasm/wasm-js.h"
 #include "src/wasm/wasm-limits.h"
 #include "src/wasm/wasm-module.h"
 #include "src/wasm/wasm-objects.h"
 #include "src/wasm/wasm-result.h"
 
 using namespace v8::internal;
 using namespace v8::internal::wasm;
 namespace base = v8::base;
 
 #define TRACE(...)                                      \
   do {                                                  \
     if (FLAG_trace_wasm_instances) PrintF(__VA_ARGS__); \
   } while (false)
 
 #define TRACE_CHAIN(instance)        \
   do {                               \
     instance->PrintInstancesChain(); \
   } while (false)
 
 namespace {
 
 static const int kInvalidSigIndex = -1;
 
 byte* raw_buffer_ptr(MaybeHandle<JSArrayBuffer> buffer, int offset) {
   return static_cast<byte*>(buffer.ToHandleChecked()->backing_store()) + offset;
 }
 
-void ReplaceReferenceInCode(Handle<Code> code, Handle<Object> old_ref,
-                            Handle<Object> new_ref) {
-  for (RelocIterator it(*code, 1 << RelocInfo::EMBEDDED_OBJECT); !it.done();
-       it.next()) {
-    if (it.rinfo()->target_object() == *old_ref) {
-      it.rinfo()->set_target_object(*new_ref);
-    }
-  }
-}
-
 static void MemoryFinalizer(const v8::WeakCallbackInfo<void>& data) {
   DisallowHeapAllocation no_gc;
   JSArrayBuffer** p = reinterpret_cast<JSArrayBuffer**>(data.GetParameter());
   JSArrayBuffer* buffer = *p;
 
   if (!buffer->was_neutered()) {
     void* memory = buffer->backing_store();
     DCHECK(memory != nullptr);
     base::OS::Free(memory,
                    RoundUp(kWasmMaxHeapOffset, base::OS::CommitPageSize()));
@@ skipping 45 matching lines @@
         ->AdjustAmountOfExternalAllocatedMemory(size);
 
     is_external = true;
     return memory;
   } else {
     void* memory = isolate->array_buffer_allocator()->Allocate(size);
     return memory;
   }
 }
 
-void RelocateMemoryReferencesInCode(Handle<FixedArray> code_table,
-                                    uint32_t num_imported_functions,
-                                    Address old_start, Address start,
-                                    uint32_t prev_size, uint32_t new_size) {
-  for (int i = static_cast<int>(num_imported_functions);
-       i < code_table->length(); ++i) {
-    DCHECK(code_table->get(i)->IsCode());
-    Handle<Code> code = Handle<Code>(Code::cast(code_table->get(i)));
-    AllowDeferredHandleDereference embedding_raw_address;
-    int mask = (1 << RelocInfo::WASM_MEMORY_REFERENCE) |
-               (1 << RelocInfo::WASM_MEMORY_SIZE_REFERENCE);
-    for (RelocIterator it(*code, mask); !it.done(); it.next()) {
-      it.rinfo()->update_wasm_memory_reference(old_start, start, prev_size,
-                                               new_size);
-    }
-  }
-}
-
-void RelocateGlobals(Handle<FixedArray> code_table, Address old_start,
-                     Address globals_start) {
-  for (int i = 0; i < code_table->length(); ++i) {
-    DCHECK(code_table->get(i)->IsCode());
-    Handle<Code> code = Handle<Code>(Code::cast(code_table->get(i)));
-    AllowDeferredHandleDereference embedding_raw_address;
-    int mask = 1 << RelocInfo::WASM_GLOBAL_REFERENCE;
-    for (RelocIterator it(*code, mask); !it.done(); it.next()) {
-      it.rinfo()->update_wasm_global_reference(old_start, globals_start);
-    }
-  }
-}
-
-void RelocateTableSizeReferences(Handle<FixedArray> code_table,
-                                 uint32_t old_size, uint32_t new_size) {
-  for (int i = 0; i < code_table->length(); ++i) {
-    DCHECK(code_table->get(i)->IsCode());
-    Handle<Code> code = Handle<Code>(Code::cast(code_table->get(i)));
-    AllowDeferredHandleDereference embedding_raw_address;
-    int mask = 1 << RelocInfo::WASM_FUNCTION_TABLE_SIZE_REFERENCE;
-    for (RelocIterator it(*code, mask); !it.done(); it.next()) {
-      it.rinfo()->update_wasm_function_table_size_reference(old_size, new_size);
-    }
-  }
-}
-
 void FlushICache(Isolate* isolate, Handle<FixedArray> code_table) {
   for (int i = 0; i < code_table->length(); ++i) {
     Handle<Code> code = code_table->GetValueChecked<Code>(isolate, i);
     Assembler::FlushICache(isolate, code->instruction_start(),
                            code->instruction_size());
   }
 }
 
 // Fetches the compilation unit of a wasm function and executes its parallel
 // phase.
@@ skipping 60 matching lines @@
       code->relocation_info()->length());
 }
 
 static void RecordStats(Isolate* isolate, Handle<FixedArray> functions) {
   DisallowHeapAllocation no_gc;
   for (int i = 0; i < functions->length(); ++i) {
     RecordStats(isolate, Code::cast(functions->get(i)));
   }
 }
 
-Address GetGlobalStartAddressFromCodeTemplate(Object* undefined,
-                                              JSObject* object) {
-  auto instance = WasmInstanceObject::cast(object);
-  Address old_address = nullptr;
-  if (instance->has_globals_buffer()) {
-    old_address =
-        static_cast<Address>(instance->globals_buffer()->backing_store());
-  }
-  return old_address;
-}
-
 void InitializeParallelCompilation(
     Isolate* isolate, const std::vector<WasmFunction>& functions,
     std::vector<compiler::WasmCompilationUnit*>& compilation_units,
     ModuleBytesEnv& module_env, ErrorThrower* thrower) {
   for (uint32_t i = FLAG_skip_compiling_wasm_funcs; i < functions.size(); ++i) {
     const WasmFunction* func = &functions[i];
     compilation_units[i] =
         func->imported ? nullptr : new compiler::WasmCompilationUnit(
                                        thrower, isolate, &module_env, func, i);
   }
@@ skipping 138 matching lines @@
       WasmName str = module_env->wire_bytes.GetName(&func);
       thrower->CompileError("Compilation of #%d:%.*s failed.", i, str.length(),
                             str.start());
       break;
     }
     // Install the code into the linker table.
     functions[i] = code;
   }
 }
 
-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;
-}
-
-void PatchContext(RelocIterator& it, Context* context) {
-  Object* old = it.rinfo()->target_object();
-  // The only context we use is the native context.
-  DCHECK_IMPLIES(old->IsContext(), old->IsNativeContext());
-  if (!old->IsNativeContext()) return;
-  it.rinfo()->set_target_object(context, UPDATE_WRITE_BARRIER,
-                                SKIP_ICACHE_FLUSH);
-}
-
-void PatchDirectCallsAndContext(Handle<FixedArray> new_functions,
-                                Handle<WasmCompiledModule> compiled_module,
-                                WasmModule* module, int start) {
-  DisallowHeapAllocation no_gc;
-  AllowDeferredHandleDereference embedding_raw_address;
-  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);
-  Context* context = compiled_module->ptr_to_native_context();
-  int mode_mask = RelocInfo::ModeMask(RelocInfo::CODE_TARGET) |
-                  RelocInfo::ModeMask(RelocInfo::EMBEDDED_OBJECT);
-
-  // Allocate decoder outside of the loop and reuse it to decode all function
-  // indexes.
-  wasm::Decoder decoder(nullptr, nullptr);
-  int num_wasm_functions = static_cast<int>(wasm_functions->size());
-  int func_index = start;
-  // We patch WASM_FUNCTION and WASM_TO_JS_FUNCTION during re-instantiation,
-  // and illegal builtins initially and after deserialization.
-  auto is_at_wasm_call = [](RelocIterator& it) {
-    Code* code = Code::GetCodeFromTargetAddress(it.rinfo()->target_address());
-    return code->kind() == Code::WASM_FUNCTION ||
-           code->kind() == Code::WASM_TO_JS_FUNCTION ||
-           code->builtin_index() == Builtins::kIllegal;
-  };
-
-  // Patch all wasm functions.
-  for (; func_index < num_wasm_functions; ++func_index) {
-    Code* wasm_function = Code::cast(new_functions->get(func_index));
-    DCHECK(wasm_function->kind() == Code::WASM_FUNCTION);
-    // 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.
-    SourcePositionTableIterator source_pos_iterator(
-        wasm_function->source_position_table());
-    const byte* func_bytes =
-        module_bytes->GetChars() +
-        compiled_module->module()->functions[func_index].code_start_offset;
-    for (RelocIterator it(wasm_function, mode_mask); !it.done(); it.next()) {
-      if (RelocInfo::IsEmbeddedObject(it.rinfo()->rmode())) {
-        PatchContext(it, context);
-        continue;
-      }
-      DCHECK(RelocInfo::IsCodeTarget(it.rinfo()->rmode()));
-      if (!is_at_wasm_call(it)) continue;
-      size_t offset = it.rinfo()->pc() - wasm_function->instruction_start();
-      int byte_pos =
-          AdvanceSourcePositionTableIterator(source_pos_iterator, offset);
-      int called_func_index =
-          ExtractDirectCallIndex(decoder, 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);
-    }
-  }
-  // Patch all exported 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, mode_mask); !it.done(); it.next()) {
-      if (RelocInfo::IsEmbeddedObject(it.rinfo()->rmode())) {
-        PatchContext(it, context);
-        continue;
-      }
-      DCHECK(RelocInfo::IsCodeTarget(it.rinfo()->rmode()));
-      if (!is_at_wasm_call(it)) continue;
-      ++num_wasm_calls;
-      Code* new_code = Code::cast(new_functions->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);
-    }
-    DCHECK_EQ(1, num_wasm_calls);
-    func_index++;
-  }
-  DCHECK_EQ(new_functions->length(), func_index);
-}
-
 static void ResetCompiledModule(Isolate* isolate, WasmInstanceObject* owner,
                                 WasmCompiledModule* compiled_module) {
   TRACE("Resetting %d\n", compiled_module->instance_id());
   Object* undefined = *isolate->factory()->undefined_value();
-  uint32_t old_mem_size = compiled_module->mem_size();
-  uint32_t default_mem_size = compiled_module->default_mem_size();
-  Object* mem_start = compiled_module->maybe_ptr_to_memory();
-  Address old_mem_address = nullptr;
-  Address globals_start =
-      GetGlobalStartAddressFromCodeTemplate(undefined, owner);
+  Object* fct_obj = compiled_module->ptr_to_code_table();
+  if (fct_obj != nullptr && fct_obj != undefined) {
+    uint32_t old_mem_size = compiled_module->mem_size();
+    uint32_t default_mem_size = compiled_module->default_mem_size();
+    Object* mem_start = compiled_module->maybe_ptr_to_memory();
 
-  // Reset function tables.
-  FixedArray* function_tables = nullptr;
-  FixedArray* empty_function_tables = nullptr;
-  if (compiled_module->has_function_tables()) {
-    function_tables = compiled_module->ptr_to_function_tables();
-    empty_function_tables = compiled_module->ptr_to_empty_function_tables();
-    compiled_module->set_ptr_to_function_tables(empty_function_tables);
-  }
+    // Patch code to update memory references, global references, and function
+    // table references.
+    Zone specialization_zone(isolate->allocator(), ZONE_NAME);
+    CodeSpecialization code_specialization(isolate, &specialization_zone);
 
-  if (old_mem_size > 0) {
-    CHECK_NE(mem_start, undefined);
-    old_mem_address =
-        static_cast<Address>(JSArrayBuffer::cast(mem_start)->backing_store());
-  }
-  int mode_mask = RelocInfo::ModeMask(RelocInfo::WASM_MEMORY_REFERENCE) |
-                  RelocInfo::ModeMask(RelocInfo::WASM_MEMORY_SIZE_REFERENCE) |
-                  RelocInfo::ModeMask(RelocInfo::WASM_GLOBAL_REFERENCE) |
-                  RelocInfo::ModeMask(RelocInfo::EMBEDDED_OBJECT);
+    if (old_mem_size > 0) {
+      CHECK_NE(mem_start, undefined);
+      Address old_mem_address =
+          static_cast<Address>(JSArrayBuffer::cast(mem_start)->backing_store());
+      code_specialization.RelocateMemoryReferences(
+          old_mem_address, old_mem_size, nullptr, default_mem_size);
+    }
 
-  // Patch code to update memory references, global references, and function
-  // table references.
-  Object* fct_obj = compiled_module->ptr_to_code_table();
-  if (fct_obj != nullptr && fct_obj != undefined &&
-      (old_mem_size > 0 || globals_start != nullptr || function_tables)) {
+    if (owner->has_globals_buffer()) {
+      Address globals_start =
+          static_cast<Address>(owner->globals_buffer()->backing_store());
+      code_specialization.RelocateGlobals(globals_start, nullptr);
+    }
+
+    // Reset function tables.
+    if (compiled_module->has_function_tables()) {
+      FixedArray* function_tables = compiled_module->ptr_to_function_tables();
+      FixedArray* empty_function_tables =
+          compiled_module->ptr_to_empty_function_tables();
+      DCHECK_EQ(function_tables->length(), empty_function_tables->length());
+      for (int i = 0, e = function_tables->length(); i < e; ++i) {
+        code_specialization.RelocateObject(
+            handle(function_tables->get(i), isolate),
+            handle(empty_function_tables->get(i), isolate));
+      }
+      compiled_module->set_ptr_to_function_tables(empty_function_tables);
+    }
+
     FixedArray* functions = FixedArray::cast(fct_obj);
-    for (int i = compiled_module->num_imported_functions();
-         i < functions->length(); ++i) {
+    for (int i = compiled_module->num_imported_functions(),
+             end = functions->length();
+         i < end; ++i) {
       Code* code = Code::cast(functions->get(i));
-      bool changed = false;
-      for (RelocIterator it(code, mode_mask); !it.done(); it.next()) {
-        RelocInfo::Mode mode = it.rinfo()->rmode();
-        if (RelocInfo::IsWasmMemoryReference(mode) ||
-            RelocInfo::IsWasmMemorySizeReference(mode)) {
-          it.rinfo()->update_wasm_memory_reference(
-              old_mem_address, nullptr, old_mem_size, default_mem_size);
-          changed = true;
-        } else if (RelocInfo::IsWasmGlobalReference(mode)) {
-          it.rinfo()->update_wasm_global_reference(globals_start, nullptr);
-          changed = true;
-        } else if (RelocInfo::IsEmbeddedObject(mode) && function_tables) {
-          Object* old = it.rinfo()->target_object();
-          for (int j = 0; j < function_tables->length(); ++j) {
-            if (function_tables->get(j) == old) {
-              it.rinfo()->set_target_object(empty_function_tables->get(j));
-              changed = true;
-            }
-          }
+      if (code->kind() != Code::WASM_FUNCTION) {
+        // From here on, there should only be wrappers for exported functions.
+        for (; i < end; ++i) {
+          DCHECK_EQ(Code::JS_TO_WASM_FUNCTION,
+                    Code::cast(functions->get(i))->kind());
         }
+        break;
       }
+      bool changed =
+          code_specialization.ApplyToWasmCode(code, SKIP_ICACHE_FLUSH);
+      // TODO(wasm): Check if this is faster than passing FLUSH_ICACHE_IF_NEEDED
+      // above.
       if (changed) {
         Assembler::FlushICache(isolate, code->instruction_start(),
                                code->instruction_size());
       }
     }
   }
   compiled_module->reset_memory();
 }
 
 static void MemoryInstanceFinalizer(Isolate* isolate,
@@ skipping 144 matching lines @@
 
   int name_chars = SNPrintF(ArrayVector(buffer), "wasm-%08x", hash);
   DCHECK(name_chars >= 0 && name_chars < kBufferSize);
   MaybeHandle<String> name_str = isolate->factory()->NewStringFromOneByte(
       Vector<const uint8_t>(reinterpret_cast<uint8_t*>(buffer), name_chars),
       TENURED);
   script->set_name(*name_str.ToHandleChecked());
 
   return script;
 }
+
 }  // namespace
 
 Handle<JSArrayBuffer> SetupArrayBuffer(Isolate* isolate, void* backing_store,
                                        size_t size, bool is_external,
                                        bool enable_guard_regions) {
   Handle<JSArrayBuffer> buffer = isolate->factory()->NewJSArrayBuffer();
   JSArrayBuffer::Setup(buffer, isolate, is_external, backing_store,
                        static_cast<int>(size));
   buffer->set_is_neuterable(false);
   buffer->set_has_guard_region(enable_guard_regions);
@@ skipping 466 matching lines @@
         TRACE("Reusing existing instance %d\n",
               compiled_module_->instance_id());
       }
       compiled_module_->set_code_table(code_table);
       compiled_module_->set_native_context(isolate_->native_context());
     }
 
     //--------------------------------------------------------------------------
     // Allocate the instance object.
     //--------------------------------------------------------------------------
+    Zone instantiation_zone(isolate_->allocator(), ZONE_NAME);
+    CodeSpecialization code_specialization(isolate_, &instantiation_zone);
     Handle<WasmInstanceObject> instance =
         WasmInstanceObject::New(isolate_, compiled_module_);
 
     //--------------------------------------------------------------------------
     // Set up the globals for the new instance.
     //--------------------------------------------------------------------------
     MaybeHandle<JSArrayBuffer> old_globals;
     uint32_t globals_size = module_->globals_size;
     if (globals_size > 0) {
       const bool enable_guard_regions = false;
       Handle<JSArrayBuffer> global_buffer =
           NewArrayBuffer(isolate_, globals_size, enable_guard_regions);
       globals_ = global_buffer;
       if (globals_.is_null()) {
         thrower_->RangeError("Out of memory: wasm globals");
         return nothing;
       }
-      Address old_address =
-          owner.is_null() ? nullptr : GetGlobalStartAddressFromCodeTemplate(
-                                          isolate_->heap()->undefined_value(),
-                                          *owner.ToHandleChecked());
-      RelocateGlobals(code_table, old_address,
-                      static_cast<Address>(global_buffer->backing_store()));
+      Address old_globals_start = nullptr;
+      if (!owner.is_null()) {
+        DCHECK(owner.ToHandleChecked()->has_globals_buffer());
+        old_globals_start = static_cast<Address>(
+            owner.ToHandleChecked()->globals_buffer()->backing_store());
+      }
+      Address new_globals_start =
+          static_cast<Address>(global_buffer->backing_store());
+      code_specialization.RelocateGlobals(old_globals_start, new_globals_start);
       instance->set_globals_buffer(*global_buffer);
     }
 
     //--------------------------------------------------------------------------
     // Prepare for initialization of function tables.
     //--------------------------------------------------------------------------
     int function_table_count =
         static_cast<int>(module_->function_tables.size());
     table_instances_.reserve(module_->function_tables.size());
     for (int index = 0; index < function_table_count; ++index) {
       table_instances_.push_back(
           {Handle<WasmTableObject>::null(), Handle<FixedArray>::null(),
            Handle<FixedArray>::null(), Handle<FixedArray>::null()});
     }
 
     //--------------------------------------------------------------------------
     // Process the imports for the module.
     //--------------------------------------------------------------------------
     int num_imported_functions = ProcessImports(code_table, instance);
     if (num_imported_functions < 0) return nothing;
 
     //--------------------------------------------------------------------------
     // Process the initialization for the module's globals.
     //--------------------------------------------------------------------------
     InitGlobals();
 
     //--------------------------------------------------------------------------
     // Set up the indirect function tables for the new instance.
     //--------------------------------------------------------------------------
-    if (function_table_count > 0) InitializeTables(code_table, instance);
+    if (function_table_count > 0)
+      InitializeTables(code_table, instance, &code_specialization);
 
     //--------------------------------------------------------------------------
     // Set up the memory for the new instance.
     //--------------------------------------------------------------------------
     MaybeHandle<JSArrayBuffer> old_memory;
 
     uint32_t min_mem_pages = module_->min_mem_pages;
     isolate_->counters()->wasm_min_mem_pages_count()->AddSample(min_mem_pages);
 
     if (!memory_.is_null()) {
@@ skipping 44 matching lines @@
       uint32_t mem_size =
           static_cast<uint32_t>(memory_->byte_length()->Number());
       LoadDataSegments(mem_start, mem_size);
 
       uint32_t old_mem_size = compiled_module_->mem_size();
       Address old_mem_start =
           compiled_module_->has_memory()
               ? static_cast<Address>(
                     compiled_module_->memory()->backing_store())
               : nullptr;
-      RelocateMemoryReferencesInCode(
-          code_table, module_->num_imported_functions, old_mem_start, mem_start,
-          old_mem_size, mem_size);
+      code_specialization.RelocateMemoryReferences(old_mem_start, old_mem_size,
+                                                   mem_start, mem_size);
       compiled_module_->set_memory(memory_);
     }
 
     //--------------------------------------------------------------------------
     // Set up the runtime support for the new instance.
     //--------------------------------------------------------------------------
     Handle<WeakCell> weak_link = factory->NewWeakCell(instance);
 
     for (int i = num_imported_functions + FLAG_skip_compiling_wasm_funcs;
          i < code_table->length(); ++i) {
@@ skipping 18 matching lines @@
     DCHECK(wasm::IsWasmInstance(*instance));
     if (instance->has_memory_object()) {
       instance->memory_object()->AddInstance(isolate_, instance);
     }
 
     //--------------------------------------------------------------------------
     // Initialize the indirect function tables.
     //--------------------------------------------------------------------------
     if (function_table_count > 0) LoadTableSegments(code_table, instance);
 
-    // Patch new call sites and the context.
-    PatchDirectCallsAndContext(code_table, compiled_module_, module_,
-                               num_imported_functions);
+    // Patch all code with the relocations registered in code_specialization.
+    {
+      code_specialization.RelocateDirectCalls(instance);
+      code_specialization.ApplyToWholeInstance(*instance, SKIP_ICACHE_FLUSH);
+    }
 
     FlushICache(isolate_, code_table);
 
     //--------------------------------------------------------------------------
     // Unpack and notify signal handler of protected instructions.
     //--------------------------------------------------------------------------
     if (FLAG_wasm_trap_handler) {
       for (int i = 0; i < code_table->length(); ++i) {
         Handle<Code> code = code_table->GetValueChecked<Code>(isolate_, i);
 
@@ skipping 674 matching lines @@
 
     if (module_->origin == kWasmOrigin) {
       v8::Maybe<bool> success = JSReceiver::SetIntegrityLevel(
           exports_object, FROZEN, Object::DONT_THROW);
       DCHECK(success.FromMaybe(false));
       USE(success);
     }
   }
 
   void InitializeTables(Handle<FixedArray> code_table,
-                        Handle<WasmInstanceObject> instance) {
+                        Handle<WasmInstanceObject> instance,
+                        CodeSpecialization* code_specialization) {
     int function_table_count =
         static_cast<int>(module_->function_tables.size());
     Handle<FixedArray> new_function_tables =
         isolate_->factory()->NewFixedArray(function_table_count);
     Handle<FixedArray> new_signature_tables =
         isolate_->factory()->NewFixedArray(function_table_count);
     for (int index = 0; index < function_table_count; ++index) {
       WasmIndirectFunctionTable& table = module_->function_tables[index];
       TableInstance& table_instance = table_instances_[index];
       int table_size = static_cast<int>(table.min_size);
 
       if (table_instance.function_table.is_null()) {
         // Create a new dispatch table if necessary.
         table_instance.function_table =
             isolate_->factory()->NewFixedArray(table_size);
         table_instance.signature_table =
             isolate_->factory()->NewFixedArray(table_size);
         for (int i = 0; i < table_size; ++i) {
           // Fill the table with invalid signature indexes so that
           // uninitialized entries will always fail the signature check.
           table_instance.signature_table->set(i,
                                               Smi::FromInt(kInvalidSigIndex));
         }
       } else {
         // Table is imported, patch table bounds check
         DCHECK(table_size <= table_instance.function_table->length());
         if (table_size < table_instance.function_table->length()) {
-          RelocateTableSizeReferences(code_table, table_size,
-                                      table_instance.function_table->length());
+          code_specialization->PatchTableSize(
+              table_size, table_instance.function_table->length());
         }
       }
 
       new_function_tables->set(static_cast<int>(index),
                                *table_instance.function_table);
       new_signature_tables->set(static_cast<int>(index),
                                 *table_instance.signature_table);
     }
 
-    // Patch all code that has references to the old indirect tables.
-    Handle<FixedArray> old_function_tables =
-        compiled_module_->function_tables();
-    Handle<FixedArray> old_signature_tables =
-        compiled_module_->signature_tables();
-    for (int i = 0; i < code_table->length(); ++i) {
-      if (!code_table->get(i)->IsCode()) continue;
-      Handle<Code> code(Code::cast(code_table->get(i)), isolate_);
-      for (int j = 0; j < function_table_count; ++j) {
-        ReplaceReferenceInCode(
-            code, Handle<Object>(old_function_tables->get(j), isolate_),
-            Handle<Object>(new_function_tables->get(j), isolate_));
-        ReplaceReferenceInCode(
-            code, Handle<Object>(old_signature_tables->get(j), isolate_),
-            Handle<Object>(new_signature_tables->get(j), isolate_));
-      }
+    FixedArray* old_function_tables =
+        compiled_module_->ptr_to_function_tables();
+    DCHECK_EQ(old_function_tables->length(), new_function_tables->length());
+    for (int i = 0, e = new_function_tables->length(); i < e; ++i) {
+      code_specialization->RelocateObject(
+          handle(old_function_tables->get(i), isolate_),
+          handle(new_function_tables->get(i), isolate_));
     }
+    FixedArray* old_signature_tables =
+        compiled_module_->ptr_to_signature_tables();
+    DCHECK_EQ(old_signature_tables->length(), new_signature_tables->length());
+    for (int i = 0, e = new_signature_tables->length(); i < e; ++i) {
+      code_specialization->RelocateObject(
+          handle(old_signature_tables->get(i), isolate_),
+          handle(new_signature_tables->get(i), isolate_));
+    }
+
     compiled_module_->set_function_tables(new_function_tables);
     compiled_module_->set_signature_tables(new_signature_tables);
   }
 
   void LoadTableSegments(Handle<FixedArray> code_table,
                          Handle<WasmInstanceObject> instance) {
     int function_table_count =
         static_cast<int>(module_->function_tables.size());
     for (int index = 0; index < function_table_count; ++index) {
       WasmIndirectFunctionTable& table = module_->function_tables[index];
@@ skipping 169 matching lines @@
   if (old_size != 0) {
     memcpy(new_mem_start, old_mem_start, old_size);
   }
   return new_buffer;
 }
 
 void UncheckedUpdateInstanceMemory(Isolate* isolate,
                                    Handle<WasmInstanceObject> instance,
                                    Address old_mem_start, uint32_t old_size) {
   DCHECK(instance->has_memory_buffer());
-  Handle<JSArrayBuffer> new_buffer(instance->memory_buffer());
-  uint32_t new_size = new_buffer->byte_length()->Number();
-  DCHECK(new_size <= std::numeric_limits<uint32_t>::max());
-  Address new_mem_start = static_cast<Address>(new_buffer->backing_store());
+  Handle<JSArrayBuffer> mem_buffer(instance->memory_buffer());
+  uint32_t new_size = mem_buffer->byte_length()->Number();
+  Address new_mem_start = static_cast<Address>(mem_buffer->backing_store());
   DCHECK_NOT_NULL(new_mem_start);
-  Handle<FixedArray> code_table = instance->compiled_module()->code_table();
-  RelocateMemoryReferencesInCode(
-      code_table, instance->compiled_module()->module()->num_imported_functions,
-      old_mem_start, new_mem_start, old_size, new_size);
+  Zone specialization_zone(isolate->allocator(), ZONE_NAME);
+  CodeSpecialization code_specialization(isolate, &specialization_zone);
+  code_specialization.RelocateMemoryReferences(old_mem_start, old_size,
+                                               new_mem_start, new_size);
+  code_specialization.ApplyToWholeInstance(*instance);
 }
 
 void DetachArrayBuffer(Isolate* isolate, Handle<JSArrayBuffer> buffer) {
   const bool has_guard_regions =
       (!buffer.is_null() && buffer->has_guard_region());
   void* backing_store = buffer->backing_store();
   if (backing_store != nullptr) {
     DCHECK(!buffer->is_neuterable());
     int64_t byte_length = NumberToSize(buffer->byte_length());
     buffer->set_is_neuterable(true);
@@ skipping 117 matching lines @@
   } else {
     return GrowWebAssemblyMemory(isolate, handle(instance_obj->memory_object()),
                                  pages);
   }
 }
 
 void wasm::GrowDispatchTables(Isolate* isolate,
                               Handle<FixedArray> dispatch_tables,
                               uint32_t old_size, uint32_t count) {
   DCHECK_EQ(0, dispatch_tables->length() % 4);
+
+  Zone specialization_zone(isolate->allocator(), ZONE_NAME);
   for (int i = 0; i < dispatch_tables->length(); i += 4) {
     Handle<FixedArray> old_function_table(
         FixedArray::cast(dispatch_tables->get(i + 2)));
     Handle<FixedArray> old_signature_table(
         FixedArray::cast(dispatch_tables->get(i + 3)));
     Handle<FixedArray> new_function_table =
         isolate->factory()->CopyFixedArrayAndGrow(old_function_table, count);
     Handle<FixedArray> new_signature_table =
         isolate->factory()->CopyFixedArrayAndGrow(old_signature_table, count);
 
-    // Get code table for the instance
-    Handle<WasmInstanceObject> instance(
-        WasmInstanceObject::cast(dispatch_tables->get(i)));
-    Handle<FixedArray> code_table(instance->compiled_module()->code_table());
-
-    // Relocate size references
-    RelocateTableSizeReferences(code_table, old_size, old_size + count);
-
-    // Replace references of old tables with new tables.
-    for (int j = 0; j < code_table->length(); ++j) {
-      if (!code_table->get(j)->IsCode()) continue;
-      Handle<Code> code = Handle<Code>(Code::cast(code_table->get(j)));
-      ReplaceReferenceInCode(code, old_function_table, new_function_table);
-      ReplaceReferenceInCode(code, old_signature_table, new_signature_table);
-    }
-
     // Update dispatch tables with new function/signature tables
     dispatch_tables->set(i + 2, *new_function_table);
     dispatch_tables->set(i + 3, *new_signature_table);
+
+    // Patch the code of the respective instance.
+    CodeSpecialization code_specialization(isolate, &specialization_zone);
+    code_specialization.PatchTableSize(old_size, old_size + count);
+    code_specialization.RelocateObject(old_function_table, new_function_table);
+    code_specialization.RelocateObject(old_signature_table,
+                                       new_signature_table);
+    code_specialization.ApplyToWholeInstance(
+        WasmInstanceObject::cast(dispatch_tables->get(i)));
   }
 }
 
 void testing::ValidateInstancesChain(Isolate* isolate,
                                      Handle<WasmModuleObject> module_obj,
                                      int instance_count) {
   CHECK_GE(instance_count, 0);
   DisallowHeapAllocation no_gc;
   WasmCompiledModule* compiled_module = module_obj->compiled_module();
   CHECK_EQ(JSObject::cast(compiled_module->ptr_to_weak_wasm_module()->value()),
@@ skipping 369 matching lines @@
   Handle<String> module_property_name =
       isolate->factory()->InternalizeUtf8String("module");
   Handle<String> instance_property_name =
       isolate->factory()->InternalizeUtf8String("instance");
   JSObject::AddProperty(ret, module_property_name, module, NONE);
   JSObject::AddProperty(ret, instance_property_name,
                         instance_object.ToHandleChecked(), NONE);
 
   ResolvePromise(isolate, promise, ret);
 }