1 files changed, 52 insertions, 4 deletions

diff --git a/src/passes/MergeBlocks.cpp b/src/passes/MergeBlocks.cpp
index 147ba9f45..e0c5e575a 100644
--- a/src/passes/MergeBlocks.cpp
+++ b/src/passes/MergeBlocks.cpp
@@ -83,9 +83,11 @@
 
 namespace wasm {
 
-// Looks for reasons we can't remove the values from breaks to an origin
-// For example, if there is a switch targeting us, we can't do it - we can't
-// remove the value from other targets
+// Looks for reasons we can't remove the values from breaks to an origin. This
+// is run when we know the value sent to that block is dropped, so the value is
+// not needed, but some corner cases stop us (for example, if there is a switch
+// targeting us, we can't do it - we can't remove the value from the switch's
+// other targets).
 struct ProblemFinder
   : public ControlFlowWalker<ProblemFinder,
                              UnifiedExpressionVisitor<ProblemFinder>> {
@@ -123,6 +125,39 @@ struct ProblemFinder
       return;
     }
 
+    if (auto* tryy = curr->dynCast<TryTable>()) {
+      auto num = tryy->catchTags.size();
+      for (Index i = 0; i < num; i++) {
+        if (tryy->catchDests[i] == origin) {
+          // This try_table branches to the origin we care about. We know the
+          // value being sent to the block is dropped, so we'd like to stop
+          // anything from being sent to it. We can stop a ref from being sent,
+          // so if that is enough to remove all the values, then we can
+          // optimize here. In other words, if this is a catch_all_ref (which
+          // can only send a ref) or this is a catch_ref of a specific tag that
+          // has no contents (so if we remove the ref, nothing remains), then we
+          // can optimize, but if this is is a catch of a tag *with* contents
+          // then those contents stop us.
+          //
+          // TODO: We could also support cases where the target block has
+          //       multiple values, and remove just ref at the end. That might
+          //       make more sense in TupleOptimization though as it would need
+          //       to track uses of parts of a tuple.
+          if (!tryy->catchTags[i] ||
+              getModule()->getTag(tryy->catchTags[i])->sig.params.size() == 0) {
+            // There must be a ref here, otherwise there is no value being sent
+            // at all, and we should not be running ProblemFinder at all.
+            assert(tryy->catchRefs[i]);
+          } else {
+            // Anything else is a problem.
+            foundProblem = true;
+            return;
+          }
+        }
+      }
+      return;
+    }
+
     // Any other branch type - switch, br_on, etc. - is not handled yet.
     BranchUtils::operateOnScopeNameUses(curr, [&](Name& name) {
       if (name == origin) {
@@ -174,6 +209,18 @@ struct BreakValueDropper : public ControlFlowWalker<BreakValueDropper> {
       replaceCurrent(curr->value);
     }
   }
+
+  void visitTryTable(TryTable* curr) {
+    auto num = curr->catchTags.size();
+    for (Index i = 0; i < num; i++) {
+      if (curr->catchDests[i] == origin) {
+        // Remove the existing ref being sent.
+        assert(curr->catchRefs[i]);
+        curr->catchRefs[i] = false;
+        curr->sentTypes[i] = Type::none;
+      }
+    }
+  }
 };
 
 // Checks for code after an unreachable element.
@@ -223,7 +270,8 @@ static bool optimizeDroppedBlock(Drop* drop,
     drop->finalize();
     block->list.back() = drop;
   }
-  block->finalize();
+  // Remove the old type, which was from the value we just removed.
+  block->finalize(Type::none);
   return true;
 }