[wasm-split] Fix instrumentation to work with memory 64 (#6013)

Correctly use the output memory's index type when generating the __write_profile
function. Requires moving some code around, but is a very small fix.

1 files changed, 23 insertions, 21 deletions

diff --git a/src/tools/wasm-split/instrumenter.cpp b/src/tools/wasm-split/instrumenter.cpp
index 38ebd271a..1c825f29f 100644
--- a/src/tools/wasm-split/instrumenter.cpp
+++ b/src/tools/wasm-split/instrumenter.cpp
@@ -177,32 +177,12 @@ void Instrumenter::instrumentFuncs() {
 // the instrumented run than funtions with larger timestamps.
 
 void Instrumenter::addProfileExport(size_t numFuncs) {
-  // Create and export a function to dump the profile into a given memory
-  // buffer. The function takes the available address and buffer size as
-  // arguments and returns the total size of the profile. It only actually
-  // writes the profile if the given space is sufficient to hold it.
-  auto name = Names::getValidFunctionName(*wasm, config.profileExport);
-  auto writeProfile = Builder::makeFunction(
-    name, Signature({Type::i32, Type::i32}, Type::i32), {});
-  writeProfile->hasExplicitName = true;
-  writeProfile->setLocalName(0, "addr");
-  writeProfile->setLocalName(1, "size");
-
   // Calculate the size of the profile:
   //   8 bytes module hash +
   //   4 bytes for the timestamp for each function
   const size_t profileSize = 8 + 4 * numFuncs;
 
-  // Create the function body
-  Builder builder(*wasm);
-  auto getAddr = [&]() { return builder.makeLocalGet(0, Type::i32); };
-  auto getSize = [&]() { return builder.makeLocalGet(1, Type::i32); };
-  auto hashConst = [&]() { return builder.makeConst(int64_t(moduleHash)); };
-  auto profileSizeConst = [&]() {
-    return builder.makeConst(int32_t(profileSize));
-  };
-
-  // Also make sure there is a memory with enough pages to write into
+  // Make sure there is a memory with enough pages to write into
   size_t pages = (profileSize + Memory::kPageSize - 1) / Memory::kPageSize;
   if (wasm->memories.empty()) {
     wasm->addMemory(Builder::makeMemory("0"));
@@ -215,6 +195,28 @@ void Instrumenter::addProfileExport(size_t numFuncs) {
     }
   }
 
+  auto ptrType = wasm->memories[0]->indexType;
+
+  // Create and export a function to dump the profile into a given memory
+  // buffer. The function takes the available address and buffer size as
+  // arguments and returns the total size of the profile. It only actually
+  // writes the profile if the given space is sufficient to hold it.
+  auto name = Names::getValidFunctionName(*wasm, config.profileExport);
+  auto writeProfile =
+    Builder::makeFunction(name, Signature({ptrType, Type::i32}, Type::i32), {});
+  writeProfile->hasExplicitName = true;
+  writeProfile->setLocalName(0, "addr");
+  writeProfile->setLocalName(1, "size");
+
+  // Create the function body
+  Builder builder(*wasm);
+  auto getAddr = [&]() { return builder.makeLocalGet(0, ptrType); };
+  auto getSize = [&]() { return builder.makeLocalGet(1, Type::i32); };
+  auto hashConst = [&]() { return builder.makeConst(int64_t(moduleHash)); };
+  auto profileSizeConst = [&]() {
+    return builder.makeConst(int32_t(profileSize));
+  };
+
   // Write the hash followed by all the time stamps
   Expression* writeData = builder.makeStore(
     8, 0, 1, getAddr(), hashConst(), Type::i64, wasm->memories[0]->name);