[NFC] factor isStandardLifetime out of HWASan

this is so we can use it for aarch64 stack tagging.

Reviewed By: eugenis

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

6 files changed, 124 insertions, 76 deletions

diff --git a/llvm/include/llvm/Transforms/Instrumentation/AddressSanitizerCommon.h b/llvm/include/llvm/Transforms/Instrumentation/AddressSanitizerCommon.h
index 0a5456c..7858a1c 100644
--- a/llvm/include/llvm/Transforms/Instrumentation/AddressSanitizerCommon.h
+++ b/llvm/include/llvm/Transforms/Instrumentation/AddressSanitizerCommon.h
@@ -47,51 +47,6 @@ public:
   Value *getPtr() { return PtrUse->get(); }
 };
 
-// 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,
-// the caller should remove Ends to ensure that work done at the other
-// exits does not happen outside of the lifetime.
-template <typename F>
-bool forAllReachableExits(const DominatorTree &DT, const PostDominatorTree &PDT,
-                          const Instruction *Start,
-                          const SmallVectorImpl<IntrinsicInst *> &Ends,
-                          const SmallVectorImpl<Instruction *> &RetVec,
-                          F Callback) {
-  if (Ends.size() == 1 && PDT.dominates(Ends[0], Start)) {
-    Callback(Ends[0]);
-    return true;
-  }
-  SmallVector<Instruction *, 8> ReachableRetVec;
-  unsigned NumCoveredExits = 0;
-  for (auto *RI : RetVec) {
-    if (!isPotentiallyReachable(Start, RI, nullptr, &DT))
-      continue;
-    ReachableRetVec.push_back(RI);
-    // TODO(fmayer): We don't support diamond shapes, where multiple lifetime
-    // ends together dominate the RI, but none of them does by itself.
-    // Check how often this happens and decide whether to support this here.
-    if (std::any_of(Ends.begin(), Ends.end(),
-                    [&](Instruction *End) { return 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()) {
-    for (auto *End : Ends)
-      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;
-}
-
 // Get AddressSanitizer parameters.
 void getAddressSanitizerParams(const Triple &TargetTriple, int LongSize,
                                bool IsKasan, uint64_t *ShadowBase,
diff --git a/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h b/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h
new file mode 100644
index 0000000..8168b51
--- /dev/null
+++ b/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h
@@ -0,0 +1,73 @@
+//===- MemoryTaggingSupport.h - helpers for memory tagging implementations ===//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares common infrastructure for HWAddressSanitizer and
+// Aarch64StackTagging.
+//
+//===----------------------------------------------------------------------===//
+#ifndef LLVM_TRANSFORMS_UTILS_MEMORYTAGGINGSUPPORT_H
+#define LLVM_TRANSFORMS_UTILS_MEMORYTAGGINGSUPPORT_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"
+
+namespace llvm {
+// 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,
+// the caller should remove Ends to ensure that work done at the other
+// exits does not happen outside of the lifetime.
+template <typename F>
+bool forAllReachableExits(const DominatorTree &DT, const PostDominatorTree &PDT,
+                          const Instruction *Start,
+                          const SmallVectorImpl<IntrinsicInst *> &Ends,
+                          const SmallVectorImpl<Instruction *> &RetVec,
+                          F Callback) {
+  if (Ends.size() == 1 && PDT.dominates(Ends[0], Start)) {
+    Callback(Ends[0]);
+    return true;
+  }
+  SmallVector<Instruction *, 8> ReachableRetVec;
+  unsigned NumCoveredExits = 0;
+  for (auto *RI : RetVec) {
+    if (!isPotentiallyReachable(Start, RI, nullptr, &DT))
+      continue;
+    ReachableRetVec.push_back(RI);
+    // TODO(fmayer): We don't support diamond shapes, where multiple lifetime
+    // ends together dominate the RI, but none of them does by itself.
+    // Check how often this happens and decide whether to support this here.
+    if (std::any_of(Ends.begin(), Ends.end(),
+                    [&](Instruction *End) { return 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()) {
+    for (auto *End : Ends)
+      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;
+}
+
+bool isStandardLifetime(const SmallVectorImpl<IntrinsicInst *> &LifetimeStart,
+                        const SmallVectorImpl<IntrinsicInst *> &LifetimeEnd,
+                        const DominatorTree &DT, size_t MaxLifetimes);
+} // namespace llvm
+
+#endif
diff --git a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp b/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
index 566c7a1..13ae502 100644
--- a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
+++ b/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
@@ -52,8 +52,8 @@
 #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 "llvm/Transforms/Utils/MemoryTaggingSupport.h"
 #include <cassert>
 #include <iterator>
 #include <utility>
diff --git a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
index 7b3741d1..b664e6a 100644
--- a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
@@ -53,6 +53,7 @@
 #include "llvm/Transforms/Instrumentation.h"
 #include "llvm/Transforms/Instrumentation/AddressSanitizerCommon.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/MemoryTaggingSupport.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 #include "llvm/Transforms/Utils/PromoteMemToReg.h"
 #include <sstream>
@@ -301,8 +302,6 @@ public:
   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);
-  static bool isStandardLifetime(const AllocaInfo &AllocaInfo,
-                                 const DominatorTree &DT);
   bool instrumentStack(
       bool ShouldDetectUseAfterScope,
       MapVector<AllocaInst *, AllocaInfo> &AllocasToInstrument,
@@ -1330,34 +1329,7 @@ bool HWAddressSanitizer::instrumentLandingPads(
   return true;
 }
 
-static bool
-maybeReachableFromEachOther(const SmallVectorImpl<IntrinsicInst *> &Insts,
-                            const DominatorTree &DT) {
-  // If we have too many lifetime ends, give up, as the algorithm below is N^2.
-  if (Insts.size() > ClMaxLifetimes)
-    return true;
-  for (size_t I = 0; I < Insts.size(); ++I) {
-    for (size_t J = 0; J < Insts.size(); ++J) {
-      if (I == J)
-        continue;
-      if (isPotentiallyReachable(Insts[I], Insts[J], nullptr, &DT))
-        return true;
-    }
-  }
-  return false;
-}
 
-// static
-bool HWAddressSanitizer::isStandardLifetime(const AllocaInfo &AllocaInfo,
-                                            const DominatorTree &DT) {
-  // An alloca that has exactly one start and end in every possible execution.
-  // If it has multiple ends, they have to be unreachable from each other, so
-  // at most one of them is actually used for each execution of the function.
-  return AllocaInfo.LifetimeStart.size() == 1 &&
-         (AllocaInfo.LifetimeEnd.size() == 1 ||
-          (AllocaInfo.LifetimeEnd.size() > 0 &&
-           !maybeReachableFromEachOther(AllocaInfo.LifetimeEnd, DT)));
-}
 
 bool HWAddressSanitizer::instrumentStack(
     bool ShouldDetectUseAfterScope,
@@ -1411,7 +1383,9 @@ bool HWAddressSanitizer::instrumentStack(
       tagAlloca(IRB, AI, UARTag, AlignedSize);
     };
     bool StandardLifetime =
-        UnrecognizedLifetimes.empty() && isStandardLifetime(Info, GetDT());
+        UnrecognizedLifetimes.empty() &&
+        isStandardLifetime(Info.LifetimeStart, Info.LifetimeEnd, GetDT(),
+                           ClMaxLifetimes);
     if (ShouldDetectUseAfterScope && StandardLifetime) {
       IntrinsicInst *Start = Info.LifetimeStart[0];
       IRB.SetInsertPoint(Start->getNextNode());
diff --git a/llvm/lib/Transforms/Utils/CMakeLists.txt b/llvm/lib/Transforms/Utils/CMakeLists.txt
index aa14ba9..ca442d8 100644
--- a/llvm/lib/Transforms/Utils/CMakeLists.txt
+++ b/llvm/lib/Transforms/Utils/CMakeLists.txt
@@ -49,6 +49,7 @@ add_llvm_component_library(LLVMTransformUtils
   LowerSwitch.cpp
   MatrixUtils.cpp
   MemoryOpRemark.cpp
+  MemoryTaggingSupport.cpp
   Mem2Reg.cpp
   MetaRenamer.cpp
   ModuleUtils.cpp
diff --git a/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp b/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp
new file mode 100644
index 0000000..a2ba28a
--- /dev/null
+++ b/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp
@@ -0,0 +1,45 @@
+//== MemoryTaggingSupport.cpp - helpers for memory tagging implementations ===//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares common infrastructure for HWAddressSanitizer and
+// Aarch64StackTagging.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Utils/MemoryTaggingSupport.h"
+
+namespace llvm {
+namespace {
+bool maybeReachableFromEachOther(const SmallVectorImpl<IntrinsicInst *> &Insts,
+                                 const DominatorTree &DT, size_t MaxLifetimes) {
+  // If we have too many lifetime ends, give up, as the algorithm below is N^2.
+  if (Insts.size() > MaxLifetimes)
+    return true;
+  for (size_t I = 0; I < Insts.size(); ++I) {
+    for (size_t J = 0; J < Insts.size(); ++J) {
+      if (I == J)
+        continue;
+      if (isPotentiallyReachable(Insts[I], Insts[J], nullptr, &DT))
+        return true;
+    }
+  }
+  return false;
+}
+} // namespace
+
+bool isStandardLifetime(const SmallVectorImpl<IntrinsicInst *> &LifetimeStart,
+                        const SmallVectorImpl<IntrinsicInst *> &LifetimeEnd,
+                        const DominatorTree &DT, size_t MaxLifetimes) {
+  // An alloca that has exactly one start and end in every possible execution.
+  // If it has multiple ends, they have to be unreachable from each other, so
+  // at most one of them is actually used for each execution of the function.
+  return LifetimeStart.size() == 1 &&
+         (LifetimeEnd.size() == 1 ||
+          (LifetimeEnd.size() > 0 &&
+           !maybeReachableFromEachOther(LifetimeEnd, DT, MaxLifetimes)));
+}
+} // namespace llvm