diff --git a/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h b/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h --- a/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h +++ b/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h @@ -19,6 +19,7 @@ #include "llvm/IR/Instruction.h" #include "llvm/IR/IntrinsicInst.h" #include "llvm/IR/Module.h" +#include "llvm/IR/ValueHandle.h" namespace llvm { namespace memtag { @@ -74,6 +75,7 @@ struct AllocaInfo { AllocaInst *AI; + TrackingVH OldAI; // Track through RAUW to replace debug uses. SmallVector LifetimeStart; SmallVector LifetimeEnd; SmallVector DbgVariableIntrinsics; diff --git a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp b/llvm/lib/Target/AArch64/AArch64StackTagging.cpp --- a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp +++ b/llvm/lib/Target/AArch64/AArch64StackTagging.cpp @@ -290,14 +290,6 @@ }; class AArch64StackTagging : public FunctionPass { - struct AllocaInfo { - AllocaInst *AI; - TrackingVH OldAI; // Track through RAUW to replace debug uses. - SmallVector LifetimeStart; - SmallVector LifetimeEnd; - SmallVector DbgVariableIntrinsics; - }; - const bool MergeInit; const bool UseStackSafety; @@ -313,7 +305,7 @@ } 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); @@ -322,9 +314,9 @@ Instruction *collectInitializers(Instruction *StartInst, Value *StartPtr, uint64_t Size, InitializerBuilder &IB); - Instruction * - insertBaseTaggedPointer(const MapVector &Allocas, - const DominatorTree *DT); + Instruction *insertBaseTaggedPointer( + const MapVector &Allocas, + const DominatorTree *DT); bool runOnFunction(Function &F) override; StringRef getPassName() const override { return "AArch64 Stack Tagging"; } @@ -466,12 +458,12 @@ } Instruction *AArch64StackTagging::insertBaseTaggedPointer( - const MapVector &InterestingAllocas, + const MapVector &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(); @@ -490,7 +482,7 @@ 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); @@ -536,63 +528,17 @@ if (MergeInit) AA = &getAnalysis().getAAResults(); - MapVector - InterestingAllocas; // need stable iteration order - SmallVector RetVec; - SmallVector UnrecognizedLifetimes; - - bool CallsReturnTwice = false; - for (Instruction &I : instructions(F)) { - if (CallInst *CI = dyn_cast(&I)) { - if (CI->canReturnTwice()) { - CallsReturnTwice = true; - } - } - if (auto *AI = dyn_cast(&I)) { - if (isInterestingAlloca(*AI)) { - InterestingAllocas[AI].AI = AI; - InterestingAllocas[AI].OldAI = AI; - } - continue; - } - - if (auto *DVI = dyn_cast(&I)) { - for (Value *V : DVI->location_ops()) - if (auto *AI = dyn_cast_or_null(V)) - if (isInterestingAlloca(*AI) && - (InterestingAllocas[AI].DbgVariableIntrinsics.empty() || - InterestingAllocas[AI].DbgVariableIntrinsics.back() != DVI)) - InterestingAllocas[AI].DbgVariableIntrinsics.push_back(DVI); - continue; - } - - auto *II = dyn_cast(&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); } @@ -602,7 +548,7 @@ if (auto *P = getAnalysisIfAvailable()) DT = &P->getDomTree(); - if (DT == nullptr && (InterestingAllocas.size() > 1 || + if (DT == nullptr && (SInfo.AllocasToInstrument.size() > 1 || !F->hasFnAttribute(Attribute::OptimizeNone))) { DeleteDT = std::make_unique(*F); DT = DeleteDT.get(); @@ -621,11 +567,11 @@ 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; @@ -642,15 +588,15 @@ 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(Start->getArgOperand(0))->getZExtValue(); @@ -660,7 +606,7 @@ 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(); } @@ -668,7 +614,7 @@ 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. @@ -686,7 +632,7 @@ // 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; diff --git a/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp b/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp --- a/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp +++ b/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp @@ -65,8 +65,10 @@ } } if (AllocaInst *AI = dyn_cast(&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(&Inst);