Simplify IfElse instructions with a single trivial block (#3114)

2 files changed, 32 insertions, 0 deletions

diff --git a/source/slang/slang-ir-eliminate-multilevel-break.cpp b/source/slang/slang-ir-eliminate-multilevel-break.cpp
index 65e388608..19c95edfd 100644
--- a/source/slang/slang-ir-eliminate-multilevel-break.cpp
+++ b/source/slang/slang-ir-eliminate-multilevel-break.cpp
@@ -193,6 +193,14 @@ struct EliminateMultiLevelBreakContext
             {
                 if (as<IRBlock>(terminator->getParent()) == block)
                 {
+                    // Don't double count instructions like
+                    // ifElse(cond, true, after, after)
+                    if(const auto ifElse = as<IRIfElse>(terminator))
+                    {
+                        if(&ifElse->afterBlock == use)
+                            continue;
+                    }
+
                     relevantUses.add(use);
                 }
             }
diff --git a/source/slang/slang-ir-simplify-cfg.cpp b/source/slang/slang-ir-simplify-cfg.cpp
index f2d0c4555..63c40b16f 100644
--- a/source/slang/slang-ir-simplify-cfg.cpp
+++ b/source/slang/slang-ir-simplify-cfg.cpp
@@ -444,6 +444,30 @@ static bool trySimplifyIfElse(IRBuilder& builder, IRIfElse* ifElseInst)
             return true;
         }
     }
+    else
+    {
+        // Otherwise, we can try to remove at least remove one of the trivial branches
+        // Remove either the true or false block if it jumps to the after block
+        // with no parameters.
+
+        const auto afterBlock = ifElseInst->getAfterBlock();
+        if(!afterBlock->getFirstParam())
+        {
+            const auto trueBlock = ifElseInst->getTrueBlock();
+            const auto falseBlock = ifElseInst->getFalseBlock();
+
+            if(isTrueBranchTrivial && trueBlock != afterBlock && !trueBlock->hasMoreThanOneUse())
+            {
+                trueBlock->replaceUsesWith(afterBlock);
+                trueBlock->removeAndDeallocate();
+            }
+            else if(isFalseBranchTrivial && falseBlock != afterBlock && !falseBlock->hasMoreThanOneUse())
+            {
+                falseBlock->replaceUsesWith(afterBlock);
+                falseBlock->removeAndDeallocate();
+            }
+        }
+    }
     return false;
 }