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

Author: khei4

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

@@ -6068,8 +6068,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 NonMonotonic = 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) {
@@ -6089,15 +6090,18 @@ SwitchLookupTable::SwitchLookupTable(
           LinearMappingPossible = false;
           break;
         }
-        Wrapped |=
+        NonMonotonic |=
             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 MayWrap = false;
+      APInt M = LinearMultiplier->getValue();
+      (void)M.smul_ov(APInt(M.getBitWidth(), TableSize - 1), MayWrap);
+      LinearMapValWrapped = NonMonotonic || MayWrap;
       Kind = LinearMapKind;
       ++NumLinearMaps;
       return;

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) {
+; 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
+  i3 0, label %cond.false
+  ]
+
+cond.false:                                       ; preds = %entry, %entry, %entry, %entry
+  %mul = shl nsw i3 %0, 1
+  br label %cond.end
+
+cond.end:                                         ; preds = %entry, %cond.false
+  %cond = phi i3 [ %mul, %cond.false ], [ 2, %entry ]
+  %conv = sext i3 %cond to i8
+  ret i8 %conv
+}