From b04fc99c347058111cfb9313546df95e10c26132 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov@redhat.com>
Date: Tue, 6 Dec 2022 16:35:24 +0100
Subject: [PATCH] [ConstantRange] Fix nsw nowrap region for 1 bit integers
 (PR59301)

The special case for V=1 was incorrect for one bit types, where
1 is also -1. Remove it, and use getNonEmpty() to handle the full
range case instead.

Adjust the exhaustive nowrap tests to test both 5 bit and 1 bit
types.

Fixes https://github.com/llvm/llvm-project/issues/59301.
---
 llvm/lib/IR/ConstantRange.cpp                      | 10 ++---
 .../Transforms/CorrelatedValuePropagation/mul.ll   |  9 +++++
 llvm/unittests/IR/ConstantRangeTest.cpp            | 47 +++++++++++-----------
 3 files changed, 36 insertions(+), 30 deletions(-)

diff --git a/llvm/lib/IR/ConstantRange.cpp b/llvm/lib/IR/ConstantRange.cpp
index 553e565..bae1e55 100644
--- a/llvm/lib/IR/ConstantRange.cpp
+++ b/llvm/lib/IR/ConstantRange.cpp
@@ -258,10 +258,9 @@ static ConstantRange makeExactMulNUWRegion(const APInt &V) {
 
 /// Exact mul nsw region for single element RHS.
 static ConstantRange makeExactMulNSWRegion(const APInt &V) {
-  // Handle special case for 0, -1 and 1. See the last for reason why we
-  // specialize -1 and 1.
+  // Handle 0 and -1 separately to avoid division by zero or overflow.
   unsigned BitWidth = V.getBitWidth();
-  if (V == 0 || V.isOne())
+  if (V == 0)
     return ConstantRange::getFull(BitWidth);
 
   APInt MinValue = APInt::getSignedMinValue(BitWidth);
@@ -278,10 +277,7 @@ static ConstantRange makeExactMulNSWRegion(const APInt &V) {
     Lower = APIntOps::RoundingSDiv(MinValue, V, APInt::Rounding::UP);
     Upper = APIntOps::RoundingSDiv(MaxValue, V, APInt::Rounding::DOWN);
   }
-  // ConstantRange ctor take a half inclusive interval [Lower, Upper + 1).
-  // Upper + 1 is guaranteed not to overflow, because |divisor| > 1. 0, -1,
-  // and 1 are already handled as special cases.
-  return ConstantRange(Lower, Upper + 1);
+  return ConstantRange::getNonEmpty(Lower, Upper + 1);
 }
 
 ConstantRange
diff --git a/llvm/test/Transforms/CorrelatedValuePropagation/mul.ll b/llvm/test/Transforms/CorrelatedValuePropagation/mul.ll
index c69d2599..b28107e 100644
--- a/llvm/test/Transforms/CorrelatedValuePropagation/mul.ll
+++ b/llvm/test/Transforms/CorrelatedValuePropagation/mul.ll
@@ -217,3 +217,12 @@ entry:
   %cmp = icmp slt i8 %c, %mul
   ret i1 %cmp
 }
+
+define i1 @one_bit(i1 %a, i1 %b) {
+; CHECK-LABEL: @one_bit(
+; CHECK-NEXT:    [[MUL:%.*]] = mul nuw i1 [[A:%.*]], [[B:%.*]]
+; CHECK-NEXT:    ret i1 [[MUL]]
+;
+  %mul = mul i1 %a, %b
+  ret i1 %mul
+}
diff --git a/llvm/unittests/IR/ConstantRangeTest.cpp b/llvm/unittests/IR/ConstantRangeTest.cpp
index 1427a6c..6d2d85c 100644
--- a/llvm/unittests/IR/ConstantRangeTest.cpp
+++ b/llvm/unittests/IR/ConstantRangeTest.cpp
@@ -1724,32 +1724,33 @@ TEST(ConstantRange, MakeGuaranteedNoWrapRegion) {
 template<typename Fn>
 void TestNoWrapRegionExhaustive(Instruction::BinaryOps BinOp,
                                 unsigned NoWrapKind, Fn OverflowFn) {
-  unsigned Bits = 5;
-  EnumerateConstantRanges(Bits, [&](const ConstantRange &CR) {
-    if (CR.isEmptySet())
-      return;
-    if (Instruction::isShift(BinOp) && CR.getUnsignedMax().uge(Bits))
-      return;
+  for (unsigned Bits : {1, 5}) {
+    EnumerateConstantRanges(Bits, [&](const ConstantRange &CR) {
+      if (CR.isEmptySet())
+        return;
+      if (Instruction::isShift(BinOp) && CR.getUnsignedMax().uge(Bits))
+        return;
+
+      ConstantRange NoWrap =
+          ConstantRange::makeGuaranteedNoWrapRegion(BinOp, CR, NoWrapKind);
+      EnumerateAPInts(Bits, [&](const APInt &N1) {
+        bool NoOverflow = true;
+        bool Overflow = true;
+        ForeachNumInConstantRange(CR, [&](const APInt &N2) {
+          if (OverflowFn(N1, N2))
+            NoOverflow = false;
+          else
+            Overflow = false;
+        });
+        EXPECT_EQ(NoOverflow, NoWrap.contains(N1));
 
-    ConstantRange NoWrap =
-        ConstantRange::makeGuaranteedNoWrapRegion(BinOp, CR, NoWrapKind);
-    EnumerateAPInts(Bits, [&](const APInt &N1) {
-      bool NoOverflow = true;
-      bool Overflow = true;
-      ForeachNumInConstantRange(CR, [&](const APInt &N2) {
-        if (OverflowFn(N1, N2))
-          NoOverflow = false;
-        else
-          Overflow = false;
+        // The no-wrap range is exact for single-element ranges.
+        if (CR.isSingleElement()) {
+          EXPECT_EQ(Overflow, !NoWrap.contains(N1));
+        }
       });
-      EXPECT_EQ(NoOverflow, NoWrap.contains(N1));
-
-      // The no-wrap range is exact for single-element ranges.
-      if (CR.isSingleElement()) {
-        EXPECT_EQ(Overflow, !NoWrap.contains(N1));
-      }
     });
-  });
+  }
 }
 
 // Show that makeGuaranteedNoWrapRegion() is maximal, and for single-element
-- 
2.7.4