[InstCombine] Don't transform code if DoTransform is false

[Carrot]

Summary

In patch https://reviews.llvm.org/D72396, it doesn't check DoTransform before transforming the code, and generates wrong result for the attached test case.

[lebedev.ri]

LG
I think that function is misdesigned, but i can't suggest an easy good alternative.

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

959

Please add a comment to the function explaining this weird DoTransform behaviour.

[Carrot]

Thanks for the quick review!
Will commit this version.

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

if (SrcWidth % DestWidth == 0) {
       Value *BitCast = Builder.CreateBitCast(VecOp, BitCastTo);
       return ExtractElementInst::Create(BitCast, Builder.getInt32(NewIdx));
     }
   }
 
   return nullptr;
 }
 
+/// Transform (zext icmp) to bitwise / integer operations in order to
+/// eliminate it. If DoTransform is false, just test whether the given
+/// (zext icmp) can be transformed.
 Instruction *InstCombinerImpl::transformZExtICmp(ICmpInst *Cmp, ZExtInst &Zext,
                                                  bool DoTransform) {
   // If we are just checking for a icmp eq of a single bit and zext'ing it
   // to an integer, then shift the bit to the appropriate place and then
   // cast to integer to avoid the comparison.
   const APInt *Op1CV;
   if (match(Cmp->getOperand(1), m_APInt(Op1CV))) {
 
     // zext (x <s  0) to i32 --> x>>u31      true if signbit set.
Instruction *InstCombinerImpl::transformZExtICmp(ICmpInst *Cmp, ZExtInst &Zext,
   if (Cmp->isEquality() && Zext.getType() == Cmp->getOperand(0)->getType()) {
     // Test if a bit is clear/set using a shifted-one mask:
     // zext (icmp eq (and X, (1 << ShAmt)), 0) --> and (lshr (not X), ShAmt), 1
     // zext (icmp ne (and X, (1 << ShAmt)), 0) --> and (lshr X, ShAmt), 1
     Value *X, *ShAmt;
     if (Cmp->hasOneUse() && match(Cmp->getOperand(1), m_ZeroInt()) &&
         match(Cmp->getOperand(0),
               m_OneUse(m_c_And(m_Shl(m_One(), m_Value(ShAmt)), m_Value(X))))) {
+      if (!DoTransform)
+        return Cmp;
+
       if (Cmp->getPredicate() == ICmpInst::ICMP_EQ)
         X = Builder.CreateNot(X);
       Value *Lshr = Builder.CreateLShr(X, ShAmt);
       Value *And1 = Builder.CreateAnd(Lshr, ConstantInt::get(X->getType(), 1));
       return replaceInstUsesWith(Zext, And1);
     }
 
     // icmp ne A, B is equal to xor A, B when A and B only really have one bit.

llvm/test/Transforms/InstCombine/zext.ll

 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %sh1 = shl i32 1, %y
   %and = and i32 %sh1, %x
   %cmp = icmp sgt i32 %and, 0
   %r = zext i1 %cmp to i32
   ret i32 %r
 }
+
+; Assert that zext(or(masked_bit_test, icmp)) can be correctly transformed to
+; or(shifted_masked_bit, zext(icmp))
+
+define void @zext_or_masked_bit_test(i32 %a, i32 %b, i32* %p) {
+; CHECK-LABEL: @zext_or_masked_bit_test
+; CHECK-NEXT:    [[LD:%.*]] = load i32, i32* %p, align 4
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[LD]], %b
+; CHECK-NEXT:    [[SHR:%.*]] = lshr i32 %a, %b
+; CHECK-NEXT:    [[AND:%.*]] = and i32 [[SHR]], 1
+; CHECK-NEXT:    [[EXT:%.*]] = zext i1 [[CMP]] to i32
+; CHECK-NEXT:    [[OR:%.*]] = or i32 [[AND]], [[EXT]]
+; CHECK-NEXT:    store i32 [[OR]], i32* %p, align 4
+; CHECK-NEXT:    ret void
+;
+  %ld = load i32, i32* %p, align 4
+  %shl = shl i32 1, %b
+  %and = and i32 %shl, %a
+  %tobool = icmp ne i32 %and, 0
+  %cmp = icmp eq i32 %ld, %b
+  %or = or i1 %tobool, %cmp
+  %conv = zext i1 %or to i32
+  store i32 %conv, i32* %p, align 4
+  ret void
+}
+