Refactor MergeBlocks to use iteration; adds Wasm GC support (#4137)

MergeBlocks was written a very long time ago, before the
iteration API, so it had a bunch of hardcoded things for
specific instructions. In particular, that did not handle GC.
This does a small refactoring to use iteration. The refactoring
is NFC, but while doing so it adds support for new relevant
instructions, including wasm GC.

Author: kripken

src/passes/MergeBlocks.cpp

@@ -74,6 +74,7 @@
 
 #include <ir/branch-utils.h>
 #include <ir/effects.h>
+#include <ir/iteration.h>
 #include <ir/utils.h>
 #include <pass.h>
 #include <wasm-builder.h>
@@ -397,7 +398,9 @@ void BreakValueDropper::visitBlock(Block* curr) {
   optimizeBlock(curr, getModule(), passOptions, branchInfo);
 }
 
-struct MergeBlocks : public WalkerPass<PostWalker<MergeBlocks>> {
+struct MergeBlocks
+  : public WalkerPass<
+      PostWalker<MergeBlocks, UnifiedExpressionVisitor<MergeBlocks>>> {
   bool isFunctionParallel() override { return true; }
 
   Pass* create() override { return new MergeBlocks; }
@@ -491,20 +494,26 @@ struct MergeBlocks : public WalkerPass<PostWalker<MergeBlocks>> {
     return outer;
   }
 
-  void visitUnary(Unary* curr) { optimize(curr, curr->value); }
-  void visitLocalSet(LocalSet* curr) { optimize(curr, curr->value); }
-  void visitLoad(Load* curr) { optimize(curr, curr->ptr); }
-  void visitReturn(Return* curr) { optimize(curr, curr->value); }
+  // Default optimizations for simple cases. Complex things are overridden
+  // below.
+  void visitExpression(Expression* curr) {
+    // Control flow need special handling. Those we can optimize are handled
+    // below.
+    if (Properties::isControlFlowStructure(curr)) {
+      return;
+    }
 
-  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 visitAtomicRMW(AtomicRMW* curr) {
-    optimize(curr, curr->value, optimize(curr, curr->ptr), &curr->ptr);
+    ChildIterator iterator(curr);
+    auto& children = iterator.children;
+    if (children.size() == 1) {
+      optimize(curr, *children[0]);
+    } else if (children.size() == 2) {
+      optimize(curr, *children[0], optimize(curr, *children[1]), children[1]);
+    } else if (children.size() == 3) {
+      optimizeTernary(curr, *children[2], *children[1], *children[0]);
+    }
   }
+
   void optimizeTernary(Expression* curr,
                        Expression*& first,
                        Expression*& second,
@@ -528,23 +537,6 @@ struct MergeBlocks : public WalkerPass<PostWalker<MergeBlocks>> {
     }
     optimize(curr, third, outer);
   }
-  void visitAtomicCmpxchg(AtomicCmpxchg* curr) {
-    optimizeTernary(curr, curr->ptr, curr->expected, curr->replacement);
-  }
-
-  void visitSelect(Select* curr) {
-    optimizeTernary(curr, curr->ifTrue, curr->ifFalse, curr->condition);
-  }
-
-  void visitDrop(Drop* curr) { optimize(curr, curr->value); }
-
-  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;

test/lit/passes/merge-blocks.wast

@@ -1,9 +1,16 @@
 ;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
-;; RUN: wasm-opt %s --merge-blocks -all -S -o - \
+;; RUN: wasm-opt %s --remove-unused-names --merge-blocks -all -S -o - \
 ;; RUN:   | filecheck %s
+;;
+;; --remove-unused-names lets --merge-blocks assume blocks without names have no
+;; branch targets.
 
 (module
  (type $anyref_=>_none (func (param anyref)))
+
+ ;; CHECK:      (type $struct (struct (field (mut i32))))
+ (type $struct (struct (field (mut i32))))
+
  ;; CHECK:      (func $br_on_to_drop
  ;; CHECK-NEXT:  (nop)
  ;; CHECK-NEXT:  (drop
@@ -31,4 +38,56 @@
    )
   )
  )
+
+ ;; CHECK:      (func $struct.set
+ ;; CHECK-NEXT:  (nop)
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (i32.const 1234)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (struct.set $struct 0
+ ;; CHECK-NEXT:   (ref.null $struct)
+ ;; CHECK-NEXT:   (i32.const 5)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (nop)
+ ;; CHECK-NEXT: )
+ (func $struct.set
+  (block
+   (nop)
+   (struct.set $struct 0
+    (block (result (ref null $struct))
+     (drop (i32.const 1234))
+     (ref.null $struct)
+    )
+    (i32.const 5)
+   )
+   (nop)
+  )
+ )
+
+ ;; CHECK:      (func $struct.get
+ ;; CHECK-NEXT:  (nop)
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (i32.const 1234)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (struct.get $struct 0
+ ;; CHECK-NEXT:    (ref.null $struct)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (nop)
+ ;; CHECK-NEXT: )
+ (func $struct.get
+  (block
+   (nop)
+   (drop
+    (struct.get $struct 0
+     (block (result (ref null $struct))
+      (drop (i32.const 1234))
+      (ref.null $struct)
+     )
+    )
+   )
+   (nop)
+  )
+ )
 )