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;
 }