hwasan: Pad arrays with non-1 size correctly.

Spotted by eugenis.

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

llvm-svn: 366171

1 files changed, 9 insertions, 5 deletions

diff --git a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
index a961c813587..450ae2f7902 100644
--- a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
@@ -1108,8 +1108,14 @@ bool HWAddressSanitizer::sanitizeFunction(Function &F) {
     uint64_t AlignedSize = alignTo(Size, Mapping.getAllocaAlignment());
     AI->setAlignment(std::max(AI->getAlignment(), 16u));
     if (Size != AlignedSize) {
+      Type *AllocatedType = AI->getAllocatedType();
+      if (AI->isArrayAllocation()) {
+        uint64_t ArraySize =
+            cast<ConstantInt>(AI->getArraySize())->getZExtValue();
+        AllocatedType = ArrayType::get(AllocatedType, ArraySize);
+      }
       Type *TypeWithPadding = StructType::get(
-          AI->getAllocatedType(), ArrayType::get(Int8Ty, AlignedSize - Size));
+          AllocatedType, ArrayType::get(Int8Ty, AlignedSize - Size));
       auto *NewAI = new AllocaInst(
           TypeWithPadding, AI->getType()->getAddressSpace(), nullptr, "", AI);
       NewAI->takeName(AI);
@@ -1117,10 +1123,8 @@ bool HWAddressSanitizer::sanitizeFunction(Function &F) {
       NewAI->setUsedWithInAlloca(AI->isUsedWithInAlloca());
       NewAI->setSwiftError(AI->isSwiftError());
       NewAI->copyMetadata(*AI);
-      Value *Zero = ConstantInt::get(Int32Ty, 0);
-      auto *GEP = GetElementPtrInst::Create(TypeWithPadding, NewAI,
-                                            {Zero, Zero}, "", AI);
-      AI->replaceAllUsesWith(GEP);
+      auto *Bitcast = new BitCastInst(NewAI, AI->getType(), "", AI);
+      AI->replaceAllUsesWith(Bitcast);
       AllocaToPaddedAllocaMap[AI] = NewAI;
     }
   }