3 files changed, 69 insertions, 26 deletions

diff --git a/llvm/lib/Analysis/GuardUtils.cpp b/llvm/lib/Analysis/GuardUtils.cpp
index 6dc2b740ac9..d4a86e18997 100644
--- a/llvm/lib/Analysis/GuardUtils.cpp
+++ b/llvm/lib/Analysis/GuardUtils.cpp
@@ -13,9 +13,9 @@
 #include "llvm/IR/PatternMatch.h"
 
 using namespace llvm;
+using namespace llvm::PatternMatch;
 
 bool llvm::isGuard(const User *U) {
-  using namespace llvm::PatternMatch;
   return match(U, m_Intrinsic<Intrinsic::experimental_guard>());
 }
 
@@ -32,7 +32,6 @@ bool llvm::isGuardAsWidenableBranch(const User *U) {
   if (!parseWidenableBranch(U, Condition, WidenableCondition, GuardedBB,
                             DeoptBB))
     return false;
-  using namespace llvm::PatternMatch;
   for (auto &Insn : *DeoptBB) {
     if (match(&Insn, m_Intrinsic<Intrinsic::experimental_deoptimize>()))
       return true;
@@ -45,17 +44,32 @@ bool llvm::isGuardAsWidenableBranch(const User *U) {
 bool llvm::parseWidenableBranch(const User *U, Value *&Condition,
                                 Value *&WidenableCondition,
                                 BasicBlock *&IfTrueBB, BasicBlock *&IfFalseBB) {
-  using namespace llvm::PatternMatch;
+  if (match(U, m_Br(m_Intrinsic<Intrinsic::experimental_widenable_condition>(),
+                    IfTrueBB, IfFalseBB)) &&
+      cast<BranchInst>(U)->getCondition()->hasOneUse()) {
+    WidenableCondition = cast<BranchInst>(U)->getCondition();
+    Condition = ConstantInt::getTrue(IfTrueBB->getContext());
+    return true;
+  }
+
+  // Check for two cases:
+  // 1) br (i1 (and A, WC())), label %IfTrue, label %IfFalse
+  // 2) br (i1 (and WC(), B)), label %IfTrue, label %IfFalse
+  // We do not check for more generalized and trees as we should canonicalize
+  // to the form above in instcombine. (TODO)
   if (!match(U, m_Br(m_And(m_Value(Condition), m_Value(WidenableCondition)),
                      IfTrueBB, IfFalseBB)))
     return false;
+  if (!match(WidenableCondition,
+             m_Intrinsic<Intrinsic::experimental_widenable_condition>())) {
+    if (!match(Condition,
+               m_Intrinsic<Intrinsic::experimental_widenable_condition>()))
+      return false;
+    std::swap(Condition, WidenableCondition);
+  }
+    
   // For the branch to be (easily) widenable, it must not correlate with other
   // branches.  Thus, the widenable condition must have a single use.
-  if (!WidenableCondition->hasOneUse() ||
-      !cast<BranchInst>(U)->getCondition()->hasOneUse())
-    return false;
-  // TODO: At the moment, we only recognize the branch if the WC call in this
-  // specific position.  We should generalize!
-  return match(WidenableCondition,
-               m_Intrinsic<Intrinsic::experimental_widenable_condition>());
+  return (WidenableCondition->hasOneUse() &&
+          cast<BranchInst>(U)->getCondition()->hasOneUse());
 }
diff --git a/llvm/lib/Transforms/Scalar/GuardWidening.cpp b/llvm/lib/Transforms/Scalar/GuardWidening.cpp
index 69d9507b356..a3eba27a4d9 100644
--- a/llvm/lib/Transforms/Scalar/GuardWidening.cpp
+++ b/llvm/lib/Transforms/Scalar/GuardWidening.cpp
@@ -84,15 +84,16 @@ static Value *getCondition(Instruction *I) {
            "Bad guard intrinsic?");
     return GI->getArgOperand(0);
   }
-  if (isGuardAsWidenableBranch(I)) {
-    auto *Cond = cast<BranchInst>(I)->getCondition();
-    return cast<BinaryOperator>(Cond)->getOperand(0);
-  }
+  Value *Cond, *WC;
+  BasicBlock *IfTrueBB, *IfFalseBB;
+  if (parseWidenableBranch(I, Cond, WC, IfTrueBB, IfFalseBB))
+    return Cond;
+
   return cast<BranchInst>(I)->getCondition();
 }
 
 // Set the condition for \p I to \p NewCond. \p I can either be a guard or a
-// conditional branch.
+// conditional branch.  
 static void setCondition(Instruction *I, Value *NewCond) {
   if (IntrinsicInst *GI = dyn_cast<IntrinsicInst>(I)) {
     assert(GI->getIntrinsicID() == Intrinsic::experimental_guard &&
diff --git a/llvm/lib/Transforms/Utils/GuardUtils.cpp b/llvm/lib/Transforms/Utils/GuardUtils.cpp
index b8c4c764e74..241bfbf80bd 100644
--- a/llvm/lib/Transforms/Utils/GuardUtils.cpp
+++ b/llvm/lib/Transforms/Utils/GuardUtils.cpp
@@ -15,10 +15,12 @@
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/MDBuilder.h"
+#include "llvm/IR/PatternMatch.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 
 using namespace llvm;
+using namespace llvm::PatternMatch;
 
 static cl::opt<uint32_t> PredicatePassBranchWeight(
     "guards-predicate-pass-branch-weight", cl::Hidden, cl::init(1 << 20),
@@ -80,23 +82,49 @@ void llvm::makeGuardControlFlowExplicit(Function *DeoptIntrinsic,
 void llvm::widenWidenableBranch(BranchInst *WidenableBR, Value *NewCond) {
   assert(isWidenableBranch(WidenableBR) && "precondition");
 
-  Instruction *WCAnd = cast<Instruction>(WidenableBR->getCondition());
-  // Condition is only guaranteed to dominate branch
-  WCAnd->moveBefore(WidenableBR);
-  Value *OldCond = WCAnd->getOperand(0);
-  IRBuilder<> B(WCAnd);
-  WCAnd->setOperand(0, B.CreateAnd(NewCond, OldCond));
+  // The tempting trivially option is to produce something like this:
+  // br (and oldcond, newcond) where oldcond is assumed to contain a widenable
+  // condition, but that doesn't match the pattern parseWidenableBranch expects
+  // so we have to be more sophisticated.
 
+  if (match(WidenableBR->getCondition(),
+            m_Intrinsic<Intrinsic::experimental_widenable_condition>())) {
+    IRBuilder<> B(WidenableBR);
+    WidenableBR->setCondition(B.CreateAnd(NewCond,
+                                          WidenableBR->getCondition()));
+  } else {
+    Instruction *WCAnd = cast<Instruction>(WidenableBR->getCondition());
+    // Condition is only guaranteed to dominate branch
+    WCAnd->moveBefore(WidenableBR);
+    IRBuilder<> B(WCAnd);
+    const bool Op0IsWC =
+      match(WCAnd->getOperand(0),
+            m_Intrinsic<Intrinsic::experimental_widenable_condition>());
+    const unsigned CondOpIdx = Op0IsWC ? 1 : 0;
+    Value *OldCond = WCAnd->getOperand(CondOpIdx);
+    NewCond = B.CreateAnd(NewCond, OldCond);
+    WCAnd->setOperand(CondOpIdx, NewCond);
+  }
   assert(isWidenableBranch(WidenableBR) && "preserve widenabiliy");
 }
 
 void llvm::setWidenableBranchCond(BranchInst *WidenableBR, Value *NewCond) {
   assert(isWidenableBranch(WidenableBR) && "precondition");
 
-  Instruction *WCAnd = cast<Instruction>(WidenableBR->getCondition());
-  // Condition is only guaranteed to dominate branch
-  WCAnd->moveBefore(WidenableBR);
-  WCAnd->setOperand(0, NewCond);
-
+  if (match(WidenableBR->getCondition(),
+            m_Intrinsic<Intrinsic::experimental_widenable_condition>())) {
+    IRBuilder<> B(WidenableBR);
+    WidenableBR->setCondition(B.CreateAnd(NewCond,
+                                          WidenableBR->getCondition()));
+  } else {
+    Instruction *WCAnd = cast<Instruction>(WidenableBR->getCondition());
+    // Condition is only guaranteed to dominate branch
+    WCAnd->moveBefore(WidenableBR);
+    const bool Op0IsWC =
+      match(WCAnd->getOperand(0),
+            m_Intrinsic<Intrinsic::experimental_widenable_condition>());
+    const unsigned CondOpIdx = Op0IsWC ? 1 : 0;
+    WCAnd->setOperand(CondOpIdx, NewCond);
+  }
   assert(isWidenableBranch(WidenableBR) && "preserve widenabiliy");
 }