[Attributor] Check range size before constant fold load (llvm#151359)

If the range size doesn't match the type size, it might read wrong data.

Author: shiltian

llvm/lib/Transforms/IPO/Attributor.cpp

@@ -272,6 +272,9 @@ AA::getInitialValueForObj(Attributor &A, const AbstractAttribute &QueryingAA,
   }
 
   if (RangePtr && !RangePtr->offsetOrSizeAreUnknown()) {
+    int64_t StorageSize = DL.getTypeStoreSize(&Ty);
+    if (StorageSize != RangePtr->Size)
+      return nullptr;
     APInt Offset = APInt(64, RangePtr->Offset);
     return ConstantFoldLoadFromConst(Initializer, &Ty, Offset, DL);
   }

llvm/test/Transforms/Attributor/range-and-constant-fold.ll

@@ -0,0 +1,38 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S -passes=attributor %s -o - | FileCheck %s
+
+@g = internal unnamed_addr addrspace(4) constant [3 x i8] c"12\00", align 16
+
+define void @foo(i32 %a, i32 %b, ptr %p.0, ptr %p.1) {
+; CHECK-LABEL: define void @foo(
+; CHECK-SAME: i32 [[A:%.*]], i32 [[B:%.*]], ptr nofree nonnull writeonly captures(none) dereferenceable(1) [[P_0:%.*]], ptr nofree nonnull writeonly align 4 captures(none) dereferenceable(8) [[P_1:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32 [[A]], [[B]]
+; CHECK-NEXT:    br i1 [[CMP]], label %[[L1:.*]], label %[[L2:.*]]
+; CHECK:       [[L1]]:
+; CHECK-NEXT:    br label %[[L3:.*]]
+; CHECK:       [[L2]]:
+; CHECK-NEXT:    br label %[[L3]]
+; CHECK:       [[L3]]:
+; CHECK-NEXT:    [[PHI:%.*]] = phi ptr addrspace(4) [ @g, %[[L1]] ], [ getelementptr inbounds nuw (i8, ptr addrspace(4) @g, i64 1), %[[L2]] ]
+; CHECK-NEXT:    [[LOAD_SMALL:%.*]] = load i8, ptr addrspace(4) [[PHI]], align 4
+; CHECK-NEXT:    store i8 [[LOAD_SMALL]], ptr [[P_0]], align 1
+; CHECK-NEXT:    [[LOAD_LARGE:%.*]] = load i64, ptr addrspace(4) [[PHI]], align 4
+; CHECK-NEXT:    store i64 [[LOAD_LARGE]], ptr [[P_1]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %cmp = icmp ne i32 %a, %b
+  br i1 %cmp, label %l1, label %l2
+l1:
+  br label %l3
+l2:
+  br label %l3
+l3:
+  %phi = phi ptr addrspace(4) [ @g, %l1 ], [ getelementptr inbounds nuw (i8, ptr addrspace(4) @g, i64 1), %l2 ]
+  %load.small = load i8, ptr addrspace(4) %phi
+  store i8 %load.small, ptr %p.0
+  %load.large = load i64, ptr addrspace(4) %phi
+  store i64 %load.large, ptr %p.1
+  ret void
+}