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.

5 files changed, 32 insertions, 15 deletions

diff --git a/src/ir/flat.h b/src/ir/flat.h
index 01a94a759..d54d4771f 100644
--- a/src/ir/flat.h
+++ b/src/ir/flat.h
@@ -64,17 +64,12 @@ namespace wasm {
 
 namespace Flat {
 
-inline bool isControlFlowStructure(Expression* curr) {
-  return curr->is<Block>() || curr->is<If>() || curr->is<Loop>() ||
-         curr->is<Try>();
-}
-
 inline void verifyFlatness(Function* func) {
   struct VerifyFlatness
     : public PostWalker<VerifyFlatness,
                         UnifiedExpressionVisitor<VerifyFlatness>> {
     void visitExpression(Expression* curr) {
-      if (isControlFlowStructure(curr)) {
+      if (Properties::isControlFlowStructure(curr)) {
         verify(!curr->type.isConcrete(),
                "control flow structures must not flow values");
       } else if (curr->is<LocalSet>()) {
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)]++;
+          }
         }
       }
     };
diff --git a/src/ir/properties.h b/src/ir/properties.h
index b4dfbbd5e..094d90bd6 100644
--- a/src/ir/properties.h
+++ b/src/ir/properties.h
@@ -59,6 +59,11 @@ inline bool isSymmetric(Binary* binary) {
   }
 }
 
+inline bool isControlFlowStructure(Expression* curr) {
+  return curr->is<Block>() || curr->is<If>() || curr->is<Loop>() ||
+         curr->is<Try>();
+}
+
 // Check if an expression is a control flow construct with a name,
 // which implies it may have breaks to it.
 inline bool isNamedControlFlow(Expression* curr) {
diff --git a/src/passes/Flatten.cpp b/src/passes/Flatten.cpp
index 13aa985aa..0a6a7022a 100644
--- a/src/passes/Flatten.cpp
+++ b/src/passes/Flatten.cpp
@@ -68,7 +68,7 @@ struct Flatten
       return;
     }
 
-    if (Flat::isControlFlowStructure(curr)) {
+    if (Properties::isControlFlowStructure(curr)) {
       // handle control flow explicitly. our children do not have control flow,
       // but they do have preludes which we need to set up in the right place
 
@@ -291,7 +291,7 @@ struct Flatten
     // next, finish up: migrate our preludes if we can
     if (!ourPreludes.empty()) {
       auto* parent = getParent();
-      if (parent && !Flat::isControlFlowStructure(parent)) {
+      if (parent && !Properties::isControlFlowStructure(parent)) {
         auto& parentPreludes = preludes[parent];
         for (auto* prelude : ourPreludes) {
           parentPreludes.push_back(prelude);
diff --git a/src/wasm/wasm-type.cpp b/src/wasm/wasm-type.cpp
index 148b55716..024894dfb 100644
--- a/src/wasm/wasm-type.cpp
+++ b/src/wasm/wasm-type.cpp
@@ -237,6 +237,19 @@ bool Type::isSubType(Type left, Type right) {
       (right == Type::anyref || left == Type::nullref)) {
     return true;
   }
+  if (left.isMulti() && right.isMulti()) {
+    const auto& leftElems = left.expand();
+    const auto& rightElems = right.expand();
+    if (leftElems.size() != rightElems.size()) {
+      return false;
+    }
+    for (size_t i = 0; i < leftElems.size(); ++i) {
+      if (!isSubType(leftElems[i], rightElems[i])) {
+        return false;
+      }
+    }
+    return true;
+  }
   return false;
 }