[WinEH] Make FuncletLayout more robust against catchret
[oota-llvm.git] / lib / CodeGen / FuncletLayout.cpp
1 //===-- FuncletLayout.cpp - Contiguously lay out funclets -----------------===//
2 //
3 //                     The LLVM Compiler Infrastructure
4 //
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
7 //
8 //===----------------------------------------------------------------------===//
9 //
10 // This file implements basic block placement transformations which result in
11 // funclets being contiguous.
12 //
13 //===----------------------------------------------------------------------===//
14 #include "llvm/CodeGen/Passes.h"
15 #include "llvm/CodeGen/MachineBasicBlock.h"
16 #include "llvm/CodeGen/MachineFunction.h"
17 #include "llvm/CodeGen/MachineFunctionPass.h"
18 #include "llvm/CodeGen/MachineModuleInfo.h"
19 #include "llvm/Target/TargetInstrInfo.h"
20 #include "llvm/Target/TargetSubtargetInfo.h"
21 using namespace llvm;
22
23 #define DEBUG_TYPE "funclet-layout"
24
25 namespace {
26 class FuncletLayout : public MachineFunctionPass {
27 public:
28   static char ID; // Pass identification, replacement for typeid
29   FuncletLayout() : MachineFunctionPass(ID) {
30     initializeFuncletLayoutPass(*PassRegistry::getPassRegistry());
31   }
32
33   bool runOnMachineFunction(MachineFunction &F) override;
34 };
35 }
36
37 static void
38 collectFuncletMembers(DenseMap<MachineBasicBlock *, int> &FuncletMembership,
39                       int Funclet, MachineBasicBlock *MBB) {
40   // Don't revisit blocks.
41   if (FuncletMembership.count(MBB) > 0) {
42     // FIXME: This is a hack, we need to assert this unconditionally.
43     bool IsProbablyUnreachableBlock =
44         MBB->empty() ||
45         (MBB->succ_empty() && !MBB->getFirstTerminator()->isReturn() &&
46          MBB->size() == 1);
47
48     if (!IsProbablyUnreachableBlock) {
49       if (FuncletMembership[MBB] != Funclet) {
50         assert(false && "MBB is part of two funclets!");
51         report_fatal_error("MBB is part of two funclets!");
52       }
53     }
54     return;
55   }
56
57   // Add this MBB to our funclet.
58   FuncletMembership[MBB] = Funclet;
59
60   bool IsReturn = false;
61   int NumTerminators = 0;
62   for (MachineInstr &MI : MBB->terminators()) {
63     IsReturn |= MI.isReturn();
64     ++NumTerminators;
65   }
66   assert((!IsReturn || NumTerminators == 1) &&
67          "Expected only one terminator when a return is present!");
68
69   // Returns are boundaries where funclet transfer can occur, don't follow
70   // successors.
71   if (IsReturn)
72     return;
73
74   for (MachineBasicBlock *SMBB : MBB->successors())
75     if (!SMBB->isEHPad())
76       collectFuncletMembers(FuncletMembership, Funclet, SMBB);
77 }
78
79 char FuncletLayout::ID = 0;
80 char &llvm::FuncletLayoutID = FuncletLayout::ID;
81 INITIALIZE_PASS(FuncletLayout, "funclet-layout",
82                 "Contiguously Lay Out Funclets", false, false)
83
84 bool FuncletLayout::runOnMachineFunction(MachineFunction &F) {
85   // We don't have anything to do if there aren't any EH pads.
86   if (!F.getMMI().hasEHFunclets())
87     return false;
88
89   const TargetInstrInfo *TII = F.getSubtarget().getInstrInfo();
90   SmallVector<MachineBasicBlock *, 16> FuncletBlocks;
91   SmallVector<std::pair<MachineBasicBlock *, int>, 16> CatchRetSuccessors;
92   for (MachineBasicBlock &MBB : F) {
93     if (MBB.isEHFuncletEntry())
94       FuncletBlocks.push_back(&MBB);
95
96     MachineBasicBlock::iterator MBBI = MBB.getFirstTerminator();
97     if (MBBI->getOpcode() != TII->getCatchReturnOpcode())
98       continue;
99
100     MachineBasicBlock *Successor = MBBI->getOperand(0).getMBB();
101     MachineBasicBlock *SuccessorColor = MBBI->getOperand(1).getMBB();
102     CatchRetSuccessors.push_back({Successor, SuccessorColor->getNumber()});
103   }
104
105   // We don't have anything to do if there aren't any EH pads.
106   if (FuncletBlocks.empty())
107     return false;
108
109   DenseMap<MachineBasicBlock *, int> FuncletMembership;
110   // Identify all the basic blocks reachable from the function entry.
111   collectFuncletMembers(FuncletMembership, F.front().getNumber(), F.begin());
112   // Next, identify all the blocks inside the funclets.
113   for (MachineBasicBlock *MBB : FuncletBlocks)
114     collectFuncletMembers(FuncletMembership, MBB->getNumber(), MBB);
115   // Finally, identify all the targets of a catchret.
116   for (std::pair<MachineBasicBlock *, int> CatchRetPair : CatchRetSuccessors)
117     collectFuncletMembers(FuncletMembership, CatchRetPair.second,
118                           CatchRetPair.first);
119
120   F.sort([&](MachineBasicBlock &x, MachineBasicBlock &y) {
121     return FuncletMembership[&x] < FuncletMembership[&y];
122   });
123
124   // Conservatively assume we changed something.
125   return true;
126 }