selectify if-elses with no control flow in them

2 files changed, 31 insertions, 1 deletions

diff --git a/src/ast_utils.h b/src/ast_utils.h
index 36895ac51..5cda48578 100644
--- a/src/ast_utils.h
+++ b/src/ast_utils.h
@@ -52,7 +52,7 @@ struct EffectAnalyzer : public PostWalker<EffectAnalyzer, Visitor<EffectAnalyzer
 
   bool accessesLocal() { return localsRead.size() + localsWritten.size() > 0; }
   bool accessesMemory() { return calls || readsMemory || writesMemory; }
-  bool hasSideEffects() { return calls || localsWritten.size() > 0 || writesMemory; }
+  bool hasSideEffects() { return calls || localsWritten.size() > 0 || writesMemory || branches; }
   bool hasAnything() { return branches || calls || accessesLocal() || readsMemory || writesMemory; }
 
   // checks if these effects would invalidate another set (e.g., if we write, we invalidate someone that reads, they can't be moved past us)
diff --git a/src/passes/RemoveUnusedBrs.cpp b/src/passes/RemoveUnusedBrs.cpp
index 1dd16d07b..f4efef173 100644
--- a/src/passes/RemoveUnusedBrs.cpp
+++ b/src/passes/RemoveUnusedBrs.cpp
@@ -175,6 +175,36 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs, Visitor<R
       }
       flows.clear();
     } while (anotherCycle);
+    // finally, we may have simplified ifs enough to turn them into selects
+    struct Selectifier : public WalkerPass<PostWalker<Selectifier, Visitor<Selectifier>>> {
+      void visitIf(If* curr) {
+        if (curr->ifFalse) {
+          // if with else, consider turning it into a select if there is no control flow
+          // TODO: estimate cost
+          EffectAnalyzer condition;
+          condition.walk(curr->condition);
+          if (!condition.hasSideEffects()) {
+            EffectAnalyzer ifTrue;
+            ifTrue.walk(curr->ifTrue);
+            if (!ifTrue.hasSideEffects()) {
+              EffectAnalyzer ifFalse;
+              ifFalse.walk(curr->ifFalse);
+              if (!ifFalse.hasSideEffects()) {
+                auto* select = getModule()->allocator.alloc<Select>();
+                select->condition = curr->condition;
+                select->ifTrue = curr->ifTrue;
+                select->ifFalse = curr->ifFalse;
+                select->finalize();
+                replaceCurrent(select);
+              }
+            }
+          }
+        }
+      }
+    };
+    Selectifier selectifier;
+    selectifier.setModule(getModule());
+    selectifier.walk(root);
   }
 };