[LVI] Get value range from mask comparison

InstCombine likes to canonicalize comparisons of the form
X == C || X == C+1 into (X & -2) == C'. Make sure LVI can still
recover the value range from this. Can of course also be useful
for proper mask comparisons.

For the sake of clarity, the implementation goes through KnownBits
to compute the range.

llvm/lib/Analysis/LazyValueInfo.cpp

@@ -36,6 +36,7 @@
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/FormattedStream.h"
+#include "llvm/Support/KnownBits.h"
 #include "llvm/Support/raw_ostream.h"
 #include <map>
 using namespace llvm;
@@ -1145,6 +1146,19 @@ static ValueLatticeElement getValueFromICmpCondition(Value *Val, ICmpInst *ICI,
   if (matchICmpOperand(Offset, RHS, Val, SwappedPred))
     return getValueFromSimpleICmpCondition(SwappedPred, LHS, Offset);
 
+  // If (Val & Mask) == C then all the masked bits are known and we can compute
+  // a value range based on that.
+  const APInt *Mask, *C;
+  if (EdgePred == ICmpInst::ICMP_EQ &&
+      match(LHS, m_And(m_Specific(Val), m_APInt(Mask))) &&
+      match(RHS, m_APInt(C))) {
+    KnownBits Known;
+    Known.Zero = ~*C & *Mask;
+    Known.One = *C & *Mask;
+    return ValueLatticeElement::getRange(
+        ConstantRange::fromKnownBits(Known, /*IsSigned*/ false));
+  }
+
   return ValueLatticeElement::getOverdefined();
 }
 

llvm/test/Transforms/CorrelatedValuePropagation/icmp.ll

@@ -995,14 +995,10 @@ define void @test_icmp_mask_two_values(i32 %a) {
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[AND]], 10
 ; CHECK-NEXT:    br i1 [[CMP]], label [[IF_TRUE:%.*]], label [[IF_FALSE:%.*]]
 ; CHECK:       if.true:
-; CHECK-NEXT:    [[CMP2:%.*]] = icmp uge i32 [[A]], 10
-; CHECK-NEXT:    call void @check1(i1 [[CMP2]])
-; CHECK-NEXT:    [[CMP3:%.*]] = icmp ule i32 [[A]], 11
-; CHECK-NEXT:    call void @check1(i1 [[CMP3]])
-; CHECK-NEXT:    [[CMP4:%.*]] = icmp ult i32 [[A]], 10
-; CHECK-NEXT:    call void @check1(i1 [[CMP4]])
-; CHECK-NEXT:    [[CMP5:%.*]] = icmp ugt i32 [[A]], 11
-; CHECK-NEXT:    call void @check1(i1 [[CMP5]])
+; CHECK-NEXT:    call void @check1(i1 true)
+; CHECK-NEXT:    call void @check1(i1 true)
+; CHECK-NEXT:    call void @check1(i1 false)
+; CHECK-NEXT:    call void @check1(i1 false)
 ; CHECK-NEXT:    ret void
 ; CHECK:       if.false:
 ; CHECK-NEXT:    ret void
@@ -1032,8 +1028,7 @@ define void @test_icmp_mask_bit_set(i32 %a) {
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[AND]], 32
 ; CHECK-NEXT:    br i1 [[CMP]], label [[IF_TRUE:%.*]], label [[IF_FALSE:%.*]]
 ; CHECK:       if.true:
-; CHECK-NEXT:    [[CMP2:%.*]] = icmp uge i32 [[A]], 32
-; CHECK-NEXT:    call void @check1(i1 [[CMP2]])
+; CHECK-NEXT:    call void @check1(i1 true)
 ; CHECK-NEXT:    [[CMP3:%.*]] = icmp uge i32 [[A]], 33
 ; CHECK-NEXT:    call void @check1(i1 [[CMP3]])
 ; CHECK-NEXT:    ret void
@@ -1061,8 +1056,7 @@ define void @test_icmp_mask_bit_unset(i32 %a) {
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[AND]], 0
 ; CHECK-NEXT:    br i1 [[CMP]], label [[IF_TRUE:%.*]], label [[IF_FALSE:%.*]]
 ; CHECK:       if.true:
-; CHECK-NEXT:    [[CMP2:%.*]] = icmp ule i32 [[A]], -33
-; CHECK-NEXT:    call void @check1(i1 [[CMP2]])
+; CHECK-NEXT:    call void @check1(i1 true)
 ; CHECK-NEXT:    [[CMP3:%.*]] = icmp ule i32 [[A]], -34
 ; CHECK-NEXT:    call void @check1(i1 [[CMP3]])
 ; CHECK-NEXT:    ret void