[MemoryBuiltins] Consider index type size when aggregating gep offsets

Main goal here is to fix some bugs seen with LowerConstantIntrinsics
pass and the lowering of llvm.objectsize.

In ObjectSizeOffsetVisitor::computeImpl we are using an external
analysis together with stripAndAccumulateConstantOffsets. The idea
is to compute the Min/Max value of individual offsets within a GEP.
The bug solved here is that when doing the Min/Max comparisons the
external analysis wasn't considering the index type size (given by
the data layout), it was simply using the type from the IR. Since a
GEP is defined as sext/truncating indices we need to consider the
index type size in the external analysis.

This solves a regression (false ubsan warnings) seen after commit
02b8ee2 (llvm#117849).

Author: bjope

llvm/lib/Analysis/MemoryBuiltins.cpp

@@ -683,29 +683,30 @@ combinePossibleConstantValues(std::optional<APInt> LHS,
 }
 
 static std::optional<APInt> aggregatePossibleConstantValuesImpl(
-    const Value *V, ObjectSizeOpts::Mode EvalMode, unsigned recursionDepth) {
+    const Value *V, ObjectSizeOpts::Mode EvalMode, unsigned BitWidth,
+    unsigned recursionDepth) {
   constexpr unsigned maxRecursionDepth = 4;
   if (recursionDepth == maxRecursionDepth)
     return std::nullopt;
 
   if (const auto *CI = dyn_cast<ConstantInt>(V)) {
-    return CI->getValue();
+    return CI->getValue().sextOrTrunc(BitWidth);
   } else if (const auto *SI = dyn_cast<SelectInst>(V)) {
     return combinePossibleConstantValues(
         aggregatePossibleConstantValuesImpl(SI->getTrueValue(), EvalMode,
-                                            recursionDepth + 1),
+                                            BitWidth, recursionDepth + 1),
         aggregatePossibleConstantValuesImpl(SI->getFalseValue(), EvalMode,
-                                            recursionDepth + 1),
+                                            BitWidth, recursionDepth + 1),
         EvalMode);
   } else if (const auto *PN = dyn_cast<PHINode>(V)) {
     unsigned Count = PN->getNumIncomingValues();
     if (Count == 0)
       return std::nullopt;
     auto Acc = aggregatePossibleConstantValuesImpl(
-        PN->getIncomingValue(0), EvalMode, recursionDepth + 1);
+        PN->getIncomingValue(0), EvalMode, BitWidth, recursionDepth + 1);
     for (unsigned I = 1; Acc && I < Count; ++I) {
       auto Tmp = aggregatePossibleConstantValuesImpl(
-          PN->getIncomingValue(I), EvalMode, recursionDepth + 1);
+          PN->getIncomingValue(I), EvalMode, BitWidth, recursionDepth + 1);
       Acc = combinePossibleConstantValues(Acc, Tmp, EvalMode);
     }
     return Acc;
@@ -715,9 +716,10 @@ static std::optional<APInt> aggregatePossibleConstantValuesImpl(
 }
 
 static std::optional<APInt>
-aggregatePossibleConstantValues(const Value *V, ObjectSizeOpts::Mode EvalMode) {
+aggregatePossibleConstantValues(const Value *V, ObjectSizeOpts::Mode EvalMode,
+                                unsigned BitWidth) {
   if (auto *CI = dyn_cast<ConstantInt>(V))
-    return CI->getValue();
+    return CI->getValue().sextOrTrunc(BitWidth);
 
   if (EvalMode != ObjectSizeOpts::Mode::Min &&
       EvalMode != ObjectSizeOpts::Mode::Max)
@@ -726,7 +728,7 @@ aggregatePossibleConstantValues(const Value *V, ObjectSizeOpts::Mode EvalMode) {
   // Not using computeConstantRange here because we cannot guarantee it's not
   // doing optimization based on UB which we want to avoid when expanding
   // __builtin_object_size.
-  return aggregatePossibleConstantValuesImpl(V, EvalMode, 0u);
+  return aggregatePossibleConstantValuesImpl(V, EvalMode, BitWidth, 0u);
 }
 
 /// Align \p Size according to \p Alignment. If \p Size is greater than
@@ -788,9 +790,14 @@ OffsetSpan ObjectSizeOffsetVisitor::computeImpl(Value *V) {
         Options.EvalMode == ObjectSizeOpts::Mode::Min
             ? ObjectSizeOpts::Mode::Max
             : ObjectSizeOpts::Mode::Min;
-    auto OffsetRangeAnalysis = [EvalMode](Value &VOffset, APInt &Offset) {
+    // For a GEPOperator the indices are first converted to offsets in the
+    // pointer’s index type, so we need to provide the index type to make sure
+    // the min/max operations are performed in correct type.
+    unsigned IdxTyBits = DL.getIndexTypeSizeInBits(V->getType());
+    auto OffsetRangeAnalysis = [EvalMode, IdxTyBits](Value &VOffset,
+                                                     APInt &Offset) {
       if (auto PossibleOffset =
-              aggregatePossibleConstantValues(&VOffset, EvalMode)) {
+              aggregatePossibleConstantValues(&VOffset, EvalMode, IdxTyBits)) {
         Offset = *PossibleOffset;
         return true;
       }
@@ -900,8 +907,9 @@ OffsetSpan ObjectSizeOffsetVisitor::visitAllocaInst(AllocaInst &I) {
     return OffsetSpan(Zero, align(Size, I.getAlign()));
 
   Value *ArraySize = I.getArraySize();
-  if (auto PossibleSize =
-          aggregatePossibleConstantValues(ArraySize, Options.EvalMode)) {
+  if (auto PossibleSize = aggregatePossibleConstantValues(
+          ArraySize, Options.EvalMode,
+          ArraySize->getType()->getScalarSizeInBits())) {
     APInt NumElems = *PossibleSize;
     if (!CheckedZextOrTrunc(NumElems))
       return ObjectSizeOffsetVisitor::unknown();
@@ -932,8 +940,8 @@ OffsetSpan ObjectSizeOffsetVisitor::visitCallBase(CallBase &CB) {
     if (!V->getType()->isIntegerTy())
       return V;
 
-    if (auto PossibleBound =
-            aggregatePossibleConstantValues(V, Options.EvalMode))
+    if (auto PossibleBound = aggregatePossibleConstantValues(
+            V, Options.EvalMode, V->getType()->getScalarSizeInBits()))
       return ConstantInt::get(V->getType(), *PossibleBound);
 
     return V;

llvm/test/Transforms/LowerConstantIntrinsics/builtin-object-size-idxsize.ll

@@ -89,19 +89,12 @@ entry:
 
 ; Indices are truncated to the pointer size in a gep. So "i32 -65526" should
 ; be truncated to "i16 10".
-; FIXME: The TST result here is incorrect!
 define i32 @possible_out_of_bounds_gep_i32_trunc(i1 %c0, i1 %c1) {
-; CHECK-REF-LABEL: define i32 @possible_out_of_bounds_gep_i32_trunc(
-; CHECK-REF-SAME: i1 [[C0:%.*]], i1 [[C1:%.*]]) {
-; CHECK-REF-NEXT:  [[ENTRY:.*:]]
-; CHECK-REF-NEXT:    [[RES:%.*]] = select i1 [[C1]], i32 3, i32 0
-; CHECK-REF-NEXT:    ret i32 [[RES]]
-;
-; CHECK-TST-LABEL: define i32 @possible_out_of_bounds_gep_i32_trunc(
-; CHECK-TST-SAME: i1 [[C0:%.*]], i1 [[C1:%.*]]) {
-; CHECK-TST-NEXT:  [[ENTRY:.*:]]
-; CHECK-TST-NEXT:    [[RES:%.*]] = select i1 [[C1]], i32 0, i32 3
-; CHECK-TST-NEXT:    ret i32 [[RES]]
+; CHECK-LABEL: define i32 @possible_out_of_bounds_gep_i32_trunc(
+; CHECK-SAME: i1 [[C0:%.*]], i1 [[C1:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[RES:%.*]] = select i1 [[C1]], i32 3, i32 0
+; CHECK-NEXT:    ret i32 [[RES]]
 ;
 entry:
   %obj = alloca [5 x i8]
@@ -207,7 +200,8 @@ define i32 @out_of_bounds_gep_i32_trunc_select(i1 %c0, i1 %c1) {
 ; CHECK-TST-LABEL: define i32 @out_of_bounds_gep_i32_trunc_select(
 ; CHECK-TST-SAME: i1 [[C0:%.*]], i1 [[C1:%.*]]) {
 ; CHECK-TST-NEXT:  [[ENTRY:.*:]]
-; CHECK-TST-NEXT:    ret i32 0
+; CHECK-TST-NEXT:    [[RES:%.*]] = select i1 [[C1]], i32 -1, i32 0
+; CHECK-TST-NEXT:    ret i32 [[RES]]
 ;
 entry:
   %obj = alloca [5 x i8]