Collect signatures from all block kinds (#2691)

Previously the signature collection mechanism responsible for
populating the type section with signatures used by instructions only
collected signatures from indirect call and block instructions. This
works as long as all other control flow constructs like ifs, loops,
and tries contain blocks with the same signature. But it is possible
to have an if with non-block children, and we would need to collect
its signature as well.

1 files changed, 11 insertions, 7 deletions

diff --git a/src/ir/module-utils.h b/src/ir/module-utils.h
index 0e18266e8..be924f742 100644
--- a/src/ir/module-utils.h
+++ b/src/ir/module-utils.h
@@ -19,6 +19,7 @@
 
 #include "ir/find_all.h"
 #include "ir/manipulation.h"
+#include "ir/properties.h"
 #include "pass.h"
 #include "support/unique_deferring_queue.h"
 #include "wasm.h"
@@ -435,16 +436,19 @@ collectSignatures(Module& wasm,
     if (func->imported()) {
       return;
     }
-    struct TypeCounter : PostWalker<TypeCounter> {
+    struct TypeCounter
+      : PostWalker<TypeCounter, UnifiedExpressionVisitor<TypeCounter>> {
       Counts& counts;
 
       TypeCounter(Counts& counts) : counts(counts) {}
-
-      void visitCallIndirect(CallIndirect* curr) { counts[curr->sig]++; }
-      void visitBlock(Block* curr) {
-        // TODO: Allow blocks to have input types as well
-        if (curr->type.isMulti()) {
-          counts[Signature(Type::none, curr->type)]++;
+      void visitExpression(Expression* curr) {
+        if (auto* call = curr->dynCast<CallIndirect>()) {
+          counts[call->sig]++;
+        } else if (Properties::isControlFlowStructure(curr)) {
+          // TODO: Allow control flow to have input types as well
+          if (curr->type.isMulti()) {
+            counts[Signature(Type::none, curr->type)]++;
+          }
         }
       }
     };