[NFC] [MTE] Use helpers for stack tagging.

[eugenis]

LGTM

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

 
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Analysis/CFG.h"
 #include "llvm/Analysis/PostDominators.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Module.h"
+#include "llvm/IR/ValueHandle.h"
 
 namespace llvm {
 namespace memtag {
 // For an alloca valid between lifetime markers Start and Ends, call the
 // Callback for all possible exits out of the lifetime in the containing
 // function, which can return from the instructions in RetVec.
 //
 // Returns whether Ends covered all possible exits. If they did not,

 bool isStandardLifetime(const SmallVectorImpl<IntrinsicInst *> &LifetimeStart,
                         const SmallVectorImpl<IntrinsicInst *> &LifetimeEnd,
                         const DominatorTree *DT, size_t MaxLifetimes);
 
 Instruction *getUntagLocationIfFunctionExit(Instruction &Inst);
 
 struct AllocaInfo {
   AllocaInst *AI;
+  TrackingVH<Instruction> OldAI; // Track through RAUW to replace debug uses.
   SmallVector<IntrinsicInst *, 2> LifetimeStart;
   SmallVector<IntrinsicInst *, 2> LifetimeEnd;
   SmallVector<DbgVariableIntrinsic *, 2> DbgVariableIntrinsics;
 };
 
 struct StackInfo {
   MapVector<AllocaInst *, AllocaInfo> AllocasToInstrument;
   SmallVector<Instruction *, 4> UnrecognizedLifetimes;

llvm/lib/Target/AArch64/AArch64StackTagging.cpp

if (VectorType *VecTy = dyn_cast<VectorType>(V->getType())) {
       }
     }
     return IRB.CreateBitOrPointerCast(
         V, IRB.getIntNTy(DL->getTypeStoreSize(V->getType()) * 8));
   }
 };
 
 class AArch64StackTagging : public FunctionPass {
-  struct AllocaInfo {
-    AllocaInst *AI;
-    TrackingVH<Instruction> OldAI; // Track through RAUW to replace debug uses.
-    SmallVector<IntrinsicInst *, 2> LifetimeStart;
-    SmallVector<IntrinsicInst *, 2> LifetimeEnd;
-    SmallVector<DbgVariableIntrinsic *, 2> DbgVariableIntrinsics;
-  };
-
   const bool MergeInit;
   const bool UseStackSafety;
 
 public:
   static char ID; // Pass ID, replacement for typeid
 
   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(AllocaInfo &Info);
+  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 *
-  insertBaseTaggedPointer(const MapVector<AllocaInst *, AllocaInfo> &Allocas,
-                          const DominatorTree *DT);
+  Instruction *insertBaseTaggedPointer(
+      const MapVector<AllocaInst *, memtag::AllocaInfo> &Allocas,
   bool runOnFunction(Function &F) override;
 
   StringRef getPassName() const override { return "AArch64 Stack Tagging"; }
 
 private:
   Function *F = nullptr;
   Function *SetTagFunc = nullptr;
   const DataLayout *DL = nullptr;

 void AArch64StackTagging::untagAlloca(AllocaInst *AI, Instruction *InsertBefore,
                                       uint64_t Size) {
   IRBuilder<> IRB(InsertBefore);
   IRB.CreateCall(SetTagFunc, {IRB.CreatePointerCast(AI, IRB.getInt8PtrTy()),
                               ConstantInt::get(IRB.getInt64Ty(), Size)});
 }
 
 Instruction *AArch64StackTagging::insertBaseTaggedPointer(
-    const MapVector<AllocaInst *, AllocaInfo> &InterestingAllocas,
+    const MapVector<AllocaInst *, memtag::AllocaInfo> &AllocasToInstrument,
     const DominatorTree *DT) {
   BasicBlock *PrologueBB = nullptr;
   // Try sinking IRG as deep as possible to avoid hurting shrink wrap.
-  for (auto &I : InterestingAllocas) {
-    const AllocaInfo &Info = I.second;
+  for (auto &I : AllocasToInstrument) {
+    const memtag::AllocaInfo &Info = I.second;
     AllocaInst *AI = Info.AI;
     if (!PrologueBB) {
       PrologueBB = AI->getParent();
       continue;
     }
     PrologueBB = DT->findNearestCommonDominator(PrologueBB, AI->getParent());
   }
   assert(PrologueBB);
 
   IRBuilder<> IRB(&PrologueBB->front());
   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(AllocaInfo &Info) {
+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;
bool AArch64StackTagging::runOnFunction(Function &Fn) {
 
   if (UseStackSafety)
     SSI = &getAnalysis<StackSafetyGlobalInfoWrapperPass>().getResult();
   F = &Fn;
   DL = &Fn.getParent()->getDataLayout();
   if (MergeInit)
     AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
 
-  MapVector<AllocaInst *, AllocaInfo>
-      InterestingAllocas; // need stable iteration order
-  SmallVector<Instruction *, 8> RetVec;
-  SmallVector<Instruction *, 4> UnrecognizedLifetimes;
-
-  bool CallsReturnTwice = false;
-  for (Instruction &I : instructions(F)) {
-    if (CallInst *CI = dyn_cast<CallInst>(&I)) {
-      if (CI->canReturnTwice()) {
-        CallsReturnTwice = true;
-      }
-    }
-    if (auto *AI = dyn_cast<AllocaInst>(&I)) {
-      if (isInterestingAlloca(*AI)) {
-        InterestingAllocas[AI].AI = AI;
-        InterestingAllocas[AI].OldAI = AI;
-      }
-      continue;
-    }
-
-    if (auto *DVI = dyn_cast<DbgVariableIntrinsic>(&I)) {
-      for (Value *V : DVI->location_ops())
-        if (auto *AI = dyn_cast_or_null<AllocaInst>(V))
-          if (isInterestingAlloca(*AI) &&
-              (InterestingAllocas[AI].DbgVariableIntrinsics.empty() ||
-               InterestingAllocas[AI].DbgVariableIntrinsics.back() != DVI))
-            InterestingAllocas[AI].DbgVariableIntrinsics.push_back(DVI);
-      continue;
-    }
-
-    auto *II = dyn_cast<IntrinsicInst>(&I);
-    if (II && (II->getIntrinsicID() == Intrinsic::lifetime_start ||
-                II->getIntrinsicID() == Intrinsic::lifetime_end)) {
-      AllocaInst *AI = findAllocaForValue(II->getArgOperand(1));
-      if (!AI) {
-        UnrecognizedLifetimes.push_back(&I);
-        continue;
-      }
-      if (!isInterestingAlloca(*AI))
-        continue;
-      if (II->getIntrinsicID() == Intrinsic::lifetime_start)
-        InterestingAllocas[AI].LifetimeStart.push_back(II);
-      else
-        InterestingAllocas[AI].LifetimeEnd.push_back(II);
-      continue;
-    }
-
-    Instruction *ExitUntag = memtag::getUntagLocationIfFunctionExit(I);
-    if (ExitUntag)
-      RetVec.push_back(ExitUntag);
-  }
+  memtag::StackInfoBuilder SIB(
+      [this](const AllocaInst &AI) { return isInterestingAlloca(AI); });
+  for (Instruction &I : instructions(F))
+    SIB.visit(I);
+  memtag::StackInfo &SInfo = SIB.get();
 
-  if (InterestingAllocas.empty())
+  if (SInfo.AllocasToInstrument.empty())
     return false;
 
-  for (auto &I : InterestingAllocas) {
-    AllocaInfo &Info = I.second;
+  for (auto &I : SInfo.AllocasToInstrument) {
+    memtag::AllocaInfo &Info = I.second;
     assert(Info.AI && isInterestingAlloca(*Info.AI));
     alignAndPadAlloca(Info);
   }
 
   std::unique_ptr<DominatorTree> DeleteDT;
   DominatorTree *DT = nullptr;
   if (auto *P = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
     DT = &P->getDomTree();
 
-  if (DT == nullptr && (InterestingAllocas.size() > 1 ||
+  if (DT == nullptr && (SInfo.AllocasToInstrument.size() > 1 ||
                         !F->hasFnAttribute(Attribute::OptimizeNone))) {
     DeleteDT = std::make_unique<DominatorTree>(*F);
     DT = DeleteDT.get();
   }
 
   std::unique_ptr<PostDominatorTree> DeletePDT;
   PostDominatorTree *PDT = nullptr;
   if (auto *P = getAnalysisIfAvailable<PostDominatorTreeWrapperPass>())
     PDT = &P->getPostDomTree();
 
   if (PDT == nullptr && !F->hasFnAttribute(Attribute::OptimizeNone)) {
     DeletePDT = std::make_unique<PostDominatorTree>(*F);
     PDT = DeletePDT.get();
   }
 
   SetTagFunc =
       Intrinsic::getDeclaration(F->getParent(), Intrinsic::aarch64_settag);
 
-  Instruction *Base = insertBaseTaggedPointer(InterestingAllocas, DT);
+  Instruction *Base = insertBaseTaggedPointer(SInfo.AllocasToInstrument, DT);
 
   int NextTag = 0;
-  for (auto &I : InterestingAllocas) {
-    const AllocaInfo &Info = I.second;
+  for (auto &I : SInfo.AllocasToInstrument) {
+    const memtag::AllocaInfo &Info = I.second;
     AllocaInst *AI = Info.AI;
     int Tag = NextTag;
     NextTag = (NextTag + 1) % 16;
     // Replace alloca with tagp(alloca).
     IRBuilder<> IRB(Info.AI->getNextNode());
     Function *TagP = Intrinsic::getDeclaration(
         F->getParent(), Intrinsic::aarch64_tagp, {Info.AI->getType()});
     Instruction *TagPCall =
         IRB.CreateCall(TagP, {Constant::getNullValue(Info.AI->getType()), Base,
                               ConstantInt::get(IRB.getInt64Ty(), Tag)});
     if (Info.AI->hasName())
       TagPCall->setName(Info.AI->getName() + ".tag");
     Info.AI->replaceAllUsesWith(TagPCall);
     TagPCall->setOperand(0, Info.AI);
 
     bool StandardLifetime =
-        UnrecognizedLifetimes.empty() &&
+        SInfo.UnrecognizedLifetimes.empty() &&
         memtag::isStandardLifetime(Info.LifetimeStart, Info.LifetimeEnd, DT,
                                    ClMaxLifetimes);
     // Calls to functions that may return twice (e.g. setjmp) confuse the
     // postdominator analysis, and will leave us to keep memory tagged after
     // function return. Work around this by always untagging at every return
     // statement if return_twice functions are called.
-    if (UnrecognizedLifetimes.empty() && StandardLifetime &&
-        !CallsReturnTwice) {
+    if (SInfo.UnrecognizedLifetimes.empty() && StandardLifetime &&
+        !SInfo.CallsReturnTwice) {
       IntrinsicInst *Start = Info.LifetimeStart[0];
       uint64_t Size =
           cast<ConstantInt>(Start->getArgOperand(0))->getZExtValue();
       Size = alignTo(Size, kTagGranuleSize);
       tagAlloca(AI, Start->getNextNode(), Start->getArgOperand(1), Size);
 
       auto TagEnd = [&](Instruction *Node) { untagAlloca(AI, Node, Size); };
       if (!DT || !PDT ||
           !memtag::forAllReachableExits(*DT, *PDT, Start, Info.LifetimeEnd,
-                                        RetVec, TagEnd)) {
+                                        SInfo.RetVec, TagEnd)) {
         for (auto *End : Info.LifetimeEnd)
           End->eraseFromParent();
       }
     } else {
       uint64_t Size = Info.AI->getAllocationSizeInBits(*DL).getValue() / 8;
       Value *Ptr = IRB.CreatePointerCast(TagPCall, IRB.getInt8PtrTy());
       tagAlloca(AI, &*IRB.GetInsertPoint(), Ptr, Size);
-      for (auto &RI : RetVec) {
+      for (auto &RI : SInfo.RetVec) {
         untagAlloca(AI, RI, Size);
       }
       // We may have inserted tag/untag outside of any lifetime interval.
       // Remove all lifetime intrinsics for this alloca.
       for (auto &II : Info.LifetimeStart)
         II->eraseFromParent();
       for (auto &II : Info.LifetimeEnd)
         II->eraseFromParent();
     }
 
     // Fixup debug intrinsics to point to the new alloca.
     for (auto DVI : Info.DbgVariableIntrinsics)
       DVI->replaceVariableLocationOp(Info.OldAI, Info.AI);
   }
 
   // If we have instrumented at least one alloca, all unrecognized lifetime
   // instrinsics have to go.
-  for (auto &I : UnrecognizedLifetimes)
+  for (auto &I : SInfo.UnrecognizedLifetimes)
     I->eraseFromParent();
 
   return true;
 }

llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp

 
 void StackInfoBuilder::visit(Instruction &Inst) {
   if (CallInst *CI = dyn_cast<CallInst>(&Inst)) {
     if (CI->canReturnTwice()) {
       Info.CallsReturnTwice = true;
     }
   }
   if (AllocaInst *AI = dyn_cast<AllocaInst>(&Inst)) {
-    if (IsInterestingAlloca(*AI))
-      Info.AllocasToInstrument.insert({AI, {}});
+    if (IsInterestingAlloca(*AI)) {
+      Info.AllocasToInstrument[AI].AI = AI;
+      Info.AllocasToInstrument[AI].OldAI = AI;
+    }
     return;
   }
   auto *II = dyn_cast<IntrinsicInst>(&Inst);
   if (II && (II->getIntrinsicID() == Intrinsic::lifetime_start ||
              II->getIntrinsicID() == Intrinsic::lifetime_end)) {
     AllocaInst *AI = findAllocaForValue(II->getArgOperand(1));
     if (!AI) {
       Info.UnrecognizedLifetimes.push_back(&Inst);