Optimize/merge duplicate function types (#1041)

17 files changed, 156 insertions, 53 deletions

diff --git a/src/passes/RemoveUnusedModuleElements.cpp b/src/passes/RemoveUnusedModuleElements.cpp
index 392096b72..259c5b942 100644
--- a/src/passes/RemoveUnusedModuleElements.cpp
+++ b/src/passes/RemoveUnusedModuleElements.cpp
@@ -16,7 +16,8 @@
 
 //
 // Removes module elements that are are never used: functions and globals,
-// which may be imported or not.
+// which may be imported or not, and function types (which we merge
+// and remove if unneeded)
 //
 
 
@@ -25,6 +26,7 @@
 #include "wasm.h"
 #include "pass.h"
 #include "ast_utils.h"
+#include "asm_v_wasm.h"
 
 namespace wasm {
 
@@ -97,8 +99,37 @@ struct ReachabilityAnalyzer : public PostWalker<ReachabilityAnalyzer> {
   }
 };
 
+// Finds function type usage
+
+struct FunctionTypeAnalyzer : public PostWalker<FunctionTypeAnalyzer> {
+  std::vector<Import*> functionImports;
+  std::vector<Function*> functions;
+  std::vector<CallIndirect*> indirectCalls;
+
+  void visitImport(Import* curr) {
+    if (curr->kind == ExternalKind::Function && curr->functionType.is()) {
+      functionImports.push_back(curr);
+    }
+  }
+
+  void visitFunction(Function* curr) {
+    if (curr->type.is()) {
+      functions.push_back(curr);
+    }
+  }
+
+  void visitCallIndirect(CallIndirect* curr) {
+    indirectCalls.push_back(curr);
+  }
+};
+
 struct RemoveUnusedModuleElements : public Pass {
   void run(PassRunner* runner, Module* module) override {
+    optimizeGlobalsAndFunctions(module);
+    optimizeFunctionTypes(module);
+  }
+
+  void optimizeGlobalsAndFunctions(Module* module) {
     std::vector<ModuleElement> roots;
     // Module start is a root.
     if (module->start.is()) {
@@ -146,6 +177,40 @@ struct RemoveUnusedModuleElements : public Pass {
     }
     module->updateMaps();
   }
+
+  void optimizeFunctionTypes(Module* module) {
+    FunctionTypeAnalyzer analyzer;
+    analyzer.walkModule(module);
+    // maps each string signature to a single canonical function type
+    std::unordered_map<std::string, FunctionType*> canonicals;
+    std::unordered_set<FunctionType*> needed;
+    auto canonicalize = [&](Name name) {
+      FunctionType* type = module->getFunctionType(name);
+      auto sig = getSig(type);
+      auto iter = canonicals.find(sig);
+      if (iter == canonicals.end()) {
+        needed.insert(type);
+        canonicals[sig] = type;
+        return type->name;
+      } else {
+        return iter->second->name;
+      }
+    };
+    // canonicalize all uses of function types
+    for (auto* import : analyzer.functionImports) {
+      import->functionType = canonicalize(import->functionType);
+    }
+    for (auto* func : analyzer.functions) {
+      func->type = canonicalize(func->type);
+    }
+    for (auto* call : analyzer.indirectCalls) {
+      call->fullType = canonicalize(call->fullType);
+    }
+    // remove no-longer used types
+    module->functionTypes.erase(std::remove_if(module->functionTypes.begin(), module->functionTypes.end(), [&needed](std::unique_ptr<FunctionType>& type) {
+      return needed.count(type.get()) == 0;
+    }), module->functionTypes.end());
+  }
 };
 
 Pass* createRemoveUnusedModuleElementsPass() {
diff --git a/test/example/c-api-kitchen-sink.txt b/test/example/c-api-kitchen-sink.txt
index a5275bad5..26e3b820f 100644
--- a/test/example/c-api-kitchen-sink.txt
+++ b/test/example/c-api-kitchen-sink.txt
@@ -1025,9 +1025,6 @@ raw:
 )
 optimized:
 (module
- (type $v (func))
- (type $vi (func (param i32)))
- (type $i (func (result i32)))
  (memory $0 0)
 )
 module loaded from binary form:
@@ -2882,9 +2879,6 @@ raw:
 optimized:
   BinaryenModulePrint(the_module);
 (module
- (type $v (func))
- (type $vi (func (param i32)))
- (type $i (func (result i32)))
  (memory $0 0)
 )
   BinaryenModuleDispose(the_module);
diff --git a/test/example/c-api-kitchen-sink.txt.txt b/test/example/c-api-kitchen-sink.txt.txt
index 4583004e3..7f281bd67 100644
--- a/test/example/c-api-kitchen-sink.txt.txt
+++ b/test/example/c-api-kitchen-sink.txt.txt
@@ -1018,8 +1018,5 @@
  )
 )
 (module
- (type $v (func))
- (type $vi (func (param i32)))
- (type $i (func (result i32)))
  (memory $0 0)
 )
