pass support

9 files changed, 398 insertions, 44 deletions

diff --git a/src/binaryen-shell.cpp b/src/binaryen-shell.cpp
index e0f524122..4ff93cf2b 100644
--- a/src/binaryen-shell.cpp
+++ b/src/binaryen-shell.cpp
@@ -10,10 +10,15 @@
 #include "wasm-s-parser.h"
 #include "wasm-interpreter.h"
 #include "wasm-validator.h"
+#include "pass.h"
 
 using namespace cashew;
 using namespace wasm;
 
+// Globals
+
+MixedArena globalAllocator;
+
 IString ASSERT_RETURN("assert_return"),
         ASSERT_TRAP("assert_trap"),
         ASSERT_INVALID("assert_invalid"),
@@ -148,18 +153,39 @@ int main(int argc, char **argv) {
 
   char *infile = nullptr;
   bool print_before = false;
+  bool print_after = false;
+  std::vector<std::string> passes;
 
   for (size_t i = 1; i < argc; i++) {
     char* curr = argv[i];
     if (curr[0] == '-') {
       std::string arg = curr;
-      if (arg == "--print-before") {
+      if (arg == "-print-before") {
         print_before = true;
+      } else if (arg == "-print-after") {
+        print_after = true;
+      } else if (arg == "--help") {
+        std::cout << "binaryen shell\n";
+        std::cout << "--------------\n\n";
+        std::cout << "options:\n";
+        std::cout << "  -print-before : print modules before processing them\n";
+        std::cout << "\n";
+        std::cout << "passes:\n";
+        auto allPasses = PassRegistry::get()->getRegisteredNames();
+        for (auto& name : allPasses) {
+          std::cout << "  -" << name << "\n";
+        }
+        exit(0);
       } else {
-        if (infile) {
-          printf("error: unrecognized argument: %s\n", curr);
+        // otherwise, assumed to be a pass
+        const char* name = curr + 1;
+        auto check = PassRegistry::get()->createPass(name);
+        if (!check) {
+          printf("error: invalid option %s\n", curr);
           exit(1);
         }
+        delete check;
+        passes.push_back(name);
       }
     } else {
       if (infile) {
@@ -215,6 +241,22 @@ int main(int argc, char **argv) {
       std::cout << wasm;
     }
 
+    MixedArena moreModuleAllocations;
+
+    if (passes.size() > 0) {
+      if (debug) std::cerr << "running passes...\n";
+      PassRunner passRunner(&moreModuleAllocations);
+      for (auto& passName : passes) {
+        passRunner.add(passName);
+      }
+      passRunner.run(&wasm);
+    }
+
+    if (print_after) {
+      if (debug) std::cerr << "printing...\n";
+      std::cout << wasm;
+    }
+
     // run asserts
     while (i < root.size()) {
       Element& curr = *root[i];
diff --git a/src/emscripten-optimizer/colors.h b/src/emscripten-optimizer/colors.h
index 96a7baf29..1746f7e79 100644
--- a/src/emscripten-optimizer/colors.h
+++ b/src/emscripten-optimizer/colors.h
@@ -2,63 +2,59 @@
 #include <cstdlib>
 #include <ostream>
 
-struct Colors {
-  static bool use;
+namespace Colors {
+  inline bool use() {
+    return (getenv("COLORS") && getenv("COLORS")[0] == '1') || // forced
+           (isatty(STDOUT_FILENO) && (!getenv("COLORS") || getenv("COLORS")[0] != '0')); // implicit
+  }
 
-  static void normal(std::ostream& stream) {
-    if (!use) return;
+  inline void normal(std::ostream& stream) {
+    if (!use()) return;
 #if defined(__linux__) || defined(__apple__)
     stream << "\033[0m";
 #endif
   }
-  static void red(std::ostream& stream) {
-    if (!use) return;
+  inline void red(std::ostream& stream) {
+    if (!use()) return;
 #if defined(__linux__) || defined(__apple__)
     stream << "\033[31m";
 #endif
   }
-  static void magenta(std::ostream& stream) {
-    if (!use) return;
+  inline void magenta(std::ostream& stream) {
+    if (!use()) return;
 #if defined(__linux__) || defined(__apple__)
     stream << "\033[35m";
 #endif
   }
-  static void orange(std::ostream& stream) {
-    if (!use) return;
+  inline void orange(std::ostream& stream) {
+    if (!use()) return;
 #if defined(__linux__) || defined(__apple__)
     stream << "\033[33m";
 #endif
   }
-  static void grey(std::ostream& stream) {
-    if (!use) return;
+  inline void grey(std::ostream& stream) {
+    if (!use()) return;
 #if defined(__linux__) || defined(__apple__)
     stream << "\033[37m";
 #endif
   }
-  static void green(std::ostream& stream) {
-    if (!use) return;
+  inline void green(std::ostream& stream) {
+    if (!use()) return;
 #if defined(__linux__) || defined(__apple__)
     stream << "\033[32m";
 #endif
   }
-  static void blue(std::ostream& stream) {
-    if (!use) return;
+  inline void blue(std::ostream& stream) {
+    if (!use()) return;
 #if defined(__linux__) || defined(__apple__)
     stream << "\033[34m";
 #endif
   }
-  static void bold(std::ostream& stream) {
-    if (!use) return;
+  inline void bold(std::ostream& stream) {
+    if (!use()) return;
 #if defined(__linux__) || defined(__apple__)
     stream << "\033[1m";
 #endif
   }
-
-  Colors() {
-    use = (getenv("COLORS") && getenv("COLORS")[0] == '1') || // forced
-          (isatty(STDOUT_FILENO) && (!getenv("COLORS") || getenv("COLORS")[0] != '0')); // implicit
-  }
-} colors;
-
-bool Colors::use;
+};
 
diff --git a/src/mixed_arena.h b/src/mixed_arena.h
index dfd9bcb09..4d0d85c67 100644
--- a/src/mixed_arena.h
+++ b/src/mixed_arena.h
@@ -1,4 +1,7 @@
 
+#ifndef mixed_arena_h
+#define mixed_arena_h
+
 #include <vector>
 
 //
@@ -36,3 +39,7 @@ struct MixedArena {
   }
 };
 
+extern MixedArena globalAllocator;
+
+#endif // mixed_arena_h
+
diff --git a/src/pass.cpp b/src/pass.cpp
new file mode 100644
index 000000000..25e7a41e5
--- /dev/null
+++ b/src/pass.cpp
@@ -0,0 +1,75 @@
+
+#include <pass.h>
+
+namespace wasm {
+
+// PassRegistry
+
+PassRegistry* PassRegistry::get() {
+  static PassRegistry* manager = nullptr;
+  if (!manager) {
+    manager = new PassRegistry();
+  }
+  return manager;
+}
+
+void PassRegistry::registerPass(const char* name, Creator create) {
+  assert(passCreatorMap.find(name) == passCreatorMap.end());
+  passCreatorMap[name] = create;
+}
+
+Pass* PassRegistry::createPass(std::string name) {
+  Creator creator = passCreatorMap[name];
+  if (!creator) return nullptr;
+  return creator();
+}
+
+std::vector<std::string> PassRegistry::getRegisteredNames() {
+  std::vector<std::string> ret;
+  for (auto pair : passCreatorMap) {
+    ret.push_back(pair.first);
+  }
+  return ret;
+}
+
+// PassRunner
+
+void PassRunner::add(std::string passName) {
+  passes.push_back(PassRegistry::get()->createPass(passName));
+}
+
+template<class P>
+void PassRunner::add() {
+  passes.push_back(new P());
+}
+
+void PassRunner::run(Module* module) {
+  for (auto pass : passes) {
+    currPass = pass;
+    pass->run(this, module);
+  }
+}
+
+template<class P>
+P* PassRunner::getLast() {
+  bool found = false;
+  P* ret;
+  for (int i = passes.size() - 1; i >= 0; i--) {
+    if (found && (ret = dynamic_cast<P*>(passes[i]))) {
+      return ret;
+    }
+    if (passes[i] == currPass) {
+      found = true;
+    }
+  }
+  return nullptr;
+}
+
+PassRunner::~PassRunner() {
+  for (auto pass : passes) {
+    delete(pass);
+  }
+}
+
+} // namespace wasm
+
diff --git a/src/pass.h b/src/pass.h
new file mode 100644
index 000000000..638052c70
--- /dev/null
+++ b/src/pass.h
@@ -0,0 +1,110 @@
+
+#include <functional>
+
+#include "wasm.h"
+#include "mixed_arena.h"
+
+namespace wasm {
+
+class Pass;
+
+//
+// Global registry of all passes in /passes/
+//
+struct PassRegistry {
+  static PassRegistry* get();
+
+  typedef std::function<Pass* ()> Creator;
+
+  void registerPass(const char* name, Creator create);
+  Pass* createPass(std::string name);
+  std::vector<std::string> getRegisteredNames();
+
+private:
+  std::map<std::string, Creator> passCreatorMap;
+};
+
+//
+// Utility class to register a pass. See pass files for usage.
+//
+template<class P>
+struct RegisterPass {
+  RegisterPass() {
+    PassRegistry::get()->registerPass(P().name, []() {
+      return new P();
+    });
+  }
+};
+
+//
+// Runs a set of passes, in order
+//
+struct PassRunner {
+  MixedArena* allocator;
+  std::vector<Pass*> passes;
+  Pass* currPass;
+
+  PassRunner(MixedArena* allocator) : allocator(allocator) {}
+
+  void add(std::string passName);
+
+  template<class P>
+  void add();
+
+  void run(Module* module);
+
+  // Get the last pass that was already executed of a certain type.
+  template<class P>
+  P* getLast();
+
+  ~PassRunner();
+};
+
+//
+// Core pass class
+//
+struct Pass : public WasmWalker {
+  const char* name;
+
+  Pass(const char* name) : name(name) {}
+
+  // Override this to perform preparation work before the pass runs
+  virtual void prepare(PassRunner* runner, Module *module) {}
+
+  void run(PassRunner* runner, Module *module) {
+    prepare(runner, module);
+    startWalk(module);
+  }
+};
+
+// Standard passes. All passes in /passes/ are runnable from the shell,
+// but registering them here in addition allows them to communicate
+// e.g. through PassRunner::getLast
+
+// Handles names in a module, in particular adding names without duplicates
+struct NameManager : public Pass {
+  NameManager() : Pass("name-manager") {}
+
+  Name getUnique(std::string prefix);
+  // TODO: getUniqueInFunction
+
+  // visitors
+  void visitBlock(Block* curr) override;
+  void visitLoop(Loop* curr) override;
+  void visitLabel(Label* curr) override;
+  void visitBreak(Break* curr) override;
+  void visitSwitch(Switch* curr) override;
+  void visitCall(Call* curr) override;
+  void visitCallImport(CallImport* curr) override;
+  void visitFunctionType(FunctionType* curr) override;
+  void visitFunction(Function* curr) override;
+  void visitImport(Import* curr) override;
+  void visitExport(Export* curr) override;
+
+private:
+  std::set<Name> names;
+  size_t counter = 0;
+};
+
+} // namespace wasm
+
diff --git a/src/passes/LowerIfElse.cpp b/src/passes/LowerIfElse.cpp
new file mode 100644
index 000000000..ff84c6844
--- /dev/null
+++ b/src/passes/LowerIfElse.cpp
@@ -0,0 +1,51 @@
+//
+// Lowers if (x) y else z into
+//
+// L: {
+//   if (x) break (y) L
+//   z
+// }
+//
+// This is useful for investigating how beneficial if_else is.
+//
+
+#include <memory>
+
+#include <wasm.h>
+#include <pass.h>
+
+namespace wasm {
+
+struct LowerIfElse : public Pass {
+  LowerIfElse() : Pass("lower-if-else") {}
+
+  MixedArena* allocator;
+  std::unique_ptr<NameManager> namer;
+
+  void prepare(PassRunner* runner, Module *module) override {
+    allocator = runner->allocator;
+    namer = std::unique_ptr<NameManager>(new NameManager());
+    namer->run(runner, module);
+  }
+
+  void visitIf(If *curr) override {
+    if (curr->ifFalse) {
+      auto block = allocator->alloc<Block>();
+      auto name = namer->getUnique("L"); // TODO: getUniqueInFunction
+      block->name = name;
+      block->list.push_back(curr);
+      block->list.push_back(curr->ifFalse);
+      curr->ifFalse = nullptr;
+      auto break_ = allocator->alloc<Break>();
+      break_->name = name;
+      break_->value = curr->ifTrue;
+      curr->ifTrue = break_;
+      replaceCurrent(block);
+    }
+  }
+};
+
+static RegisterPass<LowerIfElse> registerPass;
+
+} // namespace wasm
+
diff --git a/src/passes/NameManager.cpp b/src/passes/NameManager.cpp
new file mode 100644
index 000000000..13a6dd9f9
--- /dev/null
+++ b/src/passes/NameManager.cpp
@@ -0,0 +1,68 @@
+//
+// NameManager
+//
+
+#include <wasm.h>
+#include <pass.h>
+
+namespace wasm {
+
+Name NameManager::getUnique(std::string prefix) {
+  while (1) {
+    Name curr = cashew::IString((prefix + std::to_string(counter++)).c_str(), false);
+    if (names.find(curr) == names.end()) {
+      names.insert(curr);
+      return curr;
+    }
+  }
+}
+
+void NameManager::visitBlock(Block* curr) {
+  names.insert(curr->name);
+}
+void NameManager::visitLoop(Loop* curr) {
+  names.insert(curr->out);
+  names.insert(curr->in);
+}
+void NameManager::visitLabel(Label* curr) {
+  names.insert(curr->name);
+}
+void NameManager::visitBreak(Break* curr) {
+  names.insert(curr->name);
+}
+void NameManager::visitSwitch(Switch* curr) {
+  names.insert(curr->name);
+  names.insert(curr->default_);
+  for (auto& target : curr->targets) {
+    names.insert(target);
+  }
+}
+void NameManager::visitCall(Call* curr) {
+  names.insert(curr->target);
+}
+void NameManager::visitCallImport(CallImport* curr) {
+  names.insert(curr->target);
+}
+void NameManager::visitFunctionType(FunctionType* curr) {
+  names.insert(curr->name);
+}
+void NameManager::visitFunction(Function* curr) {
+  names.insert(curr->name);
+  for (auto& param : curr->params) {
+    names.insert(param.name);
+  }
+  for (auto& local : curr->locals) {
+    names.insert(local.name);
+  }
+}
+void NameManager::visitImport(Import* curr) {
+  names.insert(curr->name);
+}
+void NameManager::visitExport(Export* curr) {
+  names.insert(curr->name);
+}
+
+static RegisterPass<NameManager> registerPass;
+
+} // namespace wasm
+
diff --git a/src/pretty_printing.h b/src/pretty_printing.h
index 88028fdec..7389ad97f 100644
--- a/src/pretty_printing.h
+++ b/src/pretty_printing.h
@@ -3,52 +3,55 @@
 // Pretty printing helpers
 //
 
+#ifndef _pretty_printing_h
+#define _pretty_printing_h
+
 #include <ostream>
 
 #include "emscripten-optimizer/colors.h"
 
-std::ostream &doIndent(std::ostream &o, unsigned indent) {
+inline std::ostream &doIndent(std::ostream &o, unsigned indent) {
   for (unsigned i = 0; i < indent; i++) {
     o << "  ";
   }
   return o;
 }
 
-std::ostream &incIndent(std::ostream &o, unsigned& indent) {
+inline std::ostream &incIndent(std::ostream &o, unsigned& indent) {
   o << '\n';
   indent++;
   return o; 
 }
 
-std::ostream &decIndent(std::ostream &o, unsigned& indent) {
+inline std::ostream &decIndent(std::ostream &o, unsigned& indent) {
   indent--;
   doIndent(o, indent);
   return o << ')';
 }
 
-std::ostream &prepareMajorColor(std::ostream &o) {
+inline std::ostream &prepareMajorColor(std::ostream &o) {
   Colors::red(o);
   Colors::bold(o);
   return o;
 }
 
-std::ostream &prepareColor(std::ostream &o) {
+inline std::ostream &prepareColor(std::ostream &o) {
   Colors::magenta(o);
   Colors::bold(o);
   return o;
 }
 
-std::ostream &prepareMinorColor(std::ostream &o) {
+inline std::ostream &prepareMinorColor(std::ostream &o) {
   Colors::orange(o);
   return o;
 }
 
-std::ostream &restoreNormalColor(std::ostream &o) {
+inline std::ostream &restoreNormalColor(std::ostream &o) {
   Colors::normal(o);
   return o;
 }
 
-std::ostream& printText(std::ostream &o, const char *str) {
+inline std::ostream& printText(std::ostream &o, const char *str) {
   o << '"';
   Colors::green(o);
   o << str;
@@ -56,7 +59,7 @@ std::ostream& printText(std::ostream &o, const char *str) {
   return o << '"';
 }
 
-std::ostream& printOpening(std::ostream &o, const char *str, bool major=false) {
+inline std::ostream& printOpening(std::ostream &o, const char *str, bool major=false) {
   o << '(';
   major ? prepareMajorColor(o) : prepareColor(o);
   o << str;
@@ -64,7 +67,7 @@ std::ostream& printOpening(std::ostream &o, const char *str, bool major=false) {
   return o;
 }
 
-std::ostream& printMinorOpening(std::ostream &o, const char *str) {
+inline std::ostream& printMinorOpening(std::ostream &o, const char *str) {
   o << '(';
   prepareMinorColor(o);
   o << str;
@@ -72,3 +75,5 @@ std::ostream& printMinorOpening(std::ostream &o, const char *str) {
   return o;
 }
 
+#endif //  _pretty_printing_h
+
diff --git a/src/wasm.h b/src/wasm.h
index 6fc4ca5aa..2061eb5d5 100644
--- a/src/wasm.h
+++ b/src/wasm.h
@@ -58,7 +58,7 @@ enum WasmType {
   f64
 };
 
-const char* printWasmType(WasmType type) {
+inline const char* printWasmType(WasmType type) {
   switch (type) {
     case WasmType::none: return "none";
     case WasmType::i32: return "i32";
@@ -68,7 +68,7 @@ const char* printWasmType(WasmType type) {
   }
 }
 
-unsigned getWasmTypeSize(WasmType type) {
+inline unsigned getWasmTypeSize(WasmType type) {
   switch (type) {
     case WasmType::none: abort();
     case WasmType::i32: return 4;
@@ -78,7 +78,7 @@ unsigned getWasmTypeSize(WasmType type) {
   }
 }
 
-bool isWasmTypeFloat(WasmType type) {
+inline bool isWasmTypeFloat(WasmType type) {
   switch (type) {
     case f32:
     case f64: return true;
@@ -86,7 +86,7 @@ bool isWasmTypeFloat(WasmType type) {
   }
 }
 
-WasmType getWasmType(unsigned size, bool float_) {
+inline WasmType getWasmType(unsigned size, bool float_) {
   if (size < 4) return WasmType::i32;
   if (size == 4) return float_ ? WasmType::f32 : WasmType::i32;
   if (size == 8) return float_ ? WasmType::f64 : WasmType::i64;
@@ -278,7 +278,7 @@ public:
     return _id == T()._id ? (T*)this : nullptr;
   }
 
-  std::ostream& print(std::ostream &o, unsigned indent); // avoid virtual here, for performance
+  inline std::ostream& print(std::ostream &o, unsigned indent); // avoid virtual here, for performance
 
   friend std::ostream& operator<<(std::ostream &o, Expression* expression) {
     return expression->print(o, 0);