[PATCH] D119614: [HWASAN] use common alignAndPadAlloca

[Florian Mayer via Phabricator via llvm-commits]

fmayer created this revision.
Herald added a subscriber: hiraditya.
fmayer requested review of this revision.
Herald added a project: LLVM.
Herald added a subscriber: llvm-commits.

Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D119614

Files:
  llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp


Index: llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
===================================================================
--- llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
+++ llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
@@ -261,8 +261,6 @@
 
   void setSSI(const StackSafetyGlobalInfo *S) { SSI = S; }
 
-  DenseMap<AllocaInst *, AllocaInst *> padInterestingAllocas(
-      const MapVector<AllocaInst *, memtag::AllocaInfo> &AllocasToInstrument);
   bool sanitizeFunction(Function &F,
                         llvm::function_ref<const DominatorTree &()> GetDT,
                         llvm::function_ref<const PostDominatorTree &()> GetPDT);
@@ -1380,6 +1378,15 @@
           II->eraseFromParent();
       }
     }
+    memtag::alignAndPadAlloca(Info, Align(Mapping.getObjectAlignment()));
+    for (auto DVI : Info.DbgVariableIntrinsics) {
+      SmallDenseSet<Value *> LocationOps(DVI->location_ops().begin(),
+                                         DVI->location_ops().end());
+      for (Value *V : LocationOps) {
+        if (V == AI)
+          DVI->replaceVariableLocationOp(AI, Info.AI);
+      }
+    }
   }
   for (auto &I : SInfo.UnrecognizedLifetimes)
     I->eraseFromParent();
@@ -1404,39 +1411,6 @@
          !(SSI && SSI->isSafe(AI));
 }
 
-DenseMap<AllocaInst *, AllocaInst *> HWAddressSanitizer::padInterestingAllocas(
-    const MapVector<AllocaInst *, memtag::AllocaInfo> &AllocasToInstrument) {
-  DenseMap<AllocaInst *, AllocaInst *> AllocaToPaddedAllocaMap;
-  for (auto &KV : AllocasToInstrument) {
-    AllocaInst *AI = KV.first;
-    uint64_t Size = memtag::getAllocaSizeInBytes(*AI);
-    uint64_t AlignedSize = alignTo(Size, Mapping.getObjectAlignment());
-    AI->setAlignment(
-        Align(std::max(AI->getAlignment(), Mapping.getObjectAlignment())));
-    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(
-          AllocatedType, ArrayType::get(Int8Ty, AlignedSize - Size));
-      auto *NewAI = new AllocaInst(
-          TypeWithPadding, AI->getType()->getAddressSpace(), nullptr, "", AI);
-      NewAI->takeName(AI);
-      NewAI->setAlignment(AI->getAlign());
-      NewAI->setUsedWithInAlloca(AI->isUsedWithInAlloca());
-      NewAI->setSwiftError(AI->isSwiftError());
-      NewAI->copyMetadata(*AI);
-      auto *Bitcast = new BitCastInst(NewAI, AI->getType(), "", AI);
-      AI->replaceAllUsesWith(Bitcast);
-      AllocaToPaddedAllocaMap[AI] = NewAI;
-    }
-  }
-  return AllocaToPaddedAllocaMap;
-}
-
 bool HWAddressSanitizer::sanitizeFunction(
     Function &F, llvm::function_ref<const DominatorTree &()> GetDT,
     llvm::function_ref<const PostDominatorTree &()> GetPDT) {
@@ -1509,28 +1483,6 @@
     instrumentStack(DetectUseAfterScope && !SInfo.CallsReturnTwice, SIB.get(),
                     StackTag, GetDT, GetPDT);
   }
-  // Pad and align each of the allocas that we instrumented to stop small
-  // uninteresting allocas from hiding in instrumented alloca's padding and so
-  // that we have enough space to store real tags for short granules.
-  DenseMap<AllocaInst *, AllocaInst *> AllocaToPaddedAllocaMap =
-      padInterestingAllocas(SInfo.AllocasToInstrument);
-
-  if (!AllocaToPaddedAllocaMap.empty()) {
-    for (auto &Inst : instructions(F)) {
-      if (auto *DVI = dyn_cast<DbgVariableIntrinsic>(&Inst)) {
-        SmallDenseSet<Value *> LocationOps(DVI->location_ops().begin(),
-                                           DVI->location_ops().end());
-        for (Value *V : LocationOps) {
-          if (auto *AI = dyn_cast_or_null<AllocaInst>(V)) {
-            if (auto *NewAI = AllocaToPaddedAllocaMap.lookup(AI))
-              DVI->replaceVariableLocationOp(V, NewAI);
-          }
-        }
-      }
-    }
-    for (auto &P : AllocaToPaddedAllocaMap)
-      P.first->eraseFromParent();
-  }
 
   // If we split the entry block, move any allocas that were originally in the
   // entry block back into the entry block so that they aren't treated as


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D119614.408141.patch
Type: text/x-patch
Size: 4264 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20220212/efb3cc76/attachment.bin>