From f6b60b3b79606612e9df6b3ab8d4367ca673fedc Mon Sep 17 00:00:00 2001
From: Roman Lebedev <lebedev.ri@gmail.com>
Date: Fri, 25 Mar 2022 15:11:49 +0300
Subject: [PATCH] [SimplifyCFG] `FoldBranchToCommonDest()`: allow
 branch-on-select

This whole check is bogus, it's some kind of a profitability check.
For now, simply extend it to not only allow branch-on-binary-ops,
but also on poison-safe logic ops.

Refs. https://github.com/llvm/llvm-project/issues/53861
Refs. https://github.com/llvm/llvm-project/issues/54553
---
 llvm/lib/Transforms/Utils/SimplifyCFG.cpp                    |  4 +++-
 .../Transforms/SimplifyCFG/fold-branch-to-common-dest.ll     | 12 +++---------
 2 files changed, 6 insertions(+), 10 deletions(-)

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index d48c52d..29c5936 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -3531,7 +3531,9 @@ bool llvm::FoldBranchToCommonDest(BranchInst *BI, DomTreeUpdater *DTU,
 
   Instruction *Cond = dyn_cast<Instruction>(BI->getCondition());
 
-  if (!Cond || (!isa<CmpInst>(Cond) && !isa<BinaryOperator>(Cond)) ||
+  if (!Cond ||
+      (!isa<CmpInst>(Cond) && !isa<BinaryOperator>(Cond) &&
+       !isa<SelectInst>(Cond)) ||
       Cond->getParent() != BB || !Cond->hasOneUse())
     return false;
 
diff --git a/llvm/test/Transforms/SimplifyCFG/fold-branch-to-common-dest.ll b/llvm/test/Transforms/SimplifyCFG/fold-branch-to-common-dest.ll
index a3ba53e..d594596 100644
--- a/llvm/test/Transforms/SimplifyCFG/fold-branch-to-common-dest.ll
+++ b/llvm/test/Transforms/SimplifyCFG/fold-branch-to-common-dest.ll
@@ -1074,22 +1074,16 @@ define i32 @firewall(i8* %data) {
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[I]], 17
 ; CHECK-NEXT:    [[CMP3:%.*]] = icmp eq i16 [[I2]], 1
 ; CHECK-NEXT:    [[OR_COND:%.*]] = select i1 [[CMP]], i1 [[CMP3]], i1 false
-; CHECK-NEXT:    br i1 [[OR_COND]], label [[CLEANUP:%.*]], label [[IF_END:%.*]]
-; CHECK:       if.end:
 ; CHECK-NEXT:    [[CMP10:%.*]] = icmp eq i16 [[I2]], 2
 ; CHECK-NEXT:    [[OR_COND33:%.*]] = select i1 [[CMP]], i1 [[CMP10]], i1 false
-; CHECK-NEXT:    br i1 [[OR_COND33]], label [[CLEANUP]], label [[IF_END13:%.*]]
-; CHECK:       if.end13:
+; CHECK-NEXT:    [[OR_COND1:%.*]] = select i1 [[OR_COND]], i1 true, i1 [[OR_COND33]]
 ; CHECK-NEXT:    [[CMP19:%.*]] = icmp eq i16 [[I2]], 3
 ; CHECK-NEXT:    [[OR_COND34:%.*]] = select i1 [[CMP]], i1 [[CMP19]], i1 false
-; CHECK-NEXT:    br i1 [[OR_COND34]], label [[CLEANUP]], label [[IF_END22:%.*]]
-; CHECK:       if.end22:
+; CHECK-NEXT:    [[OR_COND2:%.*]] = select i1 [[OR_COND1]], i1 true, i1 [[OR_COND34]]
 ; CHECK-NEXT:    [[CMP28:%.*]] = icmp eq i16 [[I2]], 4
 ; CHECK-NEXT:    [[OR_COND35:%.*]] = select i1 [[CMP]], i1 [[CMP28]], i1 false
 ; CHECK-NEXT:    [[DOT:%.*]] = zext i1 [[OR_COND35]] to i32
-; CHECK-NEXT:    br label [[CLEANUP]]
-; CHECK:       cleanup:
-; CHECK-NEXT:    [[RETVAL_0:%.*]] = phi i32 [ 1, [[ENTRY:%.*]] ], [ 1, [[IF_END]] ], [ 1, [[IF_END13]] ], [ [[DOT]], [[IF_END22]] ]
+; CHECK-NEXT:    [[RETVAL_0:%.*]] = select i1 [[OR_COND2]], i32 1, i32 [[DOT]]
 ; CHECK-NEXT:    ret i32 [[RETVAL_0]]
 ;
 entry:
-- 
2.7.4