1 files changed, 72 insertions, 23 deletions

diff --git a/src/passes/Metrics.cpp b/src/passes/Metrics.cpp
index fe14135a4..822596cb3 100644
--- a/src/passes/Metrics.cpp
+++ b/src/passes/Metrics.cpp
@@ -19,6 +19,7 @@
 #include <pass.h>
 #include <support/colors.h>
 #include <wasm.h>
+#include <wasm-binary.h>
 
 namespace wasm {
 
@@ -26,6 +27,10 @@ using namespace std;
 
 // Prints metrics between optimization passes.
 struct Metrics : public WalkerPass<PostWalker<Metrics, UnifiedExpressionVisitor<Metrics>>> {
+  bool byFunction;
+
+  Metrics(bool byFunction) : byFunction(byFunction) {}
+
   static Metrics *lastMetricsPass;
 
   map<const char *, int> counts;
@@ -35,28 +40,25 @@ struct Metrics : public WalkerPass<PostWalker<Metrics, UnifiedExpressionVisitor<
     counts[name]++;
   }
 
-  void visitModule(Module* module) {
-    ostream &o = cout;
-    o << "Counts"
-      << "\n";
-    vector<const char*> keys;
-    int total = 0;
-    for (auto i : counts) {
-      keys.push_back(i.first);
-      total += i.second;
+  void doWalkModule(Module* module) {
+    // global things
+
+    for (auto& curr : module->functionTypes) {
+      visitFunctionType(curr.get());
     }
-    // add total
-    keys.push_back("[total]");
-    counts["[total]"] = total;
-    // add vars
-    size_t vars = 0;
-    for (auto& func : module->functions) {
-      vars += func->getNumVars();
+    for (auto& curr : module->imports) {
+      visitImport(curr.get());
+    }
+    for (auto& curr : module->exports) {
+      visitExport(curr.get());
     }
-    keys.push_back("[vars]");
-    counts["[vars]"] = vars;
+    for (auto& curr : module->globals) {
+      walkGlobal(curr.get());
+    }
+    walkTable(&module->table);
+    walkMemory(&module->memory);
+
     // add functions
-    keys.push_back("[funcs]");
     counts["[funcs]"] = module->functions.size();
     // add memory and table
     if (module->memory.exists) {
@@ -64,7 +66,6 @@ struct Metrics : public WalkerPass<PostWalker<Metrics, UnifiedExpressionVisitor<
       for (auto& segment: module->memory.segments) {
         size += segment.data.size();
       }
-      keys.push_back("[memory-data]");
       counts["[memory-data]"] = size;
     }
     if (module->table.exists) {
@@ -72,13 +73,58 @@ struct Metrics : public WalkerPass<PostWalker<Metrics, UnifiedExpressionVisitor<
       for (auto& segment: module->table.segments) {
         size += segment.data.size();
       }
-      keys.push_back("[table-data]");
       counts["[table-data]"] = size;
     }
+
+    if (byFunction) {
+      // print global
+      printCounts("global");
+      // compute binary info, so we know function sizes
+      BufferWithRandomAccess buffer;
+      WasmBinaryWriter writer(module, buffer);
+      writer.write();
+      // print for each function
+      for (Index i = 0; i < module->functions.size(); i++) {
+        auto* func = module->functions[i].get();
+        counts.clear();
+        walkFunction(func);
+        counts["[vars]"] = func->getNumVars();
+        counts["[binary-bytes]"] = writer.tableOfContents.functionBodies[i].size;
+        printCounts(std::string("func: ") + func->name.str);
+      }
+      // can't comapre detailed info between passes yet
+      lastMetricsPass = nullptr;
+    } else {
+      // add function info
+      size_t vars = 0;
+      for (auto& func : module->functions) {
+        walkFunction(func.get());
+        vars += func->getNumVars();
+      }
+      counts["[vars]"] = vars;
+      // print
+      printCounts("total");
+      // compare to next time
+      lastMetricsPass = this;
+    }
+  }
+
+  void printCounts(std::string title) {
+    ostream &o = cout;
+    vector<const char*> keys;
+    // add total
+    int total = 0;
+    for (auto i : counts) {
+      keys.push_back(i.first);
+      total += i.second;
+    }
+    keys.push_back("[total]");
+    counts["[total]"] = total;
     // sort
     sort(keys.begin(), keys.end(), [](const char* a, const char* b) -> bool {
       return strcmp(b, a) > 0;
     });
+    o << title << "\n";
     for (auto* key : keys) {
       auto value = counts[key];
       o << " " << left << setw(15) << key << ": " << setw(8)
@@ -101,12 +147,15 @@ struct Metrics : public WalkerPass<PostWalker<Metrics, UnifiedExpressionVisitor<
       }
       o << "\n";
     }
-    lastMetricsPass = this;
   }
 };
 
 Pass *createMetricsPass() {
-  return new Metrics();
+  return new Metrics(false);
+}
+
+Pass *createFunctionMetricsPass() {
+  return new Metrics(true);
 }
 
 Metrics *Metrics::lastMetricsPass;