[CodeGen] Port StackProtector to new pass manager

llvm/include/llvm/CodeGen/CodeGenPassBuilder.h

@@ -40,6 +40,7 @@
 #include "llvm/CodeGen/SelectOptimize.h"
 #include "llvm/CodeGen/ShadowStackGCLowering.h"
 #include "llvm/CodeGen/SjLjEHPrepare.h"
+#include "llvm/CodeGen/StackProtector.h"
 #include "llvm/CodeGen/UnreachableBlockElim.h"
 #include "llvm/CodeGen/WasmEHPrepare.h"
 #include "llvm/CodeGen/WinEHPrepare.h"
@@ -749,7 +750,7 @@ void CodeGenPassBuilder<Derived>::addISelPrepare(AddIRPass &addPass) const {
   // Add both the safe stack and the stack protection passes: each of them will
   // only protect functions that have corresponding attributes.
   addPass(SafeStackPass(&TM));
-  addPass(StackProtectorPass());
+  addPass(StackProtectorPass(&TM));
 
   if (Opt.PrintISelInput)
     addPass(PrintFunctionPass(dbgs(),

llvm/include/llvm/CodeGen/MachinePassRegistry.def

@@ -34,6 +34,7 @@ MODULE_PASS("shadow-stack-gc-lowering", ShadowStackGCLoweringPass, ())
 #endif
 FUNCTION_ANALYSIS("gc-function", GCFunctionAnalysis, ())
 FUNCTION_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis, (PIC))
+FUNCTION_ANALYSIS("ssp-layout", SSPLayoutAnalysis, ())
 FUNCTION_ANALYSIS("targetir", TargetIRAnalysis,
                   (std::move(TM.getTargetIRAnalysis())))
 #undef FUNCTION_ANALYSIS
@@ -65,6 +66,7 @@ FUNCTION_PASS("safe-stack", SafeStackPass, (TM))
 FUNCTION_PASS("scalarize-masked-mem-intrin", ScalarizeMaskedMemIntrinPass, ())
 FUNCTION_PASS("select-optimize", SelectOptimizePass, (TM))
 FUNCTION_PASS("sjlj-eh-prepare", SjLjEHPreparePass, (TM))
+FUNCTION_PASS("stack-protector", StackProtectorPass, (TM))
 FUNCTION_PASS("tlshoist", TLSVariableHoistPass, ())
 FUNCTION_PASS("unreachableblockelim", UnreachableBlockElimPass, ())
 FUNCTION_PASS("verify", VerifierPass, ())
@@ -134,7 +136,6 @@ MACHINE_FUNCTION_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis,
 #endif
 DUMMY_FUNCTION_PASS("atomic-expand", AtomicExpandPass, ())
 DUMMY_FUNCTION_PASS("codegenprepare", CodeGenPreparePass, ())
-DUMMY_FUNCTION_PASS("stack-protector", StackProtectorPass, ())
 #undef DUMMY_FUNCTION_PASS
 
 #ifndef DUMMY_MACHINE_MODULE_PASS

llvm/include/llvm/CodeGen/StackProtector.h

@@ -19,6 +19,7 @@
 #include "llvm/Analysis/DomTreeUpdater.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/PassManager.h"
 #include "llvm/Pass.h"
 #include "llvm/TargetParser/Triple.h"
 
@@ -30,25 +31,15 @@ class Module;
 class TargetLoweringBase;
 class TargetMachine;
 
-class StackProtector : public FunctionPass {
-private:
+class SSPLayoutInfo {
+  friend class StackProtectorPass;
+  friend class SSPLayoutAnalysis;
+  friend class StackProtector;
   static constexpr unsigned DefaultSSPBufferSize = 8;
 
   /// A mapping of AllocaInsts to their required SSP layout.
-  using SSPLayoutMap = DenseMap<const AllocaInst *,
-                                MachineFrameInfo::SSPLayoutKind>;
-
-  const TargetMachine *TM = nullptr;
-
-  /// TLI - Keep a pointer of a TargetLowering to consult for determining
-  /// target type sizes.
-  const TargetLoweringBase *TLI = nullptr;
-  Triple Trip;
-
-  Function *F = nullptr;
-  Module *M = nullptr;
-
-  std::optional<DomTreeUpdater> DTU;
+  using SSPLayoutMap =
+      DenseMap<const AllocaInst *, MachineFrameInfo::SSPLayoutKind>;
 
   /// Layout - Mapping of allocations to the required SSPLayoutKind.
   /// StackProtector analysis will update this map when determining if an
@@ -59,23 +50,59 @@ class StackProtector : public FunctionPass {
   /// protection when -fstack-protection is used.
   unsigned SSPBufferSize = DefaultSSPBufferSize;
 
+  bool RequireStackProtector = false;
+
   // A prologue is generated.
   bool HasPrologue = false;
 
   // IR checking code is generated.
   bool HasIRCheck = false;
 
-  /// InsertStackProtectors - Insert code into the prologue and epilogue of
-  /// the function.
-  ///
-  ///  - The prologue code loads and stores the stack guard onto the stack.
-  ///  - The epilogue checks the value stored in the prologue against the
-  ///    original value. It calls __stack_chk_fail if they differ.
-  bool InsertStackProtectors();
+public:
+  // Return true if StackProtector is supposed to be handled by SelectionDAG.
+  bool shouldEmitSDCheck(const BasicBlock &BB) const;
+
+  void copyToMachineFrameInfo(MachineFrameInfo &MFI) const;
+};
+
+class SSPLayoutAnalysis : public AnalysisInfoMixin<SSPLayoutAnalysis> {
+  friend class AnalysisInfoMixin<SSPLayoutAnalysis>;
+  using SSPLayoutMap = SSPLayoutInfo::SSPLayoutMap;
+
+  static AnalysisKey Key;
+
+public:
+  using Result = SSPLayoutInfo;
+
+  Result run(Function &F, FunctionAnalysisManager &FAM);
 
-  /// CreateFailBB - Create a basic block to jump to when the stack protector
-  /// check fails.
-  BasicBlock *CreateFailBB();
+  /// Check whether or not \p F needs a stack protector based upon the stack
+  /// protector level.
+  static bool requiresStackProtector(Function *F,
+                                     SSPLayoutMap *Layout = nullptr);
+};
+
+class StackProtectorPass : public PassInfoMixin<StackProtectorPass> {
+  const TargetMachine *TM;
+
+public:
+  explicit StackProtectorPass(const TargetMachine *TM) : TM(TM) {}
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM);
+};
+
+class StackProtector : public FunctionPass {
+private:
+  /// A mapping of AllocaInsts to their required SSP layout.
+  using SSPLayoutMap = SSPLayoutInfo::SSPLayoutMap;
+
+  const TargetMachine *TM = nullptr;
+
+  Function *F = nullptr;
+  Module *M = nullptr;
+
+  std::optional<DomTreeUpdater> DTU;
+
+  SSPLayoutInfo LayoutInfo;
 
 public:
   static char ID; // Pass identification, replacement for typeid.
@@ -85,16 +112,22 @@ class StackProtector : public FunctionPass {
   void getAnalysisUsage(AnalysisUsage &AU) const override;
 
   // Return true if StackProtector is supposed to be handled by SelectionDAG.
-  bool shouldEmitSDCheck(const BasicBlock &BB) const;
+  bool shouldEmitSDCheck(const BasicBlock &BB) const {
+    return LayoutInfo.shouldEmitSDCheck(BB);
+  }
 
   bool runOnFunction(Function &Fn) override;
 
-  void copyToMachineFrameInfo(MachineFrameInfo &MFI) const;
+  void copyToMachineFrameInfo(MachineFrameInfo &MFI) const {
+    LayoutInfo.copyToMachineFrameInfo(MFI);
+  }
 
   /// Check whether or not \p F needs a stack protector based upon the stack
   /// protector level.
-  static bool requiresStackProtector(Function *F, SSPLayoutMap *Layout = nullptr);
-
+  static bool requiresStackProtector(Function *F,
+                                     SSPLayoutMap *Layout = nullptr) {
+    return SSPLayoutAnalysis::requiresStackProtector(F, Layout);
+  }
 };
 
 } // end namespace llvm

llvm/lib/CodeGen/StackProtector.cpp

@@ -64,6 +64,90 @@ static cl::opt<bool> EnableSelectionDAGSP("enable-selectiondag-sp",
 static cl::opt<bool> DisableCheckNoReturn("disable-check-noreturn-call",
                                           cl::init(false), cl::Hidden);
 
+/// InsertStackProtectors - Insert code into the prologue and epilogue of the
+/// function.
+///
+///  - The prologue code loads and stores the stack guard onto the stack.
+///  - The epilogue checks the value stored in the prologue against the original
+///    value. It calls __stack_chk_fail if they differ.
+static bool InsertStackProtectors(const TargetMachine *TM, Function *F,
+                                  DomTreeUpdater *DTU, bool &HasPrologue,
+                                  bool &HasIRCheck);
+
+/// CreateFailBB - Create a basic block to jump to when the stack protector
+/// check fails.
+static BasicBlock *CreateFailBB(Function *F, const Triple &Trip);
+
+bool SSPLayoutInfo::shouldEmitSDCheck(const BasicBlock &BB) const {
+  return HasPrologue && !HasIRCheck && isa<ReturnInst>(BB.getTerminator());
+}
+
+void SSPLayoutInfo::copyToMachineFrameInfo(MachineFrameInfo &MFI) const {
+  if (Layout.empty())
+    return;
+
+  for (int I = 0, E = MFI.getObjectIndexEnd(); I != E; ++I) {
+    if (MFI.isDeadObjectIndex(I))
+      continue;
+
+    const AllocaInst *AI = MFI.getObjectAllocation(I);
+    if (!AI)
+      continue;
+
+    SSPLayoutMap::const_iterator LI = Layout.find(AI);
+    if (LI == Layout.end())
+      continue;
+
+    MFI.setObjectSSPLayout(I, LI->second);
+  }
+}
+
+SSPLayoutInfo SSPLayoutAnalysis::run(Function &F,
+                                     FunctionAnalysisManager &FAM) {
+
+  SSPLayoutInfo Info;
+  Info.RequireStackProtector =
+      SSPLayoutAnalysis::requiresStackProtector(&F, &Info.Layout);
+  Info.SSPBufferSize = F.getFnAttributeAsParsedInteger(
+      "stack-protector-buffer-size", SSPLayoutInfo::DefaultSSPBufferSize);
+  return Info;
+}
+
+AnalysisKey SSPLayoutAnalysis::Key;
+
+PreservedAnalyses StackProtectorPass::run(Function &F,
+                                          FunctionAnalysisManager &FAM) {
+  auto &Info = FAM.getResult<SSPLayoutAnalysis>(F);
+  auto *DT = FAM.getCachedResult<DominatorTreeAnalysis>(F);
+  DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
+
+  if (!Info.RequireStackProtector)
+    return PreservedAnalyses::all();
+
+  // TODO(etienneb): Functions with funclets are not correctly supported now.
+  // Do nothing if this is funclet-based personality.
+  if (F.hasPersonalityFn()) {
+    EHPersonality Personality = classifyEHPersonality(F.getPersonalityFn());
+    if (isFuncletEHPersonality(Personality))
+      return PreservedAnalyses::all();
+  }
+
+  ++NumFunProtected;
+  bool Changed = InsertStackProtectors(TM, &F, DT ? &DTU : nullptr,
+                                       Info.HasPrologue, Info.HasIRCheck);
+#ifdef EXPENSIVE_CHECKS
+  assert((!DT || DT->verify(DominatorTree::VerificationLevel::Full)) &&
+         "Failed to maintain validity of domtree!");
+#endif
+
+  if (!Changed)
+    return PreservedAnalyses::all();
+  PreservedAnalyses PA;
+  PA.preserve<SSPLayoutAnalysis>();
+  PA.preserve<DominatorTreeAnalysis>();
+  return PA;
+}
+
 char StackProtector::ID = 0;
 
 StackProtector::StackProtector() : FunctionPass(ID) {
@@ -90,14 +174,12 @@ bool StackProtector::runOnFunction(Function &Fn) {
   if (auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
     DTU.emplace(DTWP->getDomTree(), DomTreeUpdater::UpdateStrategy::Lazy);
   TM = &getAnalysis<TargetPassConfig>().getTM<TargetMachine>();
-  Trip = TM->getTargetTriple();
-  TLI = TM->getSubtargetImpl(Fn)->getTargetLowering();
-  HasPrologue = false;
-  HasIRCheck = false;
-
-  SSPBufferSize = Fn.getFnAttributeAsParsedInteger(
-      "stack-protector-buffer-size", DefaultSSPBufferSize);
-  if (!requiresStackProtector(F, &Layout))
+  LayoutInfo.HasPrologue = false;
+  LayoutInfo.HasIRCheck = false;
+
+  LayoutInfo.SSPBufferSize = Fn.getFnAttributeAsParsedInteger(
+      "stack-protector-buffer-size", SSPLayoutInfo::DefaultSSPBufferSize);
+  if (!requiresStackProtector(F, &LayoutInfo.Layout))
     return false;
 
   // TODO(etienneb): Functions with funclets are not correctly supported now.
@@ -109,7 +191,9 @@ bool StackProtector::runOnFunction(Function &Fn) {
   }
 
   ++NumFunProtected;
-  bool Changed = InsertStackProtectors();
+  bool Changed =
+      InsertStackProtectors(TM, F, DTU ? &*DTU : nullptr,
+                            LayoutInfo.HasPrologue, LayoutInfo.HasIRCheck);
 #ifdef EXPENSIVE_CHECKS
   assert((!DTU ||
           DTU->getDomTree().verify(DominatorTree::VerificationLevel::Full)) &&
@@ -284,7 +368,8 @@ static const CallInst *findStackProtectorIntrinsic(Function &F) {
 /// functions with aggregates that contain any buffer regardless of type and
 /// size, and functions that contain stack-based variables that have had their
 /// address taken.
-bool StackProtector::requiresStackProtector(Function *F, SSPLayoutMap *Layout) {
+bool SSPLayoutAnalysis::requiresStackProtector(Function *F,
+                                               SSPLayoutMap *Layout) {
   Module *M = F->getParent();
   bool Strong = false;
   bool NeedsProtector = false;
@@ -295,7 +380,7 @@ bool StackProtector::requiresStackProtector(Function *F, SSPLayoutMap *Layout) {
   SmallPtrSet<const PHINode *, 16> VisitedPHIs;
 
   unsigned SSPBufferSize = F->getFnAttributeAsParsedInteger(
-      "stack-protector-buffer-size", DefaultSSPBufferSize);
+      "stack-protector-buffer-size", SSPLayoutInfo::DefaultSSPBufferSize);
 
   if (F->hasFnAttribute(Attribute::SafeStack))
     return false;
@@ -460,13 +545,12 @@ static bool CreatePrologue(Function *F, Module *M, Instruction *CheckLoc,
   return SupportsSelectionDAGSP;
 }
 
-/// InsertStackProtectors - Insert code into the prologue and epilogue of the
-/// function.
-///
-///  - The prologue code loads and stores the stack guard onto the stack.
-///  - The epilogue checks the value stored in the prologue against the original
-///    value. It calls __stack_chk_fail if they differ.
-bool StackProtector::InsertStackProtectors() {
+bool InsertStackProtectors(const TargetMachine *TM, Function *F,
+                           DomTreeUpdater *DTU, bool &HasPrologue,
+                           bool &HasIRCheck) {
+  auto *M = F->getParent();
+  auto *TLI = TM->getSubtargetImpl(*F)->getTargetLowering();
+
   // If the target wants to XOR the frame pointer into the guard value, it's
   // impossible to emit the check in IR, so the target *must* support stack
   // protection in SDAG.
@@ -574,7 +658,7 @@ bool StackProtector::InsertStackProtectors() {
       // merge pass will merge together all of the various BB into one including
       // fail BB generated by the stack protector pseudo instruction.
       if (!FailBB)
-        FailBB = CreateFailBB();
+        FailBB = CreateFailBB(F, TM->getTargetTriple());
 
       IRBuilder<> B(CheckLoc);
       Value *Guard = getStackGuard(TLI, M, B);
@@ -589,8 +673,7 @@ bool StackProtector::InsertStackProtectors() {
                                                  SuccessProb.getNumerator());
 
       SplitBlockAndInsertIfThen(Cmp, CheckLoc,
-                                /*Unreachable=*/false, Weights,
-                                DTU ? &*DTU : nullptr,
+                                /*Unreachable=*/false, Weights, DTU,
                                 /*LI=*/nullptr, /*ThenBlock=*/FailBB);
 
       auto *BI = cast<BranchInst>(Cmp->getParent()->getTerminator());
@@ -608,9 +691,8 @@ bool StackProtector::InsertStackProtectors() {
   return HasPrologue;
 }
 
-/// CreateFailBB - Create a basic block to jump to when the stack protector
-/// check fails.
-BasicBlock *StackProtector::CreateFailBB() {
+BasicBlock *CreateFailBB(Function *F, const Triple &Trip) {
+  auto *M = F->getParent();
   LLVMContext &Context = F->getContext();
   BasicBlock *FailBB = BasicBlock::Create(Context, "CallStackCheckFailBlk", F);
   IRBuilder<> B(FailBB);
@@ -633,27 +715,3 @@ BasicBlock *StackProtector::CreateFailBB() {
   B.CreateUnreachable();
   return FailBB;
 }
-
-bool StackProtector::shouldEmitSDCheck(const BasicBlock &BB) const {
-  return HasPrologue && !HasIRCheck && isa<ReturnInst>(BB.getTerminator());
-}
-
-void StackProtector::copyToMachineFrameInfo(MachineFrameInfo &MFI) const {
-  if (Layout.empty())
-    return;
-
-  for (int I = 0, E = MFI.getObjectIndexEnd(); I != E; ++I) {
-    if (MFI.isDeadObjectIndex(I))
-      continue;
-
-    const AllocaInst *AI = MFI.getObjectAllocation(I);
-    if (!AI)
-      continue;
-
-    SSPLayoutMap::const_iterator LI = Layout.find(AI);
-    if (LI == Layout.end())
-      continue;
-
-    MFI.setObjectSSPLayout(I, LI->second);
-  }
-}