[nfc] [hwasan] factor out logic to collect info about stack

this is the first step in unifying some of the logic between hwasan and
mte stack tagging. this only moves around code, changes to converge
different implementations of the same logic follow later.

Reviewed By: eugenis

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

Author: fmayer

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

@@ -21,6 +21,7 @@
 #include "llvm/IR/Module.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.
@@ -71,6 +72,34 @@ bool isStandardLifetime(const SmallVectorImpl<IntrinsicInst *> &LifetimeStart,
 
 Instruction *getUntagLocationIfFunctionExit(Instruction &Inst);
 
+struct AllocaInfo {
+  AllocaInst *AI;
+  SmallVector<IntrinsicInst *, 2> LifetimeStart;
+  SmallVector<IntrinsicInst *, 2> LifetimeEnd;
+};
+
+struct StackInfo {
+  MapVector<AllocaInst *, AllocaInfo> AllocasToInstrument;
+  SmallVector<Instruction *, 4> UnrecognizedLifetimes;
+  DenseMap<AllocaInst *, std::vector<DbgVariableIntrinsic *>> AllocaDbgMap;
+  SmallVector<Instruction *, 8> RetVec;
+  bool CallsReturnTwice = false;
+};
+
+class StackInfoBuilder {
+public:
+  StackInfoBuilder(std::function<bool(const AllocaInst &)> IsInterestingAlloca)
+      : IsInterestingAlloca(IsInterestingAlloca) {}
+
+  void visit(Instruction &Inst);
+  StackInfo &get() { return Info; };
+
+private:
+  StackInfo Info;
+  std::function<bool(const AllocaInst &)> IsInterestingAlloca;
+};
+
+} // namespace memtag
 } // namespace llvm
 
 #endif

llvm/lib/Target/AArch64/AArch64StackTagging.cpp

@@ -583,7 +583,7 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
       continue;
     }
 
-    Instruction *ExitUntag = getUntagLocationIfFunctionExit(I);
+    Instruction *ExitUntag = memtag::getUntagLocationIfFunctionExit(I);
     if (ExitUntag)
       RetVec.push_back(ExitUntag);
   }
@@ -643,8 +643,8 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
 
     bool StandardLifetime =
         UnrecognizedLifetimes.empty() &&
-        isStandardLifetime(Info.LifetimeStart, Info.LifetimeEnd, DT,
-                           ClMaxLifetimes);
+        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
@@ -659,8 +659,8 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
 
       auto TagEnd = [&](Instruction *Node) { untagAlloca(AI, Node, Size); };
       if (!DT || !PDT ||
-          !forAllReachableExits(*DT, *PDT, Start, Info.LifetimeEnd, RetVec,
-                                TagEnd)) {
+          !memtag::forAllReachableExits(*DT, *PDT, Start, Info.LifetimeEnd,
+                                        RetVec, TagEnd)) {
         for (auto *End : Info.LifetimeEnd)
           End->eraseFromParent();
       }

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

