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

[eugenis]

RetVec should be const reference.
Please add a function scope comment explaining what it does.

[fmayer]

add comment & const

[fmayer]

In D106135#2883611, @eugenis wrote:

RetVec should be const reference.
Please add a function scope comment explaining what it does.

Sorry, confused myself by missing that this is a vector of pointers, so it can be const.
Done.

[eugenis]

LGTM

[vitalybuka]

I expected something replaced in AddressSanitizer.cpp as well?

[vitalybuka]

In D106135#2884117, @vitalybuka wrote:

I expected something replaced in AddressSanitizer.cpp as well?

My mistake.
LGTM

[vitalybuka]

llvm/include/llvm/Transforms/Instrumentation/AddressSanitizerCommon.h

61

If it's in common we need to remove all Tagging reference in the function and documentation
UntagCallback -> Callback

[fmayer]

name change.

[fmayer]

rebase

llvm/include/llvm/Transforms/Instrumentation/AddressSanitizerCommon.h

 //===--------- Definition of the AddressSanitizer class ---------*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This file declares common infrastructure for AddressSanitizer and
 // HWAddressSanitizer.
 //
 //===----------------------------------------------------------------------===//
 #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"
 
 namespace llvm {
 
 class InterestingMemoryOperand {
 public:
   Use *PtrUse;
InterestingMemoryOperand(Instruction *I, unsigned OperandNo, bool IsWrite,
     PtrUse = &I->getOperandUse(OperandNo);
   }
 
   Instruction *getInsn() { return cast<Instruction>(PtrUse->getUser()); }
 
   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

llvm/lib/Target/AArch64/AArch64StackTagging.cpp

 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/IR/DebugLoc.h"
 #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>
 #include <utility>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "aarch64-stack-tagging"
for (auto &I : Allocas) {
     if (UnrecognizedLifetimes.empty() && Info.LifetimeStart.size() == 1 &&
         Info.LifetimeEnd.size() == 1) {
       IntrinsicInst *Start = Info.LifetimeStart[0];
       IntrinsicInst *End = Info.LifetimeEnd[0];
       uint64_t Size =
           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))
     } 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) {
         untagAlloca(AI, RI, Size);
       }
       // We may have inserted tag/untag outside of any lifetime interval.