Merge pull request #540 from WebAssembly/merge-blocks

Merge blocks improvements

3 files changed, 164 insertions, 39 deletions

diff --git a/src/passes/MergeBlocks.cpp b/src/passes/MergeBlocks.cpp
index d26817e18..5794e1379 100644
--- a/src/passes/MergeBlocks.cpp
+++ b/src/passes/MergeBlocks.cpp
@@ -17,9 +17,53 @@
 //
 // Merges blocks to their parents.
 //
+// We also restructure blocks in order to enable such merging. For
+// example,
+//
+//  (i32.store
+//    (block
+//      (call $foo)
+//      (i32.load (i32.const 100))
+//    )
+//    (i32.const 0)
+//  )
+//
+// can be transformed into
+//
+//  (block
+//    (call $foo)
+//    (i32.store
+//      (block
+//        (i32.load (i32.const 100))
+//      )
+//      (i32.const 0)
+//    )
+//  )
+//
+// after which the internal block can go away, and
+// the new external block might be mergeable. This is always
+// worth it if the internal block ends up with 1 item.
+// For the second operand,
+//
+//  (i32.store
+//    (i32.const 100)
+//    (block
+//      (call $foo)
+//      (i32.load (i32.const 200))
+//    )
+//  )
+//
+// The order of operations requires that the first execute
+// before. We can do the same operation, but only if the
+// first has no side effects, or the code we are moving out
+// has no side effects.
+// If we can do this to both operands, we can generate a
+// single outside block.
+//
 
 #include <wasm.h>
 #include <pass.h>
+#include <ast_utils.h>
 
 namespace wasm {
 
@@ -50,6 +94,86 @@ struct MergeBlocks : public WalkerPass<PostWalker<MergeBlocks, Visitor<MergeBloc
       }
     }
   }
+
+  Block* optimize(Expression* curr, Expression*& child, Block* outer = nullptr, Expression** dependency1 = nullptr, Expression** dependency2 = nullptr) {
+    if (!child) return outer;
+    if (dependency1 && *dependency1 && EffectAnalyzer(*dependency1).hasSideEffects()) return outer;
+    if (dependency2 && *dependency2 && EffectAnalyzer(*dependency2).hasSideEffects()) return outer;
+    if (auto* block = child->dynCast<Block>()) {
+      if (block->list.size() >= 2) {
+        child = block->list.back();
+        if (outer == nullptr) {
+          // reuse the block, move it out
+          block->list.back() = curr;
+          block->finalize(); // last block element was our input, and is now our output, which may differ TODO optimize
+          replaceCurrent(block);
+          return block;
+        } else {
+          // append to an existing outer block
+          assert(outer->list.back() == curr);
+          outer->list.pop_back();
+          for (Index i = 0; i < block->list.size() - 1; i++) {
+            outer->list.push_back(block->list[i]);
+          }
+          outer->list.push_back(curr);
+        }
+      }
+    }
+    return outer;
+  }
+
+  void visitUnary(Unary* curr) {
+    optimize(curr, curr->value);
+  }
+  void visitSetLocal(SetLocal* curr) {
+    optimize(curr, curr->value);
+  }
+  void visitLoad(Load* curr) {
+    optimize(curr, curr->ptr);
+  }
+  void visitReturn(Return* curr) {
+    optimize(curr, curr->value);
+  }
+
+  void visitBinary(Binary* curr) {
+    optimize(curr, curr->right, optimize(curr, curr->left), &curr->left);
+  }
+  void visitStore(Store* curr) {
+    optimize(curr, curr->value, optimize(curr, curr->ptr), &curr->ptr);
+  }
+
+  void visitSelect(Select* curr) {
+    optimize(curr, curr->condition, optimize(curr, curr->ifFalse, optimize(curr, curr->ifTrue), &curr->ifTrue), &curr->ifTrue, &curr->ifFalse);
+  }
+
+  void visitBreak(Break* curr) {
+    optimize(curr, curr->condition, optimize(curr, curr->value), &curr->value);
+  }
+  void visitSwitch(Switch* curr) {
+    optimize(curr, curr->condition, optimize(curr, curr->value), &curr->value);
+  }
+
+  template<typename T>
+  void handleCall(T* curr, Block* outer = nullptr) {
+    for (Index i = 0; i < curr->operands.size(); i++) {
+      outer = optimize(curr, curr->operands[i], outer);
+      if (EffectAnalyzer(curr->operands[i]).hasSideEffects()) return;
+    }
+  }
+
+  void visitCall(Call* curr) {
+    handleCall(curr);
+  }
+
+  void visitCallImport(CallImport* curr) {
+    handleCall(curr);
+  }
+
+  void visitCallIndirect(CallIndirect* curr) {
+    auto* outer = optimize(curr, curr->target);
+    if (EffectAnalyzer(curr->target).hasSideEffects()) return;
+    handleCall(curr, outer);
+  }
 };
 
 static RegisterPass<MergeBlocks> registerPass("merge-blocks", "merges blocks to their parents");
