optimize if_else in RemoveUnusedBrs

1 files changed, 33 insertions, 8 deletions

diff --git a/src/passes/RemoveUnusedBrs.cpp b/src/passes/RemoveUnusedBrs.cpp
index 5f3018c5f..0ff48381f 100644
--- a/src/passes/RemoveUnusedBrs.cpp
+++ b/src/passes/RemoveUnusedBrs.cpp
@@ -27,14 +27,39 @@ struct RemoveUnusedBrs : public Pass {
   void visitBlock(Block *curr) override {
     if (curr->name.isNull()) return;
     if (curr->list.size() == 0) return;
-    Break* last = curr->list.back()->dyn_cast<Break>();
-    if (!last) return;
-    if (last->condition) return;
-    if (last->name == curr->name) {
-      if (!last->value) {
-        curr->list.pop_back();
-      } else {
-        curr->list[curr->list.size()-1] = last->value; // can replace with the value
+    Expression* last = curr->list.back();
+    if (Break* br = last->dyn_cast<Break>()) {
+      if (br->condition) return;
+      if (br->name == curr->name) {
+        if (!br->value) {
+          curr->list.pop_back();
+        } else {
+          curr->list[curr->list.size()-1] = br->value; // can replace with the value
+        }
+      }
+    } else if (If* ifelse = curr->list.back()->dyn_cast<If>()) {
+      if (!ifelse->ifFalse) return;
+      if (ifelse->type != none) return;
+      // an if_else that indirectly returns a value by breaking to this block can potentially remove both breaks
+      auto process = [&curr](Expression *side, bool doIt) {
+        Block* b = side->dyn_cast<Block>();
+        if (!b) return false;
+        Expression* last = b->list.back();
+        Break* br = last->dyn_cast<Break>();
+        if (!br) return false;
+        if (br->condition) return false;
+        if (!br->value) return false;
+        if (br->name != curr->name) return false;
+        if (doIt) {
+          b->list[b->list.size()-1] = br->value;
+        }
+        return true;
+      };
+      // do both, or none
+      if (process(ifelse->ifTrue, false) && process(ifelse->ifFalse, false)) {
+        process(ifelse->ifTrue, true);
+        process(ifelse->ifFalse, true);
+        ifelse->type = curr->type; // if_else now returns a value
       }
     }
   }