Fix a merge-blocks fuzz bug (#1755)

* Moving blocks into if arms may change the block type, and the code we had was written under the assumption that was not true.

* Move block sinking merge-blocks => remove-unused-brs, as it's more natural there. that pass refinalizes everything anyhow

2 files changed, 73 insertions, 58 deletions

diff --git a/src/passes/MergeBlocks.cpp b/src/passes/MergeBlocks.cpp
index 4999fd234..543c7f89c 100644
--- a/src/passes/MergeBlocks.cpp
+++ b/src/passes/MergeBlocks.cpp
@@ -335,62 +335,6 @@ struct MergeBlocks : public WalkerPass<PostWalker<MergeBlocks>> {
   Pass* create() override { return new MergeBlocks; }
 
   void visitBlock(Block *curr) {
-    // If the block has a single child which is a loop, and the block is named,
-    // then it is the exit for the loop. It's better to move it into the loop,
-    // where it can be better optimized by other passes.
-    // Similar logic for ifs: if the block is an exit for the if, we can
-    // move the block in, consider for example:
-    //    (block $label
-    //     (if (..condition1..)
-    //      (block
-    //       (br_if $label (..condition2..))
-    //       (..code..)
-    //      )
-    //     )
-    //    )
-    // After also merging the blocks, we have
-    //    (if (..condition1..)
-    //     (block $label
-    //      (br_if $label (..condition2..))
-    //      (..code..)
-    //     )
-    //    )
-    // which can be further optimized by other passes.
-    if (curr->name.is() && curr->list.size() == 1) {
-      if (auto* loop = curr->list[0]->dynCast<Loop>()) {
-        curr->list[0] = loop->body;
-        loop->body = curr;
-        auto oldOuterType = curr->type;
-        curr->finalize(curr->type);
-        loop->finalize();
-        // After the flip, the outer type must be the same
-        assert(loop->type == oldOuterType);
-        replaceCurrent(loop);
-      } else if (auto* iff = curr->list[0]->dynCast<If>()) {
-        // The label can't be used in the condition.
-        if (BranchUtils::BranchSeeker::countNamed(iff->condition, curr->name) == 0) {
-          // We can move the block into either arm, if there are no uses in the other.
-          Expression** target = nullptr;
-          if (!iff->ifFalse ||
-              BranchUtils::BranchSeeker::countNamed(iff->ifFalse, curr->name) == 0) {
-            target = &iff->ifTrue;
-          } else if (BranchUtils::BranchSeeker::countNamed(iff->ifTrue, curr->name) == 0) {
-            target = &iff->ifFalse;
-          }
-          if (target) {
-            curr->list[0] = *target;
-            *target = curr;
-            curr->finalize(curr->type);
-            iff->finalize();
-            replaceCurrent(iff);
-            // Note that the type might change, e.g. if the if condition is unreachable
-            // but the block that was on the outside had a break.
-          }
-        }
-      }
-      // Always fall through to optimize the block, which has a new child now.
-    }
-    // Otherwise, do the main merging optimizations.
     optimizeBlock(curr, getModule(), getPassOptions());
   }
 
diff --git a/src/passes/RemoveUnusedBrs.cpp b/src/passes/RemoveUnusedBrs.cpp
index ccdf55794..eb81b0a5a 100644
--- a/src/passes/RemoveUnusedBrs.cpp
+++ b/src/passes/RemoveUnusedBrs.cpp
@@ -102,9 +102,8 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
         // if without else stops the flow of values
         self->stopValueFlow();
       }
-    } else if (curr->is<Block>()) {
+    } else if (auto* block = curr->dynCast<Block>()) {
       // any breaks flowing to here are unnecessary, as we get here anyhow
-      auto* block = curr->cast<Block>();
       auto name = block->name;
       if (name.is()) {
         size_t size = flows.size();
@@ -436,6 +435,76 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
     }
   }
 
+  void sinkBlocks(Function* func) {
+    struct Sinker : public PostWalker<Sinker> {
+      bool worked = false;
+
+      void visitBlock(Block* curr) {
+        // If the block has a single child which is a loop, and the block is named,
+        // then it is the exit for the loop. It's better to move it into the loop,
+        // where it can be better optimized by other passes.
+        // Similar logic for ifs: if the block is an exit for the if, we can
+        // move the block in, consider for example:
+        //    (block $label
+        //     (if (..condition1..)
+        //      (block
+        //       (br_if $label (..condition2..))
+        //       (..code..)
+        //      )
+        //     )
+        //    )
+        // After also merging the blocks, we have
+        //    (if (..condition1..)
+        //     (block $label
+        //      (br_if $label (..condition2..))
+        //      (..code..)
+        //     )
+        //    )
+        // which can be further optimized later.
+        if (curr->name.is() && curr->list.size() == 1) {
+          if (auto* loop = curr->list[0]->dynCast<Loop>()) {
+            curr->list[0] = loop->body;
+            loop->body = curr;
+            curr->finalize(curr->type);
+            loop->finalize();
+            replaceCurrent(loop);
+            worked = true;
+          } else if (auto* iff = curr->list[0]->dynCast<If>()) {
+            // The label can't be used in the condition.
+            if (BranchUtils::BranchSeeker::countNamed(iff->condition, curr->name) == 0) {
+              // We can move the block into either arm, if there are no uses in the other.
+              Expression** target = nullptr;
+              if (!iff->ifFalse ||
+                  BranchUtils::BranchSeeker::countNamed(iff->ifFalse, curr->name) == 0) {
+                target = &iff->ifTrue;
+              } else if (BranchUtils::BranchSeeker::countNamed(iff->ifTrue, curr->name) == 0) {
+                target = &iff->ifFalse;
+              }
+              if (target) {
+                curr->list[0] = *target;
+                *target = curr;
+                // The block used to contain the if, and may have changed type from unreachable
+                // to none, for example, if the if has an unreachable condition but the arm
+                // is not unreachable.
+                curr->finalize();
+                iff->finalize();
+                replaceCurrent(iff);
+                worked = true;
+                // Note that the type might change, e.g. if the if condition is unreachable
+                // but the block that was on the outside had a break.
+              }
+            }
+          }
+        }
+      }
+    } sinker;
+
+    sinker.doWalkFunction(func);
+    if (sinker.worked) {
+      anotherCycle = true;
+    }
+  }
+
   void doWalkFunction(Function* func) {
     // multiple cycles may be needed
     bool worked = false;
@@ -463,6 +532,8 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
         anotherCycle |= optimizeLoop(loop);
       }
       loops.clear();
+      // sink blocks
+      sinkBlocks(func);
       if (anotherCycle) worked = true;
     } while (anotherCycle);