1 files changed, 68 insertions, 10 deletions

diff --git a/src/passes/RemoveUnusedNames.cpp b/src/passes/RemoveUnusedNames.cpp
index 1cd2a483f..9c6743479 100644
--- a/src/passes/RemoveUnusedNames.cpp
+++ b/src/passes/RemoveUnusedNames.cpp
@@ -15,7 +15,8 @@
  */
 
 //
-// Removes names from locations that are never branched to.
+// Removes names from locations that are never branched to, and
+// merge names when possible (by merging their blocks)
 //
 
 #include <wasm.h>
@@ -30,27 +31,84 @@ struct RemoveUnusedNames : public WalkerPass<PostWalker<RemoveUnusedNames, Visit
 
   // We maintain a list of branches that we saw in children, then when we reach
   // a parent block, we know if it was branched to
-  std::set<Name> branchesSeen;
+  std::map<Name, std::set<Expression*>> branchesSeen;
 
   void visitBreak(Break *curr) {
-    branchesSeen.insert(curr->name);
+    branchesSeen[curr->name].insert(curr);
+  }
+
+  void visitSwitch(Switch *curr) {
+    for (auto name : curr->targets) {
+      branchesSeen[name].insert(curr);
+    }
+    branchesSeen[curr->default_].insert(curr);
+  }
+
+  void handleBreakTarget(Name& name) {
+    if (name.is()) {
+      if (branchesSeen.find(name) == branchesSeen.end()) {
+        name = Name();
+      } else {
+        branchesSeen.erase(name);
+      }
+    }
   }
 
   void visitBlock(Block *curr) {
-    if (curr->name.is() && branchesSeen.count(curr->name) == 0) {
-      curr->name = Name();
+    if (curr->name.is() && curr->list.size() == 1) {
+      auto* child = curr->list[0]->dynCast<Block>();
+      if (child && child->name.is()) {
+        // we have just one child, this block, so breaking out of it goes to the same place as breaking out of us, we just need one name (and block)
+        auto& branches = branchesSeen[curr->name];
+        for (auto* branch : branches) {
+          if (Break* br = branch->dynCast<Break>()) {
+            if (br->name == curr->name) br->name = child->name;
+          } else if (Switch* sw = branch->dynCast<Switch>()) {
+            for (auto& target : sw->targets) {
+              if (target == curr->name) target = child->name;
+            }
+            if (sw->default_ == curr->name) sw->default_ = child->name;
+          } else {
+            WASM_UNREACHABLE();
+          }
+        }
+        replaceCurrent(child);
+      }
+    }
+    handleBreakTarget(curr->name);
+    if (curr->name.is() && curr->list.size() == 1) {
+      auto* child = curr->list[0]->dynCast<Loop>();
+      if (child && !child->out.is()) {
+        // we have just one child, this loop, and it lacks an out label. So this block's name is doing just that!
+        child->out = curr->name;
+        replaceCurrent(child);
+      }
     }
   }
 
-  void visitSwitch(Switch *curr) {
-    for (auto name : curr->targets) {
-      branchesSeen.insert(name);
+  void visitLoop(Loop *curr) {
+    handleBreakTarget(curr->in);
+    // Loops can have just 'in', but cannot have just 'out'
+    auto out = curr->out;
+    handleBreakTarget(curr->out);
+    if (curr->out.is() && !curr->in.is()) {
+      auto* block = getModule()->allocator.alloc<Block>();
+      block->name = out;
+      block->list.push_back(curr->body);
+      replaceCurrent(block);
+    }
+    if (curr->in.is() && !curr->out.is()) {
+      auto* child = curr->body->dynCast<Block>();
+      if (child && child->name.is()) {
+        // we have just one child, this block, and we lack an out label. So we can take the block's!
+        curr->out = child->name;
+        child->name = Name();
+      }
     }
-    branchesSeen.insert(curr->default_);
   }
 
   void visitFunction(Function *curr) {
-    branchesSeen.clear();
+    assert(branchesSeen.empty());
   }
 };