[wasm] Remember import indices with the compiled module

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/base/atomic-utils.h"
 #include "src/code-stubs.h"
 
 #include "src/macro-assembler.h"
@@ skipping 148 matching lines @@
 // TODO(mtrofin): Unnecessary once we stop using JS Heap for wasm code.
 // For now, each field is expected to have the type commented by its side.
 // The elements typed as "maybe" are optional. The others are mandatory. Since
 // the compiled module is either obtained from the current v8 instance, or from
 // a snapshot produced by a compatible (==identical) v8 instance, we simply
 // fail at instantiation time, in the face of invalid data.
 enum CompiledWasmObjectFields {
   kFunctions,        // FixedArray of Code
   kImportData,       // maybe FixedArray of FixedArray respecting the
                      // WasmImportMetadata structure.
+  kImportMap,        // FixedArray. The i-th element is the Code object used for
+                     // import i
   kExports,          // maybe FixedArray of FixedArray of WasmExportMetadata
                      // structure
   kStartupFunction,  // maybe FixedArray of WasmExportMetadata structure
   kTableOfIndirectFunctionTables,  // maybe FixedArray of FixedArray of
                                    // WasmIndirectFunctionTableMetadata
   kModuleBytes,                    // maybe String
   kFunctionNameTable,              // maybe ByteArray
   kMinRequiredMemory,              // Smi. an uint32_t
   // The following 2 are either together present or absent:
   kDataSegmentsInfo,  // maybe FixedArray of FixedArray respecting the
@@ skipping 238 matching lines @@
 }
 
 void LinkModuleFunctions(Isolate* isolate,
                          std::vector<Handle<Code>>& functions) {
   for (size_t i = 0; i < functions.size(); ++i) {
     Handle<Code> code = functions[i];
     LinkFunction(code, functions, Code::WASM_FUNCTION);
   }
 }
 
-void LinkImports(Isolate* isolate, std::vector<Handle<Code>>& functions,
-                 const std::vector<Handle<Code>>& imports) {
-  for (uint32_t i = 0; i < functions.size(); ++i) {
-    Handle<Code> code = functions[i];
-    LinkFunction(code, imports, Code::WASM_TO_JS_FUNCTION);
-  }
-}
-
 void FlushAssemblyCache(Isolate* isolate, Handle<FixedArray> functions) {
   for (int i = 0; i < functions->length(); ++i) {
     Handle<Code> code = functions->GetValueChecked<Code>(isolate, i);
     Assembler::FlushICache(isolate, code->instruction_start(),
                            code->instruction_size());
   }
 }
 
 }  // namespace
 
@@ skipping 514 matching lines @@
   Handle<FixedArray> import_data;
   if (maybe_import_data.ToHandle(&import_data)) {
     if (!CompileWrappersToImportedFunctions(isolate, ffi, import_code,
                                             import_data, thrower)) {
       return false;
     }
   }
 
   RecordStats(isolate, import_code);
 
-  Handle<FixedArray> code_table = Handle<FixedArray>(
-      FixedArray::cast(instance->GetInternalField(kWasmModuleCodeTable)));
-  // TODO(mtrofin): get the code off std::vector and on FixedArray, for
-  // consistency.
-  std::vector<Handle<Code>> function_code(code_table->length());
-  for (int i = 0; i < code_table->length(); ++i) {
-    Handle<Code> code = Handle<Code>(Code::cast(code_table->get(i)));
-    function_code[i] = code;
+  if (import_code.empty()) return true;
+
+  {
+    DisallowHeapAllocation no_gc;
+    std::map<Code*, int> import_to_index;
+    Handle<FixedArray> mapping =
+        compiled_module->GetValueChecked<FixedArray>(isolate, kImportMap);
+    for (int i = 0; i < mapping->length(); ++i) {
+      import_to_index.insert(std::make_pair(Code::cast(mapping->get(i)), i));
+    }
+
+    Handle<FixedArray> code_table = Handle<FixedArray>(
+        FixedArray::cast(instance->GetInternalField(kWasmModuleCodeTable)));
+
+    int mode_mask = RelocInfo::ModeMask(RelocInfo::CODE_TARGET);
+    AllowDeferredHandleDereference embedding_raw_address;
+    for (int i = 0; i < code_table->length(); ++i) {
+      Handle<Code> wasm_function =
+          code_table->GetValueChecked<Code>(isolate, i);
+      for (RelocIterator it(*wasm_function, mode_mask); !it.done(); it.next()) {
+        Code* target =
+            Code::GetCodeFromTargetAddress(it.rinfo()->target_address());
+        if (target->kind() == Code::WASM_TO_JS_FUNCTION) {
+          auto found_index = import_to_index.find(target);
+          CHECK(found_index != import_to_index.end());
+          int index = found_index->second;
+          CHECK(index < static_cast<int>(import_code.size()));
+          Handle<Code> new_target = import_code[index];
+          it.rinfo()->set_target_address(new_target->instruction_start(),
+                                         UPDATE_WRITE_BARRIER,
+                                         SKIP_ICACHE_FLUSH);
+        }
+      }
+    }
   }
+  Handle<FixedArray> new_mapping =
+      isolate->factory()->NewFixedArray(static_cast<int>(import_code.size()));
+  for (int i = 0; i < new_mapping->length(); ++i) {
+    new_mapping->set(i, *import_code[i]);
+  }
+  compiled_module->set(kImportMap, *new_mapping);
 
-  LinkImports(isolate, function_code, import_code);
   return true;
 }
 
 bool SetupExportsObject(Handle<FixedArray> compiled_module, Isolate* isolate,
                         Handle<JSObject> instance, ErrorThrower* thrower) {
   Factory* factory = isolate->factory();
   bool mem_export =
       static_cast<bool>(Smi::cast(compiled_module->get(kExportMem))->value());
   ModuleOrigin origin = static_cast<ModuleOrigin>(
       Smi::cast(compiled_module->get(kOrigin))->value());
@@ skipping 86 matching lines @@
 
   ModuleEnv module_env;
   module_env.module = this;
   module_env.instance = &temp_instance_for_compilation;
   module_env.origin = origin;
   InitializePlaceholders(factory, &module_env.placeholders, functions.size());
 
   Handle<FixedArray> compiled_functions =
       factory->NewFixedArray(static_cast<int>(functions.size()), TENURED);
 
-  temp_instance_for_compilation.import_code.resize(import_table.size());
-  for (uint32_t i = 0; i < import_table.size(); ++i) {
-    temp_instance_for_compilation.import_code[i] =
-        CreatePlaceholder(factory, i, Code::WASM_TO_JS_FUNCTION);
+  MaybeHandle<FixedArray> maybe_imports;
+  if (import_table.size() > 0) {
+    temp_instance_for_compilation.import_code.resize(import_table.size());
+    Handle<FixedArray> imports =
+        factory->NewFixedArray(static_cast<int>(import_table.size()));
+    for (uint32_t i = 0; i < import_table.size(); ++i) {
+      Handle<Code> placeholder =
+          CreatePlaceholder(factory, i, Code::WASM_TO_JS_FUNCTION);
+      temp_instance_for_compilation.import_code[i] = placeholder;
+      imports->set(i, *placeholder);
+    }
+    maybe_imports = imports;
   }
   isolate->counters()->wasm_functions_per_module()->AddSample(
       static_cast<int>(functions.size()));
   if (FLAG_wasm_num_compilation_tasks != 0) {
     CompileInParallel(isolate, this,
                       temp_instance_for_compilation.function_code, thrower,
                       &module_env);
   } else {
     CompileSequentially(isolate, this,
                         temp_instance_for_compilation.function_code, thrower,
@@ skipping 11 matching lines @@
   // Create the compiled module object, and populate with compiled functions
   // and information needed at instantiation time. This object needs to be
   // serializable. Instantiation may occur off a deserialized version of this
   // object.
   Handle<FixedArray> ret =
       factory->NewFixedArray(kCompiledWasmObjectTableSize, TENURED);
   ret->set(kFunctions, *compiled_functions);
   if (!indirect_table.is_null()) {
     ret->set(kTableOfIndirectFunctionTables, *indirect_table.ToHandleChecked());
   }
+  if (!maybe_imports.is_null()) {
+    ret->set(kImportMap, *maybe_imports.ToHandleChecked());
+  }
   Handle<FixedArray> import_data = GetImportsMetadata(factory, this);
   ret->set(kImportData, *import_data);
 
   // Compile export functions.
   int export_size = static_cast<int>(export_table.size());
   Handle<Code> startup_fct;
   if (export_size > 0) {
     Handle<FixedArray> exports = factory->NewFixedArray(export_size, TENURED);
     for (int i = 0; i < export_size; ++i) {
       Handle<FixedArray> export_metadata =
@@ skipping 532 matching lines @@
     return static_cast<int32_t>(HeapNumber::cast(*result)->value());
   }
   thrower->Error("WASM.compileRun() failed: Return value should be number");
   return -1;
 }
 
 }  // namespace testing
 }  // namespace wasm
 }  // namespace internal
 }  // namespace v8