[BoundsChecking] Fix merging of sizes

BoundsChecking uses ObjectSizeOffsetEvaluator to keep track of the
underlying size/offset of pointers in allocations.  However,
ObjectSizeOffsetVisitor (something ObjectSizeOffsetEvaluator
uses to check for constant sizes/offsets)
doesn't quite treat sizes and offsets the same way as
BoundsChecking.  BoundsChecking wants to know the size of the
underlying allocation and the current pointer's offset within
it, but ObjectSizeOffsetVisitor only cares about the size
from the pointer to the end of the underlying allocation.

This only comes up when merging two size/offset pairs. Add a new mode to
ObjectSizeOffsetVisitor which cares about the underlying size/offset
rather than the size from the current pointer to the end of the
allocation.

Fixes a false positive with -fsanitize=bounds.

Reviewed By: vitalybuka, asbirlea

Differential Revision: https://reviews.llvm.org/D131001

GitOrigin-RevId: 203296d642c385da07a62098050b08493ed8b236
diff --git a/include/llvm/Analysis/MemoryBuiltins.h b/include/llvm/Analysis/MemoryBuiltins.h
index 5cc85c4..949fe72 100644
--- a/include/llvm/Analysis/MemoryBuiltins.h
+++ b/include/llvm/Analysis/MemoryBuiltins.h
@@ -134,17 +134,21 @@
 struct ObjectSizeOpts {
   /// Controls how we handle conditional statements with unknown conditions.
   enum class Mode : uint8_t {
-    /// Fail to evaluate an unknown condition.
-    Exact,
+    /// All branches must be known and have the same size, starting from the
+    /// offset, to be merged.
+    ExactSizeFromOffset,
+    /// All branches must be known and have the same underlying size and offset
+    /// to be merged.
+    ExactUnderlyingSizeAndOffset,
     /// Evaluate all branches of an unknown condition. If all evaluations
     /// succeed, pick the minimum size.
     Min,
     /// Same as Min, except we pick the maximum size of all of the branches.
-    Max
+    Max,
   };
 
   /// How we want to evaluate this object's size.
-  Mode EvalMode = Mode::Exact;
+  Mode EvalMode = Mode::ExactSizeFromOffset;
   /// Whether to round the result up to the alignment of allocas, byval
   /// arguments, and global variables.
   bool RoundToAlign = false;
diff --git a/lib/Analysis/MemoryBuiltins.cpp b/lib/Analysis/MemoryBuiltins.cpp
index 413ec6d..31704c2 100644
--- a/lib/Analysis/MemoryBuiltins.cpp
+++ b/lib/Analysis/MemoryBuiltins.cpp
@@ -641,7 +641,7 @@
     EvalOptions.EvalMode =
         MaxVal ? ObjectSizeOpts::Mode::Max : ObjectSizeOpts::Mode::Min;
   else
-    EvalOptions.EvalMode = ObjectSizeOpts::Mode::Exact;
+    EvalOptions.EvalMode = ObjectSizeOpts::Mode::ExactSizeFromOffset;
 
   EvalOptions.NullIsUnknownSize =
       cast<ConstantInt>(ObjectSize->getArgOperand(2))->isOne();
@@ -999,9 +999,11 @@
     return (getSizeWithOverflow(LHS).slt(getSizeWithOverflow(RHS))) ? LHS : RHS;
   case ObjectSizeOpts::Mode::Max:
     return (getSizeWithOverflow(LHS).sgt(getSizeWithOverflow(RHS))) ? LHS : RHS;
-  case ObjectSizeOpts::Mode::Exact:
+  case ObjectSizeOpts::Mode::ExactSizeFromOffset:
     return (getSizeWithOverflow(LHS).eq(getSizeWithOverflow(RHS))) ? LHS
                                                                    : unknown();
+  case ObjectSizeOpts::Mode::ExactUnderlyingSizeAndOffset:
+    return LHS == RHS && LHS.second.eq(RHS.second) ? LHS : unknown();
   }
   llvm_unreachable("missing an eval mode");
 }
