[NFC] [MTE] Move alignAndPadAlloca to MemoryTaggingSupport.

[eugenis]

LGTM

llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h

public:
   void visit(Instruction &Inst);
   StackInfo &get() { return Info; };
 
 private:
   StackInfo Info;
   std::function<bool(const AllocaInst &)> IsInterestingAlloca;
 };
 
+uint64_t getAllocaSizeInBytes(const AllocaInst &AI);
+void alignAndPadAlloca(memtag::AllocaInfo &Info, llvm::Align Align);
+
 } // namespace memtag
 } // namespace llvm
 
 #endif

llvm/lib/Target/AArch64/AArch64StackTagging.cpp

AArch64StackTagging(bool IsOptNone = false)
       : FunctionPass(ID),
         MergeInit(ClMergeInit.getNumOccurrences() ? ClMergeInit : !IsOptNone),
         UseStackSafety(ClUseStackSafety.getNumOccurrences() ? ClUseStackSafety
                                                             : !IsOptNone) {
     initializeAArch64StackTaggingPass(*PassRegistry::getPassRegistry());
   }
 
   bool isInterestingAlloca(const AllocaInst &AI);
-  void alignAndPadAlloca(memtag::AllocaInfo &Info);
 
   void tagAlloca(AllocaInst *AI, Instruction *InsertBefore, Value *Ptr,
                  uint64_t Size);
   void untagAlloca(AllocaInst *AI, Instruction *InsertBefore, uint64_t Size);
 
   Instruction *collectInitializers(Instruction *StartInst, Value *StartPtr,
                                    uint64_t Size, InitializerBuilder &IB);
 
Instruction *AArch64StackTagging::insertBaseTaggedPointer(
   Function *IRG_SP =
       Intrinsic::getDeclaration(F->getParent(), Intrinsic::aarch64_irg_sp);
   Instruction *Base =
       IRB.CreateCall(IRG_SP, {Constant::getNullValue(IRB.getInt64Ty())});
   Base->setName("basetag");
   return Base;
 }
 
-void AArch64StackTagging::alignAndPadAlloca(memtag::AllocaInfo &Info) {
-  const Align NewAlignment =
-      max(MaybeAlign(Info.AI->getAlign()), kTagGranuleSize);
-  Info.AI->setAlignment(NewAlignment);
-
-  uint64_t Size = Info.AI->getAllocationSizeInBits(*DL).getValue() / 8;
-  uint64_t AlignedSize = alignTo(Size, kTagGranuleSize);
-  if (Size == AlignedSize)
-    return;
-
-  // Add padding to the alloca.
-  Type *AllocatedType =
-      Info.AI->isArrayAllocation()
-          ? ArrayType::get(
-                Info.AI->getAllocatedType(),
-                cast<ConstantInt>(Info.AI->getArraySize())->getZExtValue())
-          : Info.AI->getAllocatedType();
-  Type *PaddingType =
-      ArrayType::get(Type::getInt8Ty(F->getContext()), AlignedSize - Size);
-  Type *TypeWithPadding = StructType::get(AllocatedType, PaddingType);
-  auto *NewAI = new AllocaInst(
-      TypeWithPadding, Info.AI->getType()->getAddressSpace(), nullptr, "", Info.AI);
-  NewAI->takeName(Info.AI);
-  NewAI->setAlignment(Info.AI->getAlign());
-  NewAI->setUsedWithInAlloca(Info.AI->isUsedWithInAlloca());
-  NewAI->setSwiftError(Info.AI->isSwiftError());
-  NewAI->copyMetadata(*Info.AI);
-
-  auto *NewPtr = new BitCastInst(NewAI, Info.AI->getType(), "", Info.AI);
-  Info.AI->replaceAllUsesWith(NewPtr);
-  Info.AI->eraseFromParent();
-  Info.AI = NewAI;
-}
-
 // FIXME: check for MTE extension
 bool AArch64StackTagging::runOnFunction(Function &Fn) {
   if (!Fn.hasFnAttribute(Attribute::SanitizeMemTag))
     return false;
 
   if (UseStackSafety)
     SSI = &getAnalysis<StackSafetyGlobalInfoWrapperPass>().getResult();
   F = &Fn;
   DL = &Fn.getParent()->getDataLayout();
   if (MergeInit)
     AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
 
   memtag::StackInfoBuilder SIB(
       [this](const AllocaInst &AI) { return isInterestingAlloca(AI); });
   for (Instruction &I : instructions(F))
     SIB.visit(I);
   memtag::StackInfo &SInfo = SIB.get();
 
   if (SInfo.AllocasToInstrument.empty())
     return false;
 
   for (auto &I : SInfo.AllocasToInstrument) {
     memtag::AllocaInfo &Info = I.second;
     assert(Info.AI && isInterestingAlloca(*Info.AI));
-    alignAndPadAlloca(Info);
+    memtag::alignAndPadAlloca(Info, kTagGranuleSize);
   }
 
   std::unique_ptr<DominatorTree> DeleteDT;
   DominatorTree *DT = nullptr;
   if (auto *P = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
     DT = &P->getDomTree();
 
   if (DT == nullptr && (SInfo.AllocasToInstrument.size() > 1 ||

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

IRB.CreateCall(HwasanMemoryAccessCallbackSized[O.IsWrite],
                    {IRB.CreatePointerCast(Addr, IntptrTy),
                     ConstantInt::get(IntptrTy, O.TypeSize / 8)});
   }
   untagPointerOperand(O.getInsn(), Addr);
 
   return true;
 }
 
-static uint64_t getAllocaSizeInBytes(const AllocaInst &AI) {
-  auto DL = AI.getModule()->getDataLayout();
-  return AI.getAllocationSizeInBits(DL).getValue() / 8;
-}
-
 void HWAddressSanitizer::tagAlloca(IRBuilder<> &IRB, AllocaInst *AI, Value *Tag,
                                    size_t Size) {
   size_t AlignedSize = alignTo(Size, Mapping.getObjectAlignment());
   if (!UseShortGranules)
     Size = AlignedSize;
 
   Value *JustTag = IRB.CreateTrunc(Tag, IRB.getInt8Ty());
   if (InstrumentWithCalls) {
for (auto *DDI : Info.DbgVariableIntrinsics) {
       SmallVector<uint64_t, 8> NewOps = {dwarf::DW_OP_LLVM_tag_offset,
                                          retagMask(N)};
       for (size_t LocNo = 0; LocNo < DDI->getNumVariableLocationOps(); ++LocNo)
         if (DDI->getVariableLocationOp(LocNo) == AI)
           DDI->setExpression(DIExpression::appendOpsToArg(DDI->getExpression(),
                                                           NewOps, LocNo));
     }
 
-    size_t Size = getAllocaSizeInBytes(*AI);
+    size_t Size = memtag::getAllocaSizeInBytes(*AI);
     size_t AlignedSize = alignTo(Size, Mapping.getObjectAlignment());
     auto TagEnd = [&](Instruction *Node) {
       IRB.SetInsertPoint(Node);
       Value *UARTag = getUARTag(IRB, StackTag);
       tagAlloca(IRB, AI, UARTag, AlignedSize);
     };
     bool StandardLifetime =
         SInfo.UnrecognizedLifetimes.empty() &&
bool HWAddressSanitizer::instrumentStack(
   return true;
 }
 
 bool HWAddressSanitizer::isInterestingAlloca(const AllocaInst &AI) {
   return (AI.getAllocatedType()->isSized() &&
           // FIXME: instrument dynamic allocas, too
           AI.isStaticAlloca() &&
           // alloca() may be called with 0 size, ignore it.
-          getAllocaSizeInBytes(AI) > 0 &&
+          memtag::getAllocaSizeInBytes(AI) > 0 &&
           // We are only interested in allocas not promotable to registers.
           // Promotable allocas are common under -O0.
           !isAllocaPromotable(&AI) &&
           // inalloca allocas are not treated as static, and we don't want
           // dynamic alloca instrumentation for them as well.
           !AI.isUsedWithInAlloca() &&
           // swifterror allocas are register promoted by ISel
           !AI.isSwiftError()) &&
          // safe allocas are not interesting
          !(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 = getAllocaSizeInBytes(*AI);
+    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();

llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp

for (Value *V : DVI->location_ops()) {
       }
     }
   }
   Instruction *ExitUntag = getUntagLocationIfFunctionExit(Inst);
   if (ExitUntag)
     Info.RetVec.push_back(ExitUntag);
 }
 
+uint64_t getAllocaSizeInBytes(const AllocaInst &AI) {
+  auto DL = AI.getModule()->getDataLayout();
+  return AI.getAllocationSizeInBits(DL).getValue() / 8;
+}
+
+void alignAndPadAlloca(memtag::AllocaInfo &Info, llvm::Align Alignment) {
+  const Align NewAlignment = max(MaybeAlign(Info.AI->getAlign()), Alignment);
+  Info.AI->setAlignment(NewAlignment);
+  auto &Ctx = Info.AI->getFunction()->getContext();
+
+  uint64_t Size = getAllocaSizeInBytes(*Info.AI);
+  uint64_t AlignedSize = alignTo(Size, Alignment);
+  if (Size == AlignedSize)
+    return;
+
+  // Add padding to the alloca.
+  Type *AllocatedType =
+      Info.AI->isArrayAllocation()
+          ? ArrayType::get(
+                Info.AI->getAllocatedType(),
+                cast<ConstantInt>(Info.AI->getArraySize())->getZExtValue())
+          : Info.AI->getAllocatedType();
+  Type *PaddingType = ArrayType::get(Type::getInt8Ty(Ctx), AlignedSize - Size);
+  Type *TypeWithPadding = StructType::get(AllocatedType, PaddingType);
+  auto *NewAI =
+      new AllocaInst(TypeWithPadding, Info.AI->getType()->getAddressSpace(),
+                     nullptr, "", Info.AI);
+  NewAI->takeName(Info.AI);
+  NewAI->setAlignment(Info.AI->getAlign());
+  NewAI->setUsedWithInAlloca(Info.AI->isUsedWithInAlloca());
+  NewAI->setSwiftError(Info.AI->isSwiftError());
+  NewAI->copyMetadata(*Info.AI);
+
+  auto *NewPtr = new BitCastInst(NewAI, Info.AI->getType(), "", Info.AI);
+  Info.AI->replaceAllUsesWith(NewPtr);
+  Info.AI->eraseFromParent();
+  Info.AI = NewAI;
+}
+
 } // namespace memtag
 } // namespace llvm