[NFC] [MTE] helper for stack tagging lifetimes.

Reviewed By: eugenis, vitalybuka

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

include/llvm/Transforms/Instrumentation/AddressSanitizerCommon.h

@@ -14,6 +14,9 @@
 #ifndef LLVM_TRANSFORMS_INSTRUMENTATION_ADDRESSSANITIZERCOMMON_H
 #define LLVM_TRANSFORMS_INSTRUMENTATION_ADDRESSSANITIZERCOMMON_H
 
+#include "llvm/Analysis/CFG.h"
+#include "llvm/Analysis/PostDominators.h"
+#include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Module.h"
 
@@ -44,6 +47,47 @@ public:
   Value *getPtr() { return PtrUse->get(); }
 };
 
+// For an alloca valid between lifetime markers Start and End, 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 End was the only possible exit. If it wasn't, the caller
+// should remove End to ensure that work done at the other exits does not
+// happen outside of the lifetime.
+template <typename F>
+bool forAllReachableExits(DominatorTree *DT, PostDominatorTree *PDT,
+                          const Instruction *Start, Instruction *End,
+                          const SmallVectorImpl<Instruction *> &RetVec,
+                          F Callback) {
+  // We need to ensure that if we tag some object, we certainly untag it
+  // before the function exits.
+  if (PDT != nullptr && PDT->dominates(End, Start)) {
+    Callback(End);
+  } else {
+    SmallVector<Instruction *, 8> ReachableRetVec;
+    unsigned NumCoveredExits = 0;
+    for (auto &RI : RetVec) {
+      if (!isPotentiallyReachable(Start, RI, nullptr, DT))
+        continue;
+      ReachableRetVec.push_back(RI);
+      if (DT != nullptr && DT->dominates(End, RI))
+        ++NumCoveredExits;
+    }
+    // If there's a mix of covered and non-covered exits, just put the untag
+    // on exits, so we avoid the redundancy of untagging twice.
+    if (NumCoveredExits == ReachableRetVec.size()) {
+      Callback(End);
+    } else {
+      for (auto &RI : ReachableRetVec)
+        Callback(RI);
+      // We may have inserted untag outside of the lifetime interval.
+      // Signal the caller to remove the lifetime end call for this alloca.
+      return false;
+    }
+  }
+  return true;
+}
+
 } // namespace llvm
 
 #endif

lib/Target/AArch64/AArch64StackTagging.cpp

@@ -42,17 +42,18 @@
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GetElementPtrTypeIterator.h"
+#include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/IntrinsicsAArch64.h"
-#include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/Instrumentation/AddressSanitizerCommon.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include <cassert>
 #include <iterator>
@@ -648,32 +649,10 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
           cast<ConstantInt>(Start->getArgOperand(0))->getZExtValue();
       Size = alignTo(Size, kTagGranuleSize);
       tagAlloca(AI, Start->getNextNode(), Start->getArgOperand(1), Size);
-      // We need to ensure that if we tag some object, we certainly untag it
-      // before the function exits.
-      if (PDT != nullptr && PDT->dominates(End, Start)) {
-        untagAlloca(AI, End, Size);
-      } else {
-        SmallVector<Instruction *, 8> ReachableRetVec;
-        unsigned NumCoveredExits = 0;
-        for (auto &RI : RetVec) {
-          if (!isPotentiallyReachable(Start, RI, nullptr, DT))
-            continue;
-          ReachableRetVec.push_back(RI);
-          if (DT != nullptr && DT->dominates(End, RI))
-            ++NumCoveredExits;
-        }
-        // If there's a mix of covered and non-covered exits, just put the untag
-        // on exits, so we avoid the redundancy of untagging twice.
-        if (NumCoveredExits == ReachableRetVec.size()) {
-          untagAlloca(AI, End, Size);
-        } else {
-          for (auto &RI : ReachableRetVec)
-            untagAlloca(AI, RI, Size);
-          // We may have inserted untag outside of the lifetime interval.
-          // Remove the lifetime end call for this alloca.
-          End->eraseFromParent();
-        }
-      }
+
+      auto TagEnd = [&](Instruction *Node) { untagAlloca(AI, Node, Size); };
+      if (!forAllReachableExits(DT, PDT, Start, End, RetVec, TagEnd))
+        End->eraseFromParent();
     } else {
       uint64_t Size = Info.AI->getAllocationSizeInBits(*DL).getValue() / 8;
       Value *Ptr = IRB.CreatePointerCast(TagPCall, IRB.getInt8PtrTy());