diff --git a/test/merge/basics.wast.combined.finalized.opt b/test/merge/basics.wast.combined.finalized.opt
index 320ba8fa7..119851fb0 100644
--- a/test/merge/basics.wast.combined.finalized.opt
+++ b/test/merge/basics.wast.combined.finalized.opt
@@ -1,8 +1,6 @@
 (module
  (type $ii (func (param i32 i32)))
  (type $FUNCSIG$v (func))
- (type $ii$0 (func (param i32 i32)))
- (type $FUNCSIG$v$0 (func))
  (import "env" "memory" (memory $0 256))
  (import "env" "table" (table 18 anyfunc))
  (import "env" "some-func" (func $some-func))
@@ -48,12 +46,12 @@
   (nop)
   (call $willCollide)
  )
- (func $only-b (type $FUNCSIG$v$0)
+ (func $only-b (type $FUNCSIG$v)
   (nop)
   (call $only-b)
   (call $some-func-b)
   (call $some-collide$0)
-  (call_indirect $ii$0
+  (call_indirect $ii
    (i32.const 12)
    (i32.const 34)
    (i32.const 56)
@@ -79,7 +77,7 @@
    (i32.const 5678)
   )
  )
- (func $willCollide$0 (type $FUNCSIG$v$0)
+ (func $willCollide$0 (type $FUNCSIG$v)
   (nop)
   (call $willCollide$0)
  )
diff --git a/test/merge/basics.wast.combined.opt b/test/merge/basics.wast.combined.opt
index ba8ab6220..c595d00da 100644
--- a/test/merge/basics.wast.combined.opt
+++ b/test/merge/basics.wast.combined.opt
@@ -1,8 +1,6 @@
 (module
  (type $ii (func (param i32 i32)))
  (type $FUNCSIG$v (func))
- (type $ii$0 (func (param i32 i32)))
- (type $FUNCSIG$v$0 (func))
  (import "env" "memoryBase" (global $memoryBase i32))
  (import "env" "tableBase" (global $tableBase i32))
  (import "env" "memory" (memory $0 256))
@@ -56,12 +54,12 @@
   (nop)
   (call $willCollide)
  )
- (func $only-b (type $FUNCSIG$v$0)
+ (func $only-b (type $FUNCSIG$v)
   (nop)
   (call $only-b)
   (call $some-func-b)
   (call $some-collide$0)
-  (call_indirect $ii$0
+  (call_indirect $ii
    (i32.const 12)
    (i32.const 34)
    (i32.const 56)
@@ -115,7 +113,7 @@
    (i32.const 5678)
   )
  )
- (func $willCollide$0 (type $FUNCSIG$v$0)
+ (func $willCollide$0 (type $FUNCSIG$v)
   (nop)
   (call $willCollide$0)
  )
diff --git a/test/merge/dylib.wasm.combined.finalized.opt b/test/merge/dylib.wasm.combined.finalized.opt
index eda229e34..e5a26a5d3 100644
--- a/test/merge/dylib.wasm.combined.finalized.opt
+++ b/test/merge/dylib.wasm.combined.finalized.opt
@@ -1,11 +1,7 @@
 (module
- (type $0 (func (param i32 i32)))
  (type $1 (func (param i32) (result i32)))
  (type $2 (func (result i32)))
  (type $3 (func))
- (type $0$0 (func (param i32 i32)))
- (type $1$0 (func (result i32)))
- (type $2$0 (func))
  (import "env" "_puts" (func $import$1 (param i32) (result i32)))
  (import "env" "memory" (memory $0 256))
  (import "env" "table" (table 8 anyfunc))
@@ -55,7 +51,7 @@
    )
   )
  )
- (func $_foo (type $1$0) (result i32)
+ (func $_foo (type $2) (result i32)
   (local $var$0 i32)
   (block $label$0 i32
    (block $label$1 i32
@@ -66,10 +62,10 @@
    )
   )
  )
- (func $runPostSets$0 (type $2$0)
+ (func $runPostSets$0 (type $3)
   (nop)
  )
- (func $__post_instantiate$0 (type $2$0)
+ (func $__post_instantiate$0 (type $3)
   (block $label$0
    (block $label$1
     (set_global $global$0$0
diff --git a/test/merge/dylib.wasm.combined.opt b/test/merge/dylib.wasm.combined.opt
index 404281cda..1954735d9 100644
--- a/test/merge/dylib.wasm.combined.opt
+++ b/test/merge/dylib.wasm.combined.opt
@@ -1,11 +1,7 @@
 (module
- (type $0 (func (param i32 i32)))
  (type $1 (func (param i32) (result i32)))
  (type $2 (func (result i32)))
  (type $3 (func))
- (type $0$0 (func (param i32 i32)))
- (type $1$0 (func (result i32)))
- (type $2$0 (func))
  (import "env" "memoryBase" (global $import$0 i32))
  (import "env" "_puts" (func $import$1 (param i32) (result i32)))
  (import "env" "memory" (memory $0 256))
@@ -61,7 +57,7 @@
    )
   )
  )
- (func $_foo (type $1$0) (result i32)
+ (func $_foo (type $2) (result i32)
   (local $var$0 i32)
   (block $label$0 i32
    (block $label$1 i32
@@ -72,10 +68,10 @@
    )
   )
  )
- (func $runPostSets$0 (type $2$0)
+ (func $runPostSets$0 (type $3)
   (nop)
  )
- (func $__post_instantiate$0 (type $2$0)
+ (func $__post_instantiate$0 (type $3)
   (block $label$0
    (block $label$1
     (set_global $global$0$0
diff --git a/test/merge/fusing.wast.combined.finalized.opt b/test/merge/fusing.wast.combined.finalized.opt
index 202c566e1..9df0a5cdc 100644
--- a/test/merge/fusing.wast.combined.finalized.opt
+++ b/test/merge/fusing.wast.combined.finalized.opt
@@ -1,6 +1,5 @@
 (module
  (type $FUNCSIG$v (func))
- (type $FUNCSIG$v$0 (func))
  (import "env" "memory" (memory $0 256))
  (import "env" "table" (table 8 anyfunc))
  (global $a-global i32 (i32.const 0))
@@ -20,7 +19,7 @@
    (get_global $b-global)
   )
  )
- (func $bar-func (type $FUNCSIG$v$0)
+ (func $bar-func (type $FUNCSIG$v)
   (nop)
   (drop
    (get_global $a-global)
diff --git a/test/merge/fusing.wast.combined.opt b/test/merge/fusing.wast.combined.opt
index 1a09fb784..c8d58890e 100644
--- a/test/merge/fusing.wast.combined.opt
+++ b/test/merge/fusing.wast.combined.opt
@@ -1,6 +1,5 @@
 (module
  (type $FUNCSIG$v (func))
- (type $FUNCSIG$v$0 (func))
  (import "env" "memoryBase" (global $memoryBase i32))
  (import "env" "tableBase" (global $tableBase i32))
  (import "env" "memory" (memory $0 256))
@@ -23,7 +22,7 @@
    (get_global $b-global)
   )
  )
- (func $bar-func (type $FUNCSIG$v$0)
+ (func $bar-func (type $FUNCSIG$v)
   (nop)
   (drop
    (get_global $a-global)
diff --git a/test/merge/post-instantiate-a.wast.combined.finalized.opt b/test/merge/post-instantiate-a.wast.combined.finalized.opt
index 2fc7dbdb3..d0a9b4a6a 100644
--- a/test/merge/post-instantiate-a.wast.combined.finalized.opt
+++ b/test/merge/post-instantiate-a.wast.combined.finalized.opt
@@ -1,6 +1,5 @@
 (module
  (type $0 (func))
- (type $0$0 (func))
  (memory $0 1)
  (data (i32.const 1024) "")
  (export "__post_instantiate" (func $0))
diff --git a/test/merge/post-instantiate-a.wast.combined.opt b/test/merge/post-instantiate-a.wast.combined.opt
index 8f190f187..b4ecf737b 100644
--- a/test/merge/post-instantiate-a.wast.combined.opt
+++ b/test/merge/post-instantiate-a.wast.combined.opt
@@ -1,6 +1,5 @@
 (module
  (type $0 (func))
- (type $0$0 (func))
  (import "env" "memoryBase" (global $memoryBase i32))
  (import "env" "tableBase" (global $tableBase i32))
  (memory $0 0)
diff --git a/test/merge/post-instantiate-b.wast.combined.finalized.opt b/test/merge/post-instantiate-b.wast.combined.finalized.opt
index f335f08d9..367bd834b 100644
--- a/test/merge/post-instantiate-b.wast.combined.finalized.opt
+++ b/test/merge/post-instantiate-b.wast.combined.finalized.opt
@@ -1,5 +1,4 @@
 (module
- (type $0 (func))
  (type $0$0 (func))
  (memory $0 1)
  (data (i32.const 1024) "")
diff --git a/test/merge/post-instantiate-b.wast.combined.opt b/test/merge/post-instantiate-b.wast.combined.opt
index 1a85ff79d..b847fc5e8 100644
--- a/test/merge/post-instantiate-b.wast.combined.opt
+++ b/test/merge/post-instantiate-b.wast.combined.opt
@@ -1,5 +1,4 @@
 (module
- (type $0 (func))
  (type $0$0 (func))
  (import "env" "memoryBase" (global $memoryBase i32))
  (import "env" "tableBase" (global $tableBase i32))
diff --git a/test/merge/printf.wast.combined.finalized.opt b/test/merge/printf.wast.combined.finalized.opt
index 92e17f720..a801ab875 100644
--- a/test/merge/printf.wast.combined.finalized.opt
+++ b/test/merge/printf.wast.combined.finalized.opt
@@ -1,7 +1,5 @@
 (module
  (type $FUNCSIG$iii (func (param i32 i32) (result i32)))
- (type $FUNCSIG$iii$0 (func (param i32 i32) (result i32)))
- (type $1 (func))
  (import "env" "memory" (memory $0 256))
  (import "env" "table" (table 8 anyfunc))
  
diff --git a/test/merge/printf.wast.combined.opt b/test/merge/printf.wast.combined.opt
index e01710fd1..ae913e40e 100644
--- a/test/merge/printf.wast.combined.opt
+++ b/test/merge/printf.wast.combined.opt
@@ -1,7 +1,5 @@
 (module
  (type $FUNCSIG$iii (func (param i32 i32) (result i32)))
- (type $FUNCSIG$iii$0 (func (param i32 i32) (result i32)))
- (type $1 (func))
  (import "env" "memoryBase" (global $memoryBase i32))
  (import "env" "tableBase" (global $tableBase i32))
  (import "env" "memory" (memory $0 256))
diff --git a/test/passes/remove-unused-module-elements.txt b/test/passes/remove-unused-module-elements.txt
index 5fa7d005a..abce5a5aa 100644
--- a/test/passes/remove-unused-module-elements.txt
+++ b/test/passes/remove-unused-module-elements.txt
@@ -1,10 +1,14 @@
 (module
  (type $0 (func))
+ (type $1 (func (param i32)))
+ (type $2 (func (param i32) (result i32)))
  (table 1 1 anyfunc)
  (elem (i32.const 0) $called_indirect)
  (memory $0 0)
  (export "memory" (memory $0))
  (export "exported" (func $exported))
+ (export "other1" (func $other1))
+ (export "other2" (func $other2))
  (start $start)
  (func $start (type $0)
   (call $called0)
@@ -31,6 +35,49 @@
  (func $called4 (type $0)
   (call $called3)
  )
+ (func $other1 (type $1) (param $0 i32)
+  (call_indirect $0
+   (i32.const 0)
+  )
+  (call_indirect $0
+   (i32.const 0)
+  )
+  (call_indirect $0
+   (i32.const 0)
+  )
+  (call_indirect $0
+   (i32.const 0)
+  )
+  (call_indirect $1
+   (i32.const 0)
+   (i32.const 0)
+  )
+  (call_indirect $1
+   (i32.const 0)
+   (i32.const 0)
+  )
+  (drop
+   (call_indirect $2
+    (i32.const 0)
+    (i32.const 0)
+   )
+  )
+  (drop
+   (call_indirect $2
+    (i32.const 0)
+    (i32.const 0)
+   )
+  )
+  (drop
+   (call_indirect $2
+    (i32.const 0)
+    (i32.const 0)
+   )
+  )
+ )
+ (func $other2 (type $1) (param $0 i32)
+  (unreachable)
+ )
 )
 (module
  (import "env" "memory" (memory $0 256))
diff --git a/test/passes/remove-unused-module-elements.wast b/test/passes/remove-unused-module-elements.wast
index eb071ee66..e0ae7b4ec 100644
--- a/test/passes/remove-unused-module-elements.wast
+++ b/test/passes/remove-unused-module-elements.wast
@@ -2,8 +2,16 @@
   (memory 0)
   (start $start)
   (type $0 (func))
+  (type $0-dupe (func))
+  (type $1 (func (param i32)))
+  (type $1-dupe (func (param i32)))
+  (type $2 (func (param i32) (result i32)))
+  (type $2-dupe (func (param i32) (result i32)))
+  (type $2-thrupe (func (param i32) (result i32)))
   (export "memory" (memory $0))
   (export "exported" $exported)
+  (export "other1" $other1)
+  (export "other2" $other2)
   (table 1 1 anyfunc)
   (elem (i32.const 0) $called_indirect)
   (func $start (type $0)
@@ -18,26 +26,26 @@
   (func $called_indirect (type $0)
     (nop)
   )
-  (func $exported (type $0)
+  (func $exported (type $0-dupe)
     (call $called2)
   )
-  (func $called2 (type $0)
+  (func $called2 (type $0-dupe)
     (call $called2)
     (call $called3)
   )
-  (func $called3 (type $0)
+  (func $called3 (type $0-dupe)
     (call $called4)
   )
-  (func $called4 (type $0)
+  (func $called4 (type $0-dupe)
     (call $called3)
   )
-  (func $remove0 (type $0)
+  (func $remove0 (type $0-dupe)
     (call $remove1)
   )
-  (func $remove1 (type $0)
+  (func $remove1 (type $0-dupe)
     (nop)
   )
-  (func $remove2 (type $0)
+  (func $remove2 (type $0-dupe)
     (call $remove2)
   )
   (func $remove3 (type $0)
@@ -46,6 +54,20 @@
   (func $remove4 (type $0)
     (call $remove3)
   )
+  (func $other1 (param i32) (type $1)
+    (call_indirect $0 (i32.const 0))
+    (call_indirect $0 (i32.const 0))
+    (call_indirect $0-dupe (i32.const 0))
+    (call_indirect $0-dupe (i32.const 0))
+    (call_indirect $1 (i32.const 0) (i32.const 0))
+    (call_indirect $1-dupe (i32.const 0) (i32.const 0))
+    (drop (call_indirect $2 (i32.const 0) (i32.const 0)))
+    (drop (call_indirect $2-dupe (i32.const 0) (i32.const 0)))
+    (drop (call_indirect $2-thrupe (i32.const 0) (i32.const 0)))
+  )
+  (func $other2 (param i32) (type $1-dupe)
+    (unreachable)
+  )
 )
 (module ;; leave the table and memory alone
   (import "env" "memory" (memory $0 256))