[llvm] [SimplifyCFG] handle monotonic wrapped case for D150943 (PR #65882)

[Kohei Asano via llvm-commits]

https://github.com/khei4 created https://github.com/llvm/llvm-project/pull/65882:

As @mikaelholmen sees on https://reviews.llvm.org/D150943#4641707, we need to check the monotonicity of the switch results value and the maximum index's wrap. 

>From 8e0e89841af48d7eac4a82b5464549d16c99a17b Mon Sep 17 00:00:00 2001
From: khei4 <kk.asano.luxy at gmail.com>
Date: Sun, 10 Sep 2023 17:35:12 +0900
Subject: [PATCH] [SimplifyCFG] handle monotonic wrapped case for D150943

---
 llvm/lib/Transforms/Utils/SimplifyCFG.cpp     | 16 ++++++----
 .../SimplifyCFG/X86/switch_to_lookup_table.ll | 29 +++++++++++++++++++
 2 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index c6faafea9485ed6..9dc612bbb8e8111 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -6061,8 +6061,9 @@ SwitchLookupTable::SwitchLookupTable(
     bool LinearMappingPossible = true;
     APInt PrevVal;
     APInt DistToPrev;
-    // When linear map is monotonic, we can attach nsw.
-    bool Wrapped = false;
+    // When linear map is monotonic and signed overflow doesn't happen on
+    // maximum index, we can attach nsw on Add and Mul.
+    bool NMon = false;
     assert(TableSize >= 2 && "Should be a SingleValue table.");
     // Check if there is the same distance between two consecutive values.
     for (uint64_t I = 0; I < TableSize; ++I) {
@@ -6082,15 +6083,20 @@ SwitchLookupTable::SwitchLookupTable(
           LinearMappingPossible = false;
           break;
         }
-        Wrapped |=
-            Dist.isStrictlyPositive() ? Val.sle(PrevVal) : Val.sgt(PrevVal);
+        NMon |= Dist.isStrictlyPositive() ? Val.sle(PrevVal) : Val.sgt(PrevVal);
       }
       PrevVal = Val;
     }
     if (LinearMappingPossible) {
       LinearOffset = cast<ConstantInt>(TableContents[0]);
       LinearMultiplier = ConstantInt::get(M.getContext(), DistToPrev);
-      LinearMapValWrapped = Wrapped;
+      bool MayMWrap = false, MayOWrap = false;
+      APInt M = LinearMultiplier->getValue(), O = LinearOffset->getValue();
+      APInt Res = O.sadd_ov(
+          M.smul_ov(APInt(M.getBitWidth(), TableSize - 1), MayMWrap), MayOWrap);
+      // To avoid -Wunused-result.
+      (void)Res;
+      LinearMapValWrapped = NMon || MayMWrap || MayOWrap;
       Kind = LinearMapKind;
       ++NumLinearMaps;
       return;
diff --git a/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll b/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll
index a00805077aae00e..5a2557f04f7faf6 100644
--- a/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll
+++ b/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll
@@ -2039,3 +2039,32 @@ return:
   %x = phi i8 [ 3, %sw.default ], [ 124, %sw.bb3 ], [ -99, %sw.bb2 ], [ -66, %sw.bb1 ], [ -33, %entry ]
   ret i8 %x
 }
+
+define i8 @linearmap_dec_wrapped_mon(i3 %0) {                              ; Assume 1
+; CHECK-LABEL: @linearmap_dec_wrapped_mon(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SWITCH_TABLEIDX:%.*]] = sub i3 [[TMP0:%.*]], -2
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ult i3 [[SWITCH_TABLEIDX]], -4
+; CHECK-NEXT:    [[SWITCH_IDX_MULT:%.*]] = mul i3 [[SWITCH_TABLEIDX]], 2
+; CHECK-NEXT:    [[SWITCH_OFFSET:%.*]] = add i3 [[SWITCH_IDX_MULT]], -4
+; CHECK-NEXT:    [[COND:%.*]] = select i1 [[TMP1]], i3 [[SWITCH_OFFSET]], i3 2
+; CHECK-NEXT:    [[CONV:%.*]] = sext i3 [[COND]] to i8
+; CHECK-NEXT:    ret i8 [[CONV]]
+;
+entry:
+  switch i3 %0, label %cond.end [
+  i3 -1, label %cond.false
+  i3 -2, label %cond.false
+  i3 1, label %cond.false                          ; -> %cond.false
+  i3 0, label %cond.false
+  ]
+
+cond.false:                                       ; preds = %entry, %entry, %entry, %entry
+  %mul = shl nsw i3 %0, 1                            ; 1 << 1 = 2
+  br label %cond.end
+
+cond.end:                                         ; preds = %entry, %cond.false
+  %cond = phi i3 [ %mul, %cond.false ], [ 2, %entry ]; 2
+  %conv = sext i3 %cond to i8
+  ret i8 %conv
+}