de-recurse operations on nested blocks

5 files changed, 134 insertions, 24 deletions

diff --git a/src/passes/Print.cpp b/src/passes/Print.cpp
index b69f7a977..2a940359f 100644
--- a/src/passes/Print.cpp
+++ b/src/passes/Print.cpp
@@ -54,13 +54,38 @@ struct PrintSExpression : public WasmVisitor<PrintSExpression, void> {
     o << maybeNewLine;
   }
   void visitBlock(Block *curr) {
-    printOpening(o, "block");
-    if (curr->name.is()) {
-      o << ' ' << curr->name;
+    // special-case Block, because Block nesting (in their first element) can be incredibly deep
+    std::vector<Block*> stack;
+    while (1) {
+      if (stack.size() > 0) doIndent(o, indent);
+      stack.push_back(curr);
+      printOpening(o, "block");
+      if (curr->name.is()) {
+        o << ' ' << curr->name;
+      }
+      incIndent();
+      if (curr->list.size() > 0 && curr->list[0]->is<Block>()) {
+        // recurse into the first element
+        curr = curr->list[0]->cast<Block>();
+        continue;
+      } else {
+        break; // that's all we can recurse, start to unwind
+      }
     }
-    incIndent();
-    for (auto expression : curr->list) {
-      printFullLine(expression);
+    auto* top = stack.back();
+    while (stack.size() > 0) {
+      curr = stack.back();
+      stack.pop_back();
+      auto& list = curr->list;
+      for (size_t i = 0; i < list.size(); i++) {
+        if (curr != top && i == 0) {
+          // one of the block recursions we already handled
+          decIndent();
+          o << '\n';
+          continue;
+        }
+        printFullLine(list[i]);
+      }
     }
     decIndent();
   }
diff --git a/src/shared-constants.h b/src/shared-constants.h
index 03409de6e..47fa60947 100644
--- a/src/shared-constants.h
+++ b/src/shared-constants.h
@@ -67,6 +67,7 @@ cashew::IString GLOBAL("global"),
                 CALL("call"),
                 CALL_IMPORT("call_import"),
                 CALL_INDIRECT("call_indirect"),
+                BLOCK("block"),
                 BR_IF("br_if"),
                 THEN("then"),
                 ELSE("else"),
diff --git a/src/wasm-interpreter.h b/src/wasm-interpreter.h
index f6a3e1027..fb860627c 100644
--- a/src/wasm-interpreter.h
+++ b/src/wasm-interpreter.h
@@ -204,12 +204,33 @@ private:
 
       Flow visitBlock(Block *curr) {
         NOTE_ENTER("Block");
+        // special-case Block, because Block nesting (in their first element) can be incredibly deep
+        std::vector<Block*> stack;
+        stack.push_back(curr);
+        while (curr->list.size() > 0 && curr->list[0]->is<Block>()) {
+          curr = curr->list[0]->cast<Block>();
+          stack.push_back(curr);
+        }
         Flow flow;
-        for (auto expression : curr->list) {
-          flow = visit(expression);
+        auto* top = stack.back();
+        while (stack.size() > 0) {
+          curr = stack.back();
+          stack.pop_back();
           if (flow.breaking()) {
             flow.clearIf(curr->name);
-            return flow;
+            continue;
+          }
+          auto& list = curr->list;
+          for (size_t i = 0; i < list.size(); i++) {
+            if (curr != top && i == 0) {
+              // one of the block recursions we already handled
+              continue;
+            }
+            flow = visit(list[i]);
+            if (flow.breaking()) {
+              flow.clearIf(curr->name);
+              break;
+            }
           }
         }
         return flow;
diff --git a/src/wasm-s-parser.h b/src/wasm-s-parser.h
index cfabbf52d..db57dda9f 100644
--- a/src/wasm-s-parser.h
+++ b/src/wasm-s-parser.h
@@ -699,22 +699,55 @@ private:
   }
 
   Expression* makeBlock(Element& s) {
-    auto ret = allocator.alloc<Block>();
-    size_t i = 1;
-    if (i >= s.size()) return ret; // empty block
-    if (s[i]->isStr()) {
-      ret->name = s[i]->str();
-      i++;
-    } else {
-      ret->name = getPrefixedName("block");
+    // special-case Block, because Block nesting (in their first element) can be incredibly deep
+    auto curr = allocator.alloc<Block>();
+    auto* sp = &s;
+    std::vector<std::pair<Element*, Block*>> stack;
+    while (1) {
+      stack.emplace_back(sp, curr);
+      auto& s = *sp;
+      size_t i = 1;
+      if (i < s.size() && s[i]->isStr()) {
+        curr->name = s[i]->str();
+        i++;
+      } else {
+        curr->name = getPrefixedName("block");
+      }
+      labelStack.push_back(curr->name);
+      if (i >= s.size()) break; // labeled empty block
+      auto& first = *s[i];
+      if (first[0]->str() == BLOCK) {
+        // recurse
+        curr = allocator.alloc<Block>();
+        sp = &first;
+        continue;
+      }
+      break;
     }
-    labelStack.push_back(ret->name);
-    for (; i < s.size(); i++) {
-      ret->list.push_back(parseExpression(s[i]));
+    // we now have a stack of Blocks, with their labels, but no contents yet
+    for (int t = int(stack.size()) - 1; t >= 0; t--) {
+      auto* sp = stack[t].first;
+      auto* curr = stack[t].second;
+      auto& s = *sp;
+      size_t i = 1;
+      if (i < s.size()) {
+        if (s[i]->isStr()) {
+          i++;
+        }
+        if (t < int(stack.size()) - 1) {
+          // first child is one of our recursions
+          curr->list.push_back(stack[t + 1].second);
+          i++;
+        }
+        for (; i < s.size(); i++) {
+          curr->list.push_back(parseExpression(s[i]));
+        }
+      }
+      assert(labelStack.back() == curr->name);
+      labelStack.pop_back();
+      curr->finalize();
     }
-    labelStack.pop_back();
-    ret->finalize();
-    return ret;
+    return stack[0].second;
   }
 
   // Similar to block, but the label is handled by the enclosing if (since there might not be a then or else, ick)
diff --git a/src/wasm.h b/src/wasm.h
index 4bcb4bc87..b52803799 100644
--- a/src/wasm.h
+++ b/src/wasm.h
@@ -1408,7 +1408,37 @@ struct WasmWalker : public WasmWalkerBase<SubType, ReturnType> {
   void walk(Expression*& curr) override {
     if (!curr) return;
 
-    ChildWalker<WasmWalker<SubType, ReturnType>>(*this).visit(curr);
+    // special-case Block, because Block nesting (in their first element) can be incredibly deep
+    if (curr->is<Block>()) {
+      auto* block = curr->dyn_cast<Block>();
+      std::vector<Block*> stack;
+      stack.push_back(block);
+      while (block->list.size() > 0 && block->list[0]->is<Block>()) {
+        block = block->list[0]->cast<Block>();
+        stack.push_back(block);
+      }
+      // walk all the children
+      for (int i = int(stack.size()) - 1; i >= 0; i--) {
+        auto* block = stack[i];
+        auto& children = block->list;
+        for (size_t j = 0; j < children.size(); j++) {
+          if (i < int(stack.size()) - 1 && j == 0) {
+            // this is one of the stacked blocks, no need to walk its children, we are doing that ourselves
+            this->visit(children[0]);
+            if (replace) {
+              children[0] = replace;
+              replace = nullptr;
+            }
+          } else {
+            this->walk(children[j]);
+          }
+        }
+      }
+      // we walked all the children, and can rejoin later below to visit this node itself
+    } else {
+      // generic child-walking
+      ChildWalker<WasmWalker<SubType, ReturnType>>(*this).visit(curr);
+    }
 
     this->visit(curr);