Review comments + fix some conditions

Change-Id: I4b71adfd8bffdda4d2b0d1cba85a3fd73a105a28

Author: jrbyrnes

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

@@ -3593,10 +3593,16 @@ static Value *foldOrOfInversions(BinaryOperator &I,
   return nullptr;
 }
 
+// A decomposition of ((A & N) ? 0 : N * C) . Where X = A, Factor = C, Mask = N.
+// The NUW / NSW bools
+// Note that we can decompose equivalent forms of this expression (e.g. ((A & N)
+// * C))
 struct DecomposedBitMaskMul {
   Value *X;
   APInt Factor;
   APInt Mask;
+  bool NUW;
+  bool NSW;
 };
 
 static std::optional<DecomposedBitMaskMul> matchBitmaskMul(Value *V) {
@@ -3606,28 +3612,29 @@ static std::optional<DecomposedBitMaskMul> matchBitmaskMul(Value *V) {
 
   Value *MulOp = nullptr;
   const APInt *MulConst = nullptr;
+
+  // Decompose (A & N) * C) into BitMaskMul
   if (match(Op, m_Mul(m_Value(MulOp), m_APInt(MulConst)))) {
     Value *Original = nullptr;
     const APInt *Mask = nullptr;
-    if (!MulConst->isStrictlyPositive())
+    if (MulConst->isZero())
       return std::nullopt;
 
     if (match(MulOp, m_And(m_Value(Original), m_APInt(Mask)))) {
-      if (!Mask->isStrictlyPositive())
+      if (Mask->isZero())
         return std::nullopt;
-      DecomposedBitMaskMul Ret;
-      Ret.X = Original;
-      Ret.Mask = *Mask;
-      Ret.Factor = *MulConst;
-      return Ret;
+      return std::optional<DecomposedBitMaskMul>(
+          {Original, *MulConst, *Mask,
+           cast<BinaryOperator>(Op)->hasNoUnsignedWrap(),
+           cast<BinaryOperator>(Op)->hasNoSignedWrap()});
     }
     return std::nullopt;
   }
 
   Value *Cond = nullptr;
   const APInt *EqZero = nullptr, *NeZero = nullptr;
 
-  //  (!(A & N) ? 0 : N * C) + (!(A & M) ? 0 : M * C) -> A & (N + M) * C
+  // Decompose ((A & N) ? 0 : N * C) into BitMaskMul
   if (match(Op, m_Select(m_Value(Cond), m_APInt(EqZero), m_APInt(NeZero)))) {
     auto ICmpDecompose =
         decomposeBitTest(Cond, /*LookThruTrunc=*/true,
@@ -3638,22 +3645,20 @@ static std::optional<DecomposedBitMaskMul> matchBitmaskMul(Value *V) {
     if (ICmpDecompose->Pred == ICmpInst::ICMP_NE)
       std::swap(EqZero, NeZero);
 
-    if (!EqZero->isZero() || !NeZero->isStrictlyPositive())
+    if (!EqZero->isZero() || NeZero->isZero())
       return std::nullopt;
 
     if (!ICmpInst::isEquality(ICmpDecompose->Pred) ||
         !ICmpDecompose->C.isZero() || !ICmpDecompose->Mask.isPowerOf2() ||
-        ICmpDecompose->Mask.isNegative())
+        ICmpDecompose->Mask.isZero())
       return std::nullopt;
 
     if (!NeZero->urem(ICmpDecompose->Mask).isZero())
       return std::nullopt;
 
-    DecomposedBitMaskMul Ret;
-    Ret.X = ICmpDecompose->X;
-    Ret.Mask = ICmpDecompose->Mask;
-    Ret.Factor = NeZero->udiv(ICmpDecompose->Mask);
-    return Ret;
+    return std::optional<DecomposedBitMaskMul>(
+        {ICmpDecompose->X, NeZero->udiv(ICmpDecompose->Mask),
+         ICmpDecompose->Mask, /*NUW=*/false, /*NSW=*/false});
   }
 
   return std::nullopt;
@@ -3741,19 +3746,26 @@ Instruction *InstCombinerImpl::visitOr(BinaryOperator &I) {
                                    /*NSW=*/true, /*NUW=*/true))
       return R;
 
-    auto Decomp0 = matchBitmaskMul(I.getOperand(0));
+    // (!(A & N) ? 0 : N * C) + (!(A & M) ? 0 : M * C) -> A & (N + M) * C
+    // This also accepts the equivalent mul form of (A & N) ? 0 : N * C)
+    // expressions i.e. (A & N) * C
     auto Decomp1 = matchBitmaskMul(I.getOperand(1));
-
-    if (Decomp0 && Decomp1) {
-      if (Decomp0->X == Decomp1->X &&
+    if (Decomp1) {
+      auto Decomp0 = matchBitmaskMul(I.getOperand(0));
+      if (Decomp0 && Decomp0->X == Decomp1->X &&
           (Decomp0->Mask & Decomp1->Mask).isZero() &&
           Decomp0->Factor == Decomp1->Factor) {
+
         auto NewAnd = Builder.CreateAnd(
             Decomp0->X, ConstantInt::get(Decomp0->X->getType(),
                                          (Decomp0->Mask + Decomp1->Mask)));
 
-        return BinaryOperator::CreateMul(
+        auto Combined = BinaryOperator::CreateMul(
             NewAnd, ConstantInt::get(NewAnd->getType(), Decomp1->Factor));
+
+        Combined->setHasNoUnsignedWrap(Decomp0->NUW && Decomp1->NUW);
+        Combined->setHasNoSignedWrap(Decomp0->NSW && Decomp1->NSW);
+        return Combined;
       }
     }
   }

llvm/test/Transforms/InstCombine/or-bitmask.ll

@@ -37,8 +37,8 @@ define i32 @add_select_cmp_and2(i32 %in) {
 define i32 @add_select_cmp_and3(i32 %in) {
 ; CHECK-LABEL: @add_select_cmp_and3(
 ; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[IN:%.*]], 7
-; CHECK-NEXT:    [[TEMP:%.*]] = mul nuw nsw i32 [[TMP1]], 72
-; CHECK-NEXT:    ret i32 [[TEMP]]
+; CHECK-NEXT:    [[TEMP1:%.*]] = mul nuw nsw i32 [[TMP1]], 72
+; CHECK-NEXT:    ret i32 [[TEMP1]]
 ;
   %bitop0 = and i32 %in, 1
   %cmp0 = icmp eq i32 %bitop0, 0
@@ -57,8 +57,8 @@ define i32 @add_select_cmp_and3(i32 %in) {
 define i32 @add_select_cmp_and4(i32 %in) {
 ; CHECK-LABEL: @add_select_cmp_and4(
 ; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[IN:%.*]], 15
-; CHECK-NEXT:    [[TEMP3:%.*]] = mul nuw nsw i32 [[TMP2]], 72
-; CHECK-NEXT:    ret i32 [[TEMP3]]
+; CHECK-NEXT:    [[TEMP2:%.*]] = mul nuw nsw i32 [[TMP2]], 72
+; CHECK-NEXT:    ret i32 [[TEMP2]]
 ;
   %bitop0 = and i32 %in, 1
   %cmp0 = icmp eq i32 %bitop0, 0