Update InstrumentMemory pass to support memory.grow instructions

Also added instruction filter that allows to limit instrumented instructions.

Author: loganek

src/passes/InstrumentMemory.cpp

@@ -59,6 +59,7 @@
 
 #include "asmjs/shared-constants.h"
 #include "shared-constants.h"
+#include "support/string.h"
 #include <pass.h>
 #include <wasm-builder.h>
 #include <wasm.h>
@@ -93,14 +94,26 @@ static Name array_set_val_f32("array_set_val_f32");
 static Name array_set_val_f64("array_set_val_f64");
 static Name array_get_index("array_get_index");
 static Name array_set_index("array_set_index");
+static Name memory_grow_pre("memory_grow_pre");
+static Name memory_grow_post("memory_grow_post");
 
 // TODO: Add support for atomicRMW/cmpxchg
 
-struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
-  // Adds calls to new imports.
-  bool addsEffects() override { return true; }
+using InstructionFilter = std::optional<std::unordered_set<std::string>>;
+
+#define CHECK_EXPRESSION(expr)                                                 \
+  do {                                                                         \
+    if (!checkExpression(expr)) {                                              \
+      return;                                                                  \
+    }                                                                          \
+  } while (false)
 
+struct AddInstrumentation : public WalkerPass<PostWalker<AddInstrumentation>> {
+  explicit AddInstrumentation(InstructionFilter filter)
+    : _filter(std::move(filter)) {}
   void visitLoad(Load* curr) {
+    CHECK_EXPRESSION(curr);
+
     id++;
     Builder builder(*getModule());
     auto mem = getModule()->getMemory(curr->memory);
@@ -134,6 +147,8 @@ struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
   }
 
   void visitStore(Store* curr) {
+    CHECK_EXPRESSION(curr);
+
     id++;
     Builder builder(*getModule());
     auto mem = getModule()->getMemory(curr->memory);
@@ -167,6 +182,8 @@ struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
   }
 
   void visitStructGet(StructGet* curr) {
+    CHECK_EXPRESSION(curr);
+
     Builder builder(*getModule());
     Name target;
     if (curr->type == Type::i32) {
@@ -185,6 +202,8 @@ struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
   }
 
   void visitStructSet(StructSet* curr) {
+    CHECK_EXPRESSION(curr);
+
     Builder builder(*getModule());
     Name target;
     if (curr->value->type == Type::i32) {
@@ -205,6 +224,8 @@ struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
   }
 
   void visitArrayGet(ArrayGet* curr) {
+    CHECK_EXPRESSION(curr);
+
     Builder builder(*getModule());
     curr->index =
       builder.makeCall(array_get_index,
@@ -227,6 +248,8 @@ struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
   }
 
   void visitArraySet(ArraySet* curr) {
+    CHECK_EXPRESSION(curr);
+
     Builder builder(*getModule());
     curr->index =
       builder.makeCall(array_set_index,
@@ -250,10 +273,28 @@ struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
                        curr->value->type);
   }
 
+  void visitMemoryGrow(MemoryGrow* curr) {
+    CHECK_EXPRESSION(curr);
+
+    id++;
+    Builder builder(*getModule());
+    auto addressType = getModule()->getMemory(curr->memory)->addressType;
+    curr->delta =
+      builder.makeCall(memory_grow_pre,
+                       {builder.makeConst(int32_t(id)), curr->delta},
+                       addressType);
+    replaceCurrent(builder.makeCall(
+      memory_grow_post, {builder.makeConst(int32_t(id)), curr}, addressType));
+  }
+
   void visitModule(Module* curr) {
     auto addressType =
       curr->memories.empty() ? Type::i32 : curr->memories[0]->addressType;
 
+    // Grow.
+    addImport(curr, memory_grow_pre, {Type::i32, addressType}, addressType);
+    addImport(curr, memory_grow_post, {Type::i32, addressType}, addressType);
+
     // Load.
     addImport(curr,
               load_ptr,
@@ -300,14 +341,61 @@ struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
   }
 
 private:
-  Index id;
+  Index id = 0;
+  InstructionFilter _filter;
 
   void addImport(Module* curr, Name name, Type params, Type results) {
     auto import = Builder::makeFunction(name, Signature(params, results), {});
     import->module = ENV;
     import->base = name;
     curr->addFunction(std::move(import));
   }
+
+  template<typename T> bool checkExpression(T* expr) {
+    if (_filter.has_value()) {
+      return _filter->count(expressionToString(expr)) > 0;
+    }
+    return true;
+  }
+
+  std::string loadStoreTypeToString(Type type) { return type.toString(); }
+
+  std::string expressionToString(Load* expr) {
+    const auto suffix =
+      (expr->bytes == expr->type.getByteSize()
+         ? ""
+         : std::to_string(expr->bytes * 8) + "_" + (expr->signed_ ? 's' : 'u'));
+
+    return expr->type.toString() + ".load" + suffix;
+  }
+
+  std::string expressionToString(Store* expr) {
+    const auto suffix = expr->bytes == expr->valueType.getByteSize()
+                          ? ""
+                          : std::to_string(expr->bytes * 8);
+
+    return expr->valueType.toString() + ".store" + suffix;
+  }
+
+  std::string expressionToString(StructGet* expr) { return "struct.get"; }
+  std::string expressionToString(StructSet* expr) { return "struct.set"; }
+  std::string expressionToString(ArrayGet* expr) { return "array.get"; }
+  std::string expressionToString(ArraySet* expr) { return "array.set"; }
+  std::string expressionToString(MemoryGrow* expr) { return "memory.grow"; }
+};
+struct InstrumentMemory : Pass {
+  // Adds calls to new imports.
+  bool addsEffects() override { return true; }
+
+  void run(Module* module) override {
+    auto arg = getArgumentOrDefault("instrument-memory", "");
+    InstructionFilter instructions = std::nullopt;
+    if (arg.size() > 0) {
+      String::Split s(arg, ",");
+      instructions = std::unordered_set<std::string>{s.begin(), s.end()};
+    }
+    AddInstrumentation(std::move(instructions)).run(getPassRunner(), module);
+  }
 };
 
 Pass* createInstrumentMemoryPass() { return new InstrumentMemory(); }