diff --git a/src/tools/binaryen-shell.cpp b/src/tools/binaryen-shell.cpp
index 8f5fbe4e5..96f506934 100644
--- a/src/tools/binaryen-shell.cpp
+++ b/src/tools/binaryen-shell.cpp
@@ -198,50 +198,51 @@ int main(int argc, const char* argv[]) {
 
   auto input(read_file<std::vector<char>>(options.extra["infile"], Flags::Text, options.debug ? Flags::Debug : Flags::Release));
 
-  if (options.debug) std::cerr << "parsing text to s-expressions...\n";
-  SExpressionParser parser(input.data());
-  Element& root = *parser.root;
-
-  // A .wast may have multiple modules, with some asserts after them
   bool checked = false;
-  size_t i = 0;
-  while (i < root.size()) {
-    Element& curr = *root[i];
-    IString id = curr[0]->str();
-    if (id == MODULE) {
-      if (options.debug) std::cerr << "parsing s-expressions to wasm...\n";
-      Module wasm;
-      std::unique_ptr<SExpressionWasmBuilder> builder;
-      try {
+
+  try {
+    if (options.debug) std::cerr << "parsing text to s-expressions...\n";
+    SExpressionParser parser(input.data());
+    Element& root = *parser.root;
+
+    // A .wast may have multiple modules, with some asserts after them
+    size_t i = 0;
+    while (i < root.size()) {
+      Element& curr = *root[i];
+      IString id = curr[0]->str();
+      if (id == MODULE) {
+        if (options.debug) std::cerr << "parsing s-expressions to wasm...\n";
+        Module wasm;
+        std::unique_ptr<SExpressionWasmBuilder> builder;
         builder = make_unique<SExpressionWasmBuilder>(wasm, *root[i]);
-      } catch (ParseException& p) {
-        p.dump(std::cerr);
-        abort();
-      }
-      i++;
-      assert(WasmValidator().validate(wasm));
-
-      MixedArena moreModuleAllocations;
-
-      if (passes.size() > 0) {
-        if (options.debug) std::cerr << "running passes...\n";
-        PassRunner passRunner(&wasm);
-        if (options.debug) passRunner.setDebug(true);
-        for (auto& passName : passes) {
-          if (passName == "O") {
-            passRunner.addDefaultOptimizationPasses();
-          } else {
-            passRunner.add(passName);
+        i++;
+        assert(WasmValidator().validate(wasm));
+
+        MixedArena moreModuleAllocations;
+
+        if (passes.size() > 0) {
+          if (options.debug) std::cerr << "running passes...\n";
+          PassRunner passRunner(&wasm);
+          if (options.debug) passRunner.setDebug(true);
+          for (auto& passName : passes) {
+            if (passName == "O") {
+              passRunner.addDefaultOptimizationPasses();
+            } else {
+              passRunner.add(passName);
+            }
           }
+          passRunner.run();
+          assert(WasmValidator().validate(wasm));
         }
-        passRunner.run();
-        assert(WasmValidator().validate(wasm));
-      }
 
-      run_asserts(&i, &checked, &wasm, &root, &builder, entry);
-    } else {
-      run_asserts(&i, &checked, nullptr, &root, nullptr, entry);
+        run_asserts(&i, &checked, &wasm, &root, &builder, entry);
+      } else {
+        run_asserts(&i, &checked, nullptr, &root, nullptr, entry);
+      }
     }
+  } catch (ParseException& p) {
+    p.dump(std::cerr);
+    abort();
   }
 
   if (checked) {
diff --git a/src/wasm-s-parser.h b/src/wasm-s-parser.h
index dae321866..bb6fe3fd4 100644
--- a/src/wasm-s-parser.h
+++ b/src/wasm-s-parser.h
@@ -173,7 +173,7 @@ private:
         curr->list().push_back(parseString());
       }
     }
-    assert(stack.size() == 0); 
+    if (stack.size() != 0) throw ParseException("stack is not empty", curr->line, curr->col);
     return curr;
   }