3 files changed, 80 insertions, 14 deletions

diff --git a/src/asm2wasm.h b/src/asm2wasm.h
index e91a9344e..0320b7924 100644
--- a/src/asm2wasm.h
+++ b/src/asm2wasm.h
@@ -1573,6 +1573,7 @@ void Asm2WasmBuilder::optimize() {
   PassRunner passRunner(&allocator);
   passRunner.add("remove-unused-brs");
   passRunner.add("remove-unused-names");
+  passRunner.add("merge-blocks");
   passRunner.run(&wasm);
 }
 
diff --git a/src/passes/MergeBlocks.cpp b/src/passes/MergeBlocks.cpp
new file mode 100644
index 000000000..3d7afc4cb
--- /dev/null
+++ b/src/passes/MergeBlocks.cpp
@@ -0,0 +1,56 @@
+/*
+ * Copyright 2015 WebAssembly Community Group participants
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+//
+// Merges blocks to their parents.
+//
+
+#include <wasm.h>
+#include <pass.h>
+
+namespace wasm {
+
+struct MergeBlocks : public Pass {
+  void visitBlock(Block *curr) override {
+    bool more = true;
+    while (more) {
+      more = false;
+      for (size_t i = 0; i < curr->list.size(); i++) {
+        Block* child = curr->list[i]->dyn_cast<Block>();
+        if (!child) continue;
+        if (child->name.is()) continue; // named blocks can have breaks to them (and certainly do, if we ran RemoveUnusedNames and RemoveUnusedBrs)
+        ExpressionList merged;
+        for (size_t j = 0; j < i; j++) {
+          merged.push_back(curr->list[j]);
+        }
+        for (auto item : child->list) {
+          merged.push_back(item);
+        }
+        for (size_t j = i + 1; j < curr->list.size(); j++) {
+          merged.push_back(curr->list[j]);
+        }
+        curr->list = merged;
+        more = true;
+        break;
+      }
+    }
+  }
+};
+
+static RegisterPass<MergeBlocks> registerPass("merge-blocks", "merges blocks to their parents");
+
+} // namespace wasm
+
diff --git a/src/wasm-s-parser.h b/src/wasm-s-parser.h
index f204d933a..dfef00d9b 100644
--- a/src/wasm-s-parser.h
+++ b/src/wasm-s-parser.h
@@ -297,6 +297,7 @@ private:
   size_t localIndex; // params and locals
   size_t otherIndex;
   std::vector<Name> labelStack;
+  bool hasReturn;
 
   Name getPrefixedName(std::string prefix) {
     return IString((prefix + std::to_string(otherIndex++)).c_str(), false);
@@ -316,7 +317,9 @@ private:
     func->body = nullptr;
     localIndex = 0;
     otherIndex = 0;
+    hasReturn = false;
     std::vector<NameType> typeParams; // we may have both params and a type. store the type info here
+    Block* autoBlock = nullptr; // we may need to add a block for the very top level
     for (;i < s.size(); i++) {
       Element& curr = *s[i];
       IString id = curr[0]->str();
@@ -368,17 +371,30 @@ private:
         if (!func->body) {
           func->body = ex;
         } else {
-          auto block = func->body->dyn_cast<Block>();
-          if (!block) {
-            block = allocator.alloc<Block>();
-            block->list.push_back(func->body);
-            func->body = block;
+          if (!autoBlock) {
+            autoBlock = allocator.alloc<Block>();
+            autoBlock->list.push_back(func->body);
+            func->body = autoBlock;
           }
-          block->list.push_back(ex);
+          autoBlock->list.push_back(ex);
         }
       }
     }
     if (!func->body) func->body = allocator.alloc<Nop>();
+    if (hasReturn) {
+      Block* body = func->body->dyn_cast<Block>();
+      if (!body) {
+        body = allocator.alloc<Block>();
+        body->list.push_back(func->body);
+        func->body = body;
+      }
+      if (body->name.is()) {
+        body = allocator.alloc<Block>();
+        body->list.push_back(func->body);
+        func->body = body;
+      }
+      body->name = FAKE_RETURN;
+    }
     wasm.addFunction(func);
     currLocalTypes.clear();
     labelStack.clear();
@@ -880,14 +896,7 @@ private:
 
   Expression* makeReturn(Element& s) {
     // return will likely not remain in wasm, but is in the testcases, for now. fake it
-    Block *temp;
-    if (!(currFunction->body && (temp = currFunction->body->dyn_cast<Block>()) && temp->name == FAKE_RETURN)) {
-      Expression* old = currFunction->body;
-      temp = allocator.alloc<Block>();
-      temp->name = FAKE_RETURN;
-      if (old) temp->list.push_back(old);
-      currFunction->body = temp;
-    }
+    hasReturn = true;
     auto ret = allocator.alloc<Break>();
     ret->name = FAKE_RETURN;
     if (s.size() >= 2) {