@@ -247,13 +247,6 @@ bool shouldDetectUseAfterScope(const Triple &TargetTriple) {
 /// An instrumentation pass implementing detection of addressability bugs
 /// using tagged pointers.
 class HWAddressSanitizer {
-private:
-  struct AllocaInfo {
-    AllocaInst *AI;
-    SmallVector<IntrinsicInst *, 2> LifetimeStart;
-    SmallVector<IntrinsicInst *, 2> LifetimeEnd;
-  };
-
 public:
   HWAddressSanitizer(Module &M, bool CompileKernel, bool Recover,
                      const StackSafetyGlobalInfo *SSI)
@@ -269,7 +262,7 @@ class HWAddressSanitizer {
   void setSSI(const StackSafetyGlobalInfo *S) { SSI = S; }
 
   DenseMap<AllocaInst *, AllocaInst *> padInterestingAllocas(
-      const MapVector<AllocaInst *, AllocaInfo> &AllocasToInstrument);
+      const MapVector<AllocaInst *, memtag::AllocaInfo> &AllocasToInstrument);
   bool sanitizeFunction(Function &F,
                         llvm::function_ref<const DominatorTree &()> GetDT,
                         llvm::function_ref<const PostDominatorTree &()> GetPDT);
@@ -304,14 +297,10 @@ class HWAddressSanitizer {
   void tagAlloca(IRBuilder<> &IRB, AllocaInst *AI, Value *Tag, size_t Size);
   Value *tagPointer(IRBuilder<> &IRB, Type *Ty, Value *PtrLong, Value *Tag);
   Value *untagPointer(IRBuilder<> &IRB, Value *PtrLong);
-  bool instrumentStack(
-      bool ShouldDetectUseAfterScope,
-      MapVector<AllocaInst *, AllocaInfo> &AllocasToInstrument,
-      SmallVector<Instruction *, 4> &UnrecognizedLifetimes,
-      DenseMap<AllocaInst *, std::vector<DbgVariableIntrinsic *>> &AllocaDbgMap,
-      SmallVectorImpl<Instruction *> &RetVec, Value *StackTag,
-      llvm::function_ref<const DominatorTree &()> GetDT,
-      llvm::function_ref<const PostDominatorTree &()> GetPDT);
+  bool instrumentStack(bool ShouldDetectUseAfterScope, memtag::StackInfo &Info,
+                       Value *StackTag,
+                       llvm::function_ref<const DominatorTree &()> GetDT,
+                       llvm::function_ref<const PostDominatorTree &()> GetPDT);
   Value *readRegister(IRBuilder<> &IRB, StringRef Name);
   bool instrumentLandingPads(SmallVectorImpl<Instruction *> &RetVec);
   Value *getNextTagWithCall(IRBuilder<> &IRB);
@@ -1325,11 +1314,7 @@ bool HWAddressSanitizer::instrumentLandingPads(
 }
 
 bool HWAddressSanitizer::instrumentStack(
-    bool ShouldDetectUseAfterScope,
-    MapVector<AllocaInst *, AllocaInfo> &AllocasToInstrument,
-    SmallVector<Instruction *, 4> &UnrecognizedLifetimes,
-    DenseMap<AllocaInst *, std::vector<DbgVariableIntrinsic *>> &AllocaDbgMap,
-    SmallVectorImpl<Instruction *> &RetVec, Value *StackTag,
+    bool ShouldDetectUseAfterScope, memtag::StackInfo &SInfo, Value *StackTag,
     llvm::function_ref<const DominatorTree &()> GetDT,
     llvm::function_ref<const PostDominatorTree &()> GetPDT) {
   // Ideally, we want to calculate tagged stack base pointer, and rewrite all
@@ -1339,10 +1324,10 @@ bool HWAddressSanitizer::instrumentStack(
   // This generates one extra instruction per alloca use.
   unsigned int I = 0;
 
-  for (auto &KV : AllocasToInstrument) {
+  for (auto &KV : SInfo.AllocasToInstrument) {
     auto N = I++;
     auto *AI = KV.first;
-    AllocaInfo &Info = KV.second;
+    memtag::AllocaInfo &Info = KV.second;
     IRBuilder<> IRB(AI->getNextNode());
 
     // Replace uses of the alloca with tagged address.
@@ -1356,7 +1341,7 @@ bool HWAddressSanitizer::instrumentStack(
     AI->replaceUsesWithIf(Replacement,
                           [AILong](Use &U) { return U.getUser() != AILong; });
 
-    for (auto *DDI : AllocaDbgMap.lookup(AI)) {
+    for (auto *DDI : SInfo.AllocaDbgMap.lookup(AI)) {
       // Prepend "tag_offset, N" to the dwarf expression.
       // Tag offset logically applies to the alloca pointer, and it makes sense
       // to put it at the beginning of the expression.
@@ -1376,21 +1361,22 @@ bool HWAddressSanitizer::instrumentStack(
       tagAlloca(IRB, AI, UARTag, AlignedSize);
     };
     bool StandardLifetime =
-        UnrecognizedLifetimes.empty() &&
-        isStandardLifetime(Info.LifetimeStart, Info.LifetimeEnd, &GetDT(),
-                           ClMaxLifetimes);
+        SInfo.UnrecognizedLifetimes.empty() &&
+        memtag::isStandardLifetime(Info.LifetimeStart, Info.LifetimeEnd,
+                                   &GetDT(), ClMaxLifetimes);
     if (ShouldDetectUseAfterScope && StandardLifetime) {
       IntrinsicInst *Start = Info.LifetimeStart[0];
       IRB.SetInsertPoint(Start->getNextNode());
       tagAlloca(IRB, AI, Tag, Size);
-      if (!forAllReachableExits(GetDT(), GetPDT(), Start, Info.LifetimeEnd,
-                                RetVec, TagEnd)) {
+      if (!memtag::forAllReachableExits(GetDT(), GetPDT(), Start,
+                                        Info.LifetimeEnd, SInfo.RetVec,
+                                        TagEnd)) {
         for (auto *End : Info.LifetimeEnd)
           End->eraseFromParent();
       }
     } else {
       tagAlloca(IRB, AI, Tag, Size);
-      for (auto *RI : RetVec)
+      for (auto *RI : SInfo.RetVec)
         TagEnd(RI);
       if (!StandardLifetime) {
         for (auto &II : Info.LifetimeStart)
@@ -1400,7 +1386,7 @@ bool HWAddressSanitizer::instrumentStack(
       }
     }
   }
-  for (auto &I : UnrecognizedLifetimes)
+  for (auto &I : SInfo.UnrecognizedLifetimes)
     I->eraseFromParent();
   return true;
 }
@@ -1424,7 +1410,7 @@ bool HWAddressSanitizer::isInterestingAlloca(const AllocaInst &AI) {
 }
 
 DenseMap<AllocaInst *, AllocaInst *> HWAddressSanitizer::padInterestingAllocas(
-    const MapVector<AllocaInst *, AllocaInfo> &AllocasToInstrument) {
+    const MapVector<AllocaInst *, memtag::AllocaInfo> &AllocasToInstrument) {
   DenseMap<AllocaInst *, AllocaInst *> AllocaToPaddedAllocaMap;
   for (auto &KV : AllocasToInstrument) {
     AllocaInst *AI = KV.first;
@@ -1469,52 +1455,13 @@ bool HWAddressSanitizer::sanitizeFunction(
 
   SmallVector<InterestingMemoryOperand, 16> OperandsToInstrument;
   SmallVector<MemIntrinsic *, 16> IntrinToInstrument;
-  MapVector<AllocaInst *, AllocaInfo> AllocasToInstrument;
-  SmallVector<Instruction *, 8> RetVec;
   SmallVector<Instruction *, 8> LandingPadVec;
-  SmallVector<Instruction *, 4> UnrecognizedLifetimes;
-  DenseMap<AllocaInst *, std::vector<DbgVariableIntrinsic *>> AllocaDbgMap;
-  bool CallsReturnTwice = false;
+
+  memtag::StackInfoBuilder SIB(
+      [this](const AllocaInst &AI) { return isInterestingAlloca(AI); });
   for (auto &Inst : instructions(F)) {
-    if (CallInst *CI = dyn_cast<CallInst>(&Inst)) {
-      if (CI->canReturnTwice()) {
-        CallsReturnTwice = true;
-      }
-    }
     if (InstrumentStack) {
-      if (AllocaInst *AI = dyn_cast<AllocaInst>(&Inst)) {
-        if (isInterestingAlloca(*AI))
-          AllocasToInstrument.insert({AI, {}});
-        continue;
-      }
-      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) {
-          UnrecognizedLifetimes.push_back(&Inst);
-          continue;
-        }
-        if (!isInterestingAlloca(*AI))
-          continue;
-        if (II->getIntrinsicID() == Intrinsic::lifetime_start)
-          AllocasToInstrument[AI].LifetimeStart.push_back(II);
-        else
-          AllocasToInstrument[AI].LifetimeEnd.push_back(II);
-        continue;
-      }
-    }
-
-    Instruction *ExitUntag = getUntagLocationIfFunctionExit(Inst);
-    if (ExitUntag)
-      RetVec.push_back(ExitUntag);
-
-    if (auto *DVI = dyn_cast<DbgVariableIntrinsic>(&Inst)) {
-      for (Value *V : DVI->location_ops()) {
-        if (auto *Alloca = dyn_cast_or_null<AllocaInst>(V))
-          if (!AllocaDbgMap.count(Alloca) || AllocaDbgMap[Alloca].back() != DVI)
-            AllocaDbgMap[Alloca].push_back(DVI);
-      }
+      SIB.visit(Inst);
     }
 
     if (InstrumentLandingPads && isa<LandingPadInst>(Inst))
@@ -1527,22 +1474,24 @@ bool HWAddressSanitizer::sanitizeFunction(
         IntrinToInstrument.push_back(MI);
   }
 
+  memtag::StackInfo &SInfo = SIB.get();
+
   initializeCallbacks(*F.getParent());
 
   bool Changed = false;
 
   if (!LandingPadVec.empty())
     Changed |= instrumentLandingPads(LandingPadVec);
 
-  if (AllocasToInstrument.empty() && F.hasPersonalityFn() &&
+  if (SInfo.AllocasToInstrument.empty() && F.hasPersonalityFn() &&
       F.getPersonalityFn()->getName() == kHwasanPersonalityThunkName) {
     // __hwasan_personality_thunk is a no-op for functions without an
     // instrumented stack, so we can drop it.
     F.setPersonalityFn(nullptr);
     Changed = true;
   }
 
-  if (AllocasToInstrument.empty() && OperandsToInstrument.empty() &&
+  if (SInfo.AllocasToInstrument.empty() && OperandsToInstrument.empty() &&
       IntrinToInstrument.empty())
     return Changed;
 
@@ -1552,24 +1501,24 @@ bool HWAddressSanitizer::sanitizeFunction(
   IRBuilder<> EntryIRB(InsertPt);
   emitPrologue(EntryIRB,
                /*WithFrameRecord*/ ClRecordStackHistory &&
-                   Mapping.WithFrameRecord && !AllocasToInstrument.empty());
+                   Mapping.WithFrameRecord &&
+                   !SInfo.AllocasToInstrument.empty());
 
-  if (!AllocasToInstrument.empty()) {
+  if (!SInfo.AllocasToInstrument.empty()) {
     Value *StackTag =
         ClGenerateTagsWithCalls ? nullptr : getStackBaseTag(EntryIRB);
     // 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.
-    instrumentStack(DetectUseAfterScope && !CallsReturnTwice,
-                    AllocasToInstrument, UnrecognizedLifetimes, AllocaDbgMap,
-                    RetVec, StackTag, GetDT, GetPDT);
+    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(AllocasToInstrument);
+      padInterestingAllocas(SInfo.AllocasToInstrument);
 
   if (!AllocaToPaddedAllocaMap.empty()) {
     for (auto &Inst : instructions(F)) {

llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp

@@ -12,7 +12,10 @@
 
 #include "llvm/Transforms/Utils/MemoryTaggingSupport.h"
 
+#include "llvm/Analysis/ValueTracking.h"
+
 namespace llvm {
+namespace memtag {
 namespace {
 bool maybeReachableFromEachOther(const SmallVectorImpl<IntrinsicInst *> &Insts,
                                  const DominatorTree *DT, size_t MaxLifetimes) {
@@ -54,4 +57,46 @@ Instruction *getUntagLocationIfFunctionExit(Instruction &Inst) {
   }
   return nullptr;
 }
+
+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, {}});
+    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);
+      return;
+    }
+    if (!IsInterestingAlloca(*AI))
+      return;
+    if (II->getIntrinsicID() == Intrinsic::lifetime_start)
+      Info.AllocasToInstrument[AI].LifetimeStart.push_back(II);
+    else
+      Info.AllocasToInstrument[AI].LifetimeEnd.push_back(II);
+    return;
+  }
+  if (auto *DVI = dyn_cast<DbgVariableIntrinsic>(&Inst)) {
+    for (Value *V : DVI->location_ops()) {
+      if (auto *Alloca = dyn_cast_or_null<AllocaInst>(V))
+        if (!Info.AllocaDbgMap.count(Alloca) ||
+            Info.AllocaDbgMap[Alloca].back() != DVI)
+          Info.AllocaDbgMap[Alloca].push_back(DVI);
+    }
+  }
+  Instruction *ExitUntag = getUntagLocationIfFunctionExit(Inst);
+  if (ExitUntag)
+    Info.RetVec.push_back(ExitUntag);
+}
+
+} // namespace memtag
 } // namespace llvm