[wasm] Extract a CloneCodeHelper class for cloning and patching code

The refactoring extracts code responsible for cloning and patching
code when creating a NativeModule clone. This extraction makes it easier
to reuse the code.

Change-Id: Id4543ee07ed85c1b44472723b0bce983a97da03e
Reviewed-on: https://chromium-review.googlesource.com/975302
Commit-Queue: Kim-Anh Tran <kimanh@google.com>
Reviewed-by: Clemens Hammacher <clemensh@chromium.org>
Reviewed-by: Andreas Haas <ahaas@chromium.org>
Cr-Commit-Position: refs/heads/master@{#52168}

src/wasm/wasm-code-manager.cc

@@ -354,6 +354,83 @@ WasmCode::~WasmCode() {
   }
 }
 
+NativeModule::CloneCodeHelper::CloneCodeHelper(
+    NativeModule* source_native_module, NativeModule* cloning_native_module)
+    : source_native_module_(source_native_module),
+      cloning_native_module_(cloning_native_module) {
+  for (auto& pair : source_native_module_->trampolines_) {
+    Address old_dest = pair.second;
+    auto local = cloning_native_module_->trampolines_.find(pair.first);
+    DCHECK(local != cloning_native_module_->trampolines_.end());
+    Address new_dest = local->second;
+    reverse_lookup_.emplace(old_dest, new_dest);
+  }
+
+  for (auto& pair : source_native_module_->stubs_) {
+    Address old_dest = pair.second->instructions().start();
+    auto local = cloning_native_module_->stubs_.find(pair.first);
+    DCHECK(local != cloning_native_module_->stubs_.end());
+    Address new_dest = local->second->instructions().start();
+    reverse_lookup_.emplace(old_dest, new_dest);
+  }
+}
+
+void NativeModule::CloneCodeHelper::SelectForCloning(int32_t code_index) {
+  selection_.emplace_back(code_index);
+}
+
+void NativeModule::CloneCodeHelper::CloneAndPatchCode(
+    bool patch_stub_to_stub_calls) {
+  if (patch_stub_to_stub_calls) {
+    PatchStubToStubCalls();
+  }
+
+  WasmCode* anonymous_lazy_builtin = nullptr;
+  for (uint32_t index : selection_) {
+    const WasmCode* original_code = source_native_module_->GetCode(index);
+    switch (original_code->kind()) {
+      case WasmCode::kLazyStub: {
+        // Use the first anonymous lazy compile stub hit in this loop as the
+        // canonical copy for all further ones by remembering it locally via
+        // the {anonymous_lazy_builtin} variable. All non-anonymous such stubs
+        // are just cloned directly via {CloneLazyBuiltinInto} below.
+        if (!original_code->IsAnonymous()) {
+          WasmCode* new_code = cloning_native_module_->CloneCode(
+              original_code, WasmCode::kNoFlushICache);
+          PatchTrampolineAndStubCalls(original_code, new_code, reverse_lookup_,
+                                      WasmCode::kFlushICache);
+          break;
+        }
+        if (anonymous_lazy_builtin == nullptr) {
+          WasmCode* new_code = cloning_native_module_->CloneCode(
+              original_code, WasmCode::kNoFlushICache);
+          PatchTrampolineAndStubCalls(original_code, new_code, reverse_lookup_,
+                                      WasmCode::kFlushICache);
+          anonymous_lazy_builtin = new_code;
+        }
+        cloning_native_module_->code_table_[index] = anonymous_lazy_builtin;
+      } break;
+      case WasmCode::kFunction: {
+        WasmCode* new_code = cloning_native_module_->CloneCode(
+            original_code, WasmCode::kNoFlushICache);
+        PatchTrampolineAndStubCalls(original_code, new_code, reverse_lookup_,
+                                    WasmCode::kFlushICache);
+      } break;
+      default:
+        UNREACHABLE();
+    }
+  }
+}
+
+void NativeModule::CloneCodeHelper::PatchStubToStubCalls() {
+  for (auto& pair : cloning_native_module_->stubs_) {
+    WasmCode* new_stub = pair.second;
+    WasmCode* old_stub = source_native_module_->stubs_.find(pair.first)->second;
+    PatchTrampolineAndStubCalls(old_stub, new_stub, reverse_lookup_,
+                                WasmCode::kFlushICache);
+  }
+}
+
 base::AtomicNumber<size_t> NativeModule::next_id_;
 
 NativeModule::NativeModule(uint32_t num_functions, uint32_t num_imports,
@@ -960,65 +1037,13 @@ std::unique_ptr<NativeModule> NativeModule::Clone() {
   // needed yet.
   ret->CloneTrampolinesAndStubs(this, WasmCode::kNoFlushICache);
 
-  std::unordered_map<Address, Address, AddressHasher> reverse_lookup;
-  for (auto& pair : trampolines_) {
-    Address old_dest = pair.second;
-    auto local = ret->trampolines_.find(pair.first);
-    DCHECK(local != ret->trampolines_.end());
-    Address new_dest = local->second;
-    reverse_lookup.emplace(old_dest, new_dest);
-  }
-
-  for (auto& pair : stubs_) {
-    Address old_dest = pair.second->instructions().start();
-    auto local = ret->stubs_.find(pair.first);
-    DCHECK(local != ret->stubs_.end());
-    Address new_dest = local->second->instructions().start();
-    reverse_lookup.emplace(old_dest, new_dest);
-  }
-
-  for (auto& pair : ret->stubs_) {
-    WasmCode* new_stub = pair.second;
-    WasmCode* old_stub = stubs_.find(pair.first)->second;
-    PatchTrampolineAndStubCalls(old_stub, new_stub, reverse_lookup,
-                                WasmCode::kFlushICache);
-  }
-
-  WasmCode* anonymous_lazy_builtin = nullptr;
+  // Create a helper for cloning and patching code.
+  CloneCodeHelper helper(this, ret.get());
   for (uint32_t i = num_imported_functions(), e = FunctionCount(); i < e; ++i) {
-    const WasmCode* original_code = GetCode(i);
-    switch (original_code->kind()) {
-      case WasmCode::kLazyStub: {
-        // Use the first anonymous lazy compile stub hit in this loop as the
-        // canonical copy for all further ones by remembering it locally via
-        // the {anonymous_lazy_builtin} variable. All non-anonymous such stubs
-        // are just cloned directly via {CloneLazyBuiltinInto} below.
-        if (!original_code->IsAnonymous()) {
-          WasmCode* new_code =
-              ret->CloneCode(original_code, WasmCode::kNoFlushICache);
-          PatchTrampolineAndStubCalls(original_code, new_code, reverse_lookup,
-                                      WasmCode::kFlushICache);
-          break;
-        }
-        if (anonymous_lazy_builtin == nullptr) {
-          WasmCode* new_code =
-              ret->CloneCode(original_code, WasmCode::kNoFlushICache);
-          PatchTrampolineAndStubCalls(original_code, new_code, reverse_lookup,
-                                      WasmCode::kFlushICache);
-          anonymous_lazy_builtin = new_code;
-        }
-        ret->code_table_[i] = anonymous_lazy_builtin;
-      } break;
-      case WasmCode::kFunction: {
-        WasmCode* new_code =
-            ret->CloneCode(original_code, WasmCode::kNoFlushICache);
-        PatchTrampolineAndStubCalls(original_code, new_code, reverse_lookup,
-                                    WasmCode::kFlushICache);
-      } break;
-      default:
-        UNREACHABLE();
-    }
+    helper.SelectForCloning(i);
   }
+  helper.CloneAndPatchCode(true);
+
   return ret;
 }
 

src/wasm/wasm-code-manager.h

@@ -202,6 +202,26 @@ class WasmCodeManager;
 // WasmCodeManager::Commit.
 class V8_EXPORT_PRIVATE NativeModule final {
  public:
+  // Helper class to selectively clone and patch code from a
+  // {source_native_module} into a {cloning_native_module}.
+  class CloneCodeHelper {
+   public:
+    explicit CloneCodeHelper(NativeModule* source_native_module,
+                             NativeModule* cloning_native_module);
+
+    void SelectForCloning(int32_t code_index);
+
+    void CloneAndPatchCode(bool patch_stub_to_stub_calls);
+
+   private:
+    void PatchStubToStubCalls();
+
+    NativeModule* source_native_module_;
+    NativeModule* cloning_native_module_;
+    std::vector<uint32_t> selection_;
+    std::unordered_map<Address, Address, AddressHasher> reverse_lookup_;
+  };
+
   std::unique_ptr<NativeModule> Clone();
 
   WasmCode* AddCode(const CodeDesc& desc, uint32_t frame_count, uint32_t index,