diff --git a/lib/Transforms/Instrumentation/BoundsChecking.cpp b/lib/Transforms/Instrumentation/BoundsChecking.cpp
index 1eadafb..81ffa628 100644
--- a/lib/Transforms/Instrumentation/BoundsChecking.cpp
+++ b/lib/Transforms/Instrumentation/BoundsChecking.cpp
@@ -146,6 +146,7 @@
   const DataLayout &DL = F.getParent()->getDataLayout();
   ObjectSizeOpts EvalOpts;
   EvalOpts.RoundToAlign = true;
+  EvalOpts.EvalMode = ObjectSizeOpts::Mode::ExactUnderlyingSizeAndOffset;
   ObjectSizeOffsetEvaluator ObjSizeEval(DL, &TLI, F.getContext(), EvalOpts);
 
   // check HANDLE_MEMORY_INST in include/llvm/Instruction.def for memory
diff --git a/test/Instrumentation/BoundsChecking/simple.ll b/test/Instrumentation/BoundsChecking/simple.ll
index a2af719..5785861 100644
--- a/test/Instrumentation/BoundsChecking/simple.ll
+++ b/test/Instrumentation/BoundsChecking/simple.ll
@@ -330,7 +330,6 @@
 }
 
 ; Check that merging sizes in a phi works.
-; FIXME: bounds-checking thinks that %alloc has an underlying size of 0.
 define i8 @f14(i1 %i) {
 ; CHECK-LABEL: @f14(
 ; CHECK-NEXT:  entry:
@@ -340,16 +339,18 @@
 ; CHECK-NEXT:    [[G:%.*]] = getelementptr i8, ptr [[A]], i32 32
 ; CHECK-NEXT:    br label [[BB2]]
 ; CHECK:       bb2:
-; CHECK-NEXT:    [[ALLOC:%.*]] = phi ptr [ null, [[ENTRY:%.*]] ], [ [[G]], [[BB1]] ]
+; CHECK-NEXT:    [[TMP0:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ 32, [[BB1]] ]
+; CHECK-NEXT:    [[TMP1:%.*]] = phi i64 [ 0, [[ENTRY]] ], [ 32, [[BB1]] ]
+; CHECK-NEXT:    [[ALLOC:%.*]] = phi ptr [ null, [[ENTRY]] ], [ [[G]], [[BB1]] ]
 ; CHECK-NEXT:    [[IND:%.*]] = phi i64 [ 0, [[ENTRY]] ], [ -4, [[BB1]] ]
-; CHECK-NEXT:    [[TMP0:%.*]] = add i64 0, [[IND]]
+; CHECK-NEXT:    [[TMP2:%.*]] = add i64 [[TMP1]], [[IND]]
 ; CHECK-NEXT:    [[P:%.*]] = getelementptr i8, ptr [[ALLOC]], i64 [[IND]]
-; CHECK-NEXT:    [[TMP1:%.*]] = sub i64 0, [[TMP0]]
-; CHECK-NEXT:    [[TMP2:%.*]] = icmp ult i64 0, [[TMP0]]
-; CHECK-NEXT:    [[TMP3:%.*]] = icmp ult i64 [[TMP1]], 1
-; CHECK-NEXT:    [[TMP4:%.*]] = or i1 [[TMP2]], [[TMP3]]
-; CHECK-NEXT:    br i1 [[TMP4]], label [[TRAP:%.*]], label [[TMP5:%.*]]
-; CHECK:       5:
+; CHECK-NEXT:    [[TMP3:%.*]] = sub i64 [[TMP0]], [[TMP2]]
+; CHECK-NEXT:    [[TMP4:%.*]] = icmp ult i64 [[TMP0]], [[TMP2]]
+; CHECK-NEXT:    [[TMP5:%.*]] = icmp ult i64 [[TMP3]], 1
+; CHECK-NEXT:    [[TMP6:%.*]] = or i1 [[TMP4]], [[TMP5]]
+; CHECK-NEXT:    br i1 [[TMP6]], label [[TRAP:%.*]], label [[TMP7:%.*]]
+; CHECK:       7:
 ; CHECK-NEXT:    [[RET:%.*]] = load i8, ptr [[P]], align 1
 ; CHECK-NEXT:    ret i8 [[RET]]
 ; CHECK:       trap: