1 files changed, 197 insertions, 2 deletions

diff --git a/src/tools/wasm-split.cpp b/src/tools/wasm-split.cpp
index 5f6207e9a..61f52d5fa 100644
--- a/src/tools/wasm-split.cpp
+++ b/src/tools/wasm-split.cpp
@@ -19,10 +19,13 @@
 
 #include "ir/module-splitting.h"
 #include "ir/module-utils.h"
+#include "ir/names.h"
 #include "support/name.h"
 #include "support/utilities.h"
 #include "tool-options.h"
+#include "wasm-builder.h"
 #include "wasm-io.h"
+#include "wasm-type.h"
 #include "wasm-validator.h"
 #include <sstream>
 
@@ -43,6 +46,7 @@ std::set<Name> parseNameList(const std::string& list) {
 
 struct WasmSplitOptions : ToolOptions {
   bool verbose = false;
+  bool emitBinary = true;
 
   bool instrument = false;
 
@@ -162,6 +166,18 @@ WasmSplitOptions::WasmSplitOptions()
            verbose = true;
            quiet = false;
          })
+    .add("--emit-text",
+         "-S",
+         "Emit text instead of binary for the output file or files.",
+         Options::Arguments::Zero,
+         [&](Options* o, const std::string& argument) { emitBinary = false; })
+    .add("--debuginfo",
+         "-g",
+         "Emit names section in wasm binary (or full debuginfo in wast)",
+         Options::Arguments::Zero,
+         [&](Options* o, const std::string& arguments) {
+           passOptions.debugInfo = true;
+         })
     .add_positional(
       "INFILE",
       Options::Arguments::One,
@@ -245,8 +261,186 @@ void parseInput(Module& wasm, const WasmSplitOptions& options) {
   options.applyFeatures(wasm);
 }
 
+// Add a global monotonic counter and a timestamp global for each function, code
+// at the beginning of each function to set its timestamp, and a new exported
+// function for dumping the profile data.
+struct Instrumenter : public Pass {
+  PassRunner* runner = nullptr;
+  Module* wasm = nullptr;
+
+  const std::string& profileExport;
+  uint64_t moduleHash;
+
+  Name counterGlobal;
+  std::vector<Name> functionGlobals;
+
+  Instrumenter(const std::string& profileExport, uint64_t moduleHash);
+
+  void run(PassRunner* runner, Module* wasm) override;
+  void addGlobals();
+  void instrumentFuncs();
+  void addProfileExport();
+};
+
+Instrumenter::Instrumenter(const std::string& profileExport,
+                           uint64_t moduleHash)
+  : profileExport(profileExport), moduleHash(moduleHash) {}
+
+void Instrumenter::run(PassRunner* runner, Module* wasm) {
+  this->runner = runner;
+  this->wasm = wasm;
+  addGlobals();
+  instrumentFuncs();
+  addProfileExport();
+}
+
+void Instrumenter::addGlobals() {
+  // Create fresh global names (over-reserves, but that's ok)
+  counterGlobal = Names::getValidGlobalName(*wasm, "monotonic_counter");
+  functionGlobals.reserve(wasm->functions.size());
+  ModuleUtils::iterDefinedFunctions(*wasm, [&](Function* func) {
+    functionGlobals.push_back(Names::getValidGlobalName(
+      *wasm, std::string(func->name.c_str()) + "_timestamp"));
+  });
+
+  // Create and add new globals
+  auto addGlobal = [&](Name name) {
+    auto global = Builder::makeGlobal(
+      name,
+      Type::i32,
+      Builder(*wasm).makeConst(Literal::makeZero(Type::i32)),
+      Builder::Mutable);
+    global->hasExplicitName = true;
+    wasm->addGlobal(std::move(global));
+  };
+  addGlobal(counterGlobal);
+  for (auto& name : functionGlobals) {
+    addGlobal(name);
+  }
+}
+
+void Instrumenter::instrumentFuncs() {
+  // Inject the following code at the beginning of each function to advance the
+  // monotonic counter and set the function's timestamp if it hasn't already
+  // been set.
+  //
+  //   (if (i32.eqz (global.get $timestamp))
+  //     (block
+  //       (global.set $monotonic_counter
+  //         (i32.add
+  //           (global.get $monotonic_counter)
+  //           (i32.const 1)
+  //         )
+  //       )
+  //       (global.set $timestamp
+  //         (global.get $monotonic_counter)
+  //       )
+  //     )
+  //   )
+  Builder builder(*wasm);
+  auto globalIt = functionGlobals.begin();
+  ModuleUtils::iterDefinedFunctions(*wasm, [&](Function* func) {
+    func->body = builder.makeSequence(
+      builder.makeIf(
+        builder.makeUnary(EqZInt32,
+                          builder.makeGlobalGet(*globalIt, Type::i32)),
+        builder.makeSequence(
+          builder.makeGlobalSet(
+            counterGlobal,
+            builder.makeBinary(AddInt32,
+                               builder.makeGlobalGet(counterGlobal, Type::i32),
+                               builder.makeConst(Literal::makeOne(Type::i32)))),
+          builder.makeGlobalSet(
+            *globalIt, builder.makeGlobalGet(counterGlobal, Type::i32)))),
+      func->body,
+      func->body->type);
+    ++globalIt;
+  });
+}
+
+void Instrumenter::addProfileExport() {
+  // 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, 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 * functionGlobals.size();
+
+  // 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));
+  };
+
+  // Write the hash followed by all the time stamps
+  Expression* writeData =
+    builder.makeStore(8, 0, 1, getAddr(), hashConst(), Type::i64);
+
+  uint32_t offset = 8;
+  for (const auto& global : functionGlobals) {
+    writeData = builder.blockify(
+      writeData,
+      builder.makeStore(4,
+                        offset,
+                        1,
+                        getAddr(),
+                        builder.makeGlobalGet(global, Type::i32),
+                        Type::i32));
+    offset += 4;
+  }
+
+  writeProfile->body = builder.makeSequence(
+    builder.makeIf(builder.makeBinary(GeUInt32, getSize(), profileSizeConst()),
+                   writeData),
+    profileSizeConst());
+
+  // Create an export for the function
+  wasm->addFunction(std::move(writeProfile));
+  wasm->addExport(
+    Builder::makeExport(profileExport, name, ExternalKind::Function));
+
+  // Also make sure there is a memory with enough pages to write into
+  size_t pages = (profileSize + Memory::kPageSize - 1) / Memory::kPageSize;
+  if (!wasm->memory.exists) {
+    wasm->memory.exists = true;
+    wasm->memory.initial = pages;
+    wasm->memory.max = pages;
+  } else if (wasm->memory.initial < pages) {
+    wasm->memory.initial = pages;
+    if (wasm->memory.max < pages) {
+      wasm->memory.max = pages;
+    }
+  }
+}
+
 void instrumentModule(Module& wasm, const WasmSplitOptions& options) {
-  Fatal() << "TODO: implement instrumentation\n";
+  // Check that the profile export name is not already taken
+  if (wasm.getExportOrNull(options.profileExport) != nullptr) {
+    Fatal() << "error: Export " << options.profileExport << " already exists.";
+  }
+
+  // TODO: calculate module hash.
+  uint64_t moduleHash = 0;
+  PassRunner runner(&wasm, options.passOptions);
+  Instrumenter(options.profileExport, moduleHash).run(&runner, &wasm);
+
+  // Write the output modules
+  ModuleWriter writer;
+  writer.setBinary(options.emitBinary);
+  writer.setDebugInfo(options.passOptions.debugInfo);
+  writer.write(wasm, options.output);
 }
 
 void splitModule(Module& wasm, const WasmSplitOptions& options) {
@@ -337,7 +531,8 @@ void splitModule(Module& wasm, const WasmSplitOptions& options) {
 
   // Write the output modules
   ModuleWriter writer;
-  writer.setBinary(true);
+  writer.setBinary(options.emitBinary);
+  writer.setDebugInfo(options.passOptions.debugInfo);
   writer.write(wasm, options.primaryOutput);
   writer.write(*secondary, options.secondaryOutput);
 }