1f22ef01cSRoman Divacky //===- UnifyFunctionExitNodes.cpp - Make all functions have a single exit -===//
2f22ef01cSRoman Divacky //
3f22ef01cSRoman Divacky //                     The LLVM Compiler Infrastructure
4f22ef01cSRoman Divacky //
5f22ef01cSRoman Divacky // This file is distributed under the University of Illinois Open Source
6f22ef01cSRoman Divacky // License. See LICENSE.TXT for details.
7f22ef01cSRoman Divacky //
8f22ef01cSRoman Divacky //===----------------------------------------------------------------------===//
9f22ef01cSRoman Divacky //
10f22ef01cSRoman Divacky // This pass is used to ensure that functions have at most one return
11f22ef01cSRoman Divacky // instruction in them.  Additionally, it keeps track of which node is the new
12f22ef01cSRoman Divacky // exit node of the CFG.  If there are no exit nodes in the CFG, the getExitNode
13f22ef01cSRoman Divacky // method will return a null pointer.
14f22ef01cSRoman Divacky //
15f22ef01cSRoman Divacky //===----------------------------------------------------------------------===//
16f22ef01cSRoman Divacky 
17f22ef01cSRoman Divacky #include "llvm/Transforms/Utils/UnifyFunctionExitNodes.h"
18139f7f9bSDimitry Andric #include "llvm/IR/BasicBlock.h"
19139f7f9bSDimitry Andric #include "llvm/IR/Function.h"
20139f7f9bSDimitry Andric #include "llvm/IR/Instructions.h"
21139f7f9bSDimitry Andric #include "llvm/IR/Type.h"
22*4ba319b5SDimitry Andric #include "llvm/Transforms/Utils.h"
23f22ef01cSRoman Divacky using namespace llvm;
24f22ef01cSRoman Divacky 
25f22ef01cSRoman Divacky char UnifyFunctionExitNodes::ID = 0;
26e580952dSDimitry Andric INITIALIZE_PASS(UnifyFunctionExitNodes, "mergereturn",
272754fe60SDimitry Andric                 "Unify function exit nodes", false, false)
28f22ef01cSRoman Divacky 
createUnifyFunctionExitNodesPass()29f22ef01cSRoman Divacky Pass *llvm::createUnifyFunctionExitNodesPass() {
30f22ef01cSRoman Divacky   return new UnifyFunctionExitNodes();
31f22ef01cSRoman Divacky }
32f22ef01cSRoman Divacky 
getAnalysisUsage(AnalysisUsage & AU) const33f22ef01cSRoman Divacky void UnifyFunctionExitNodes::getAnalysisUsage(AnalysisUsage &AU) const{
34f22ef01cSRoman Divacky   // We preserve the non-critical-edgeness property
35f22ef01cSRoman Divacky   AU.addPreservedID(BreakCriticalEdgesID);
36f22ef01cSRoman Divacky   // This is a cluster of orthogonal Transforms
37f22ef01cSRoman Divacky   AU.addPreservedID(LowerSwitchID);
38f22ef01cSRoman Divacky }
39f22ef01cSRoman Divacky 
40f22ef01cSRoman Divacky // UnifyAllExitNodes - Unify all exit nodes of the CFG by creating a new
41f22ef01cSRoman Divacky // BasicBlock, and converting all returns to unconditional branches to this
42f22ef01cSRoman Divacky // new basic block.  The singular exit node is returned.
43f22ef01cSRoman Divacky //
44f22ef01cSRoman Divacky // If there are no return stmts in the Function, a null pointer is returned.
45f22ef01cSRoman Divacky //
runOnFunction(Function & F)46f22ef01cSRoman Divacky bool UnifyFunctionExitNodes::runOnFunction(Function &F) {
47f22ef01cSRoman Divacky   // Loop over all of the blocks in a function, tracking all of the blocks that
48f22ef01cSRoman Divacky   // return.
49f22ef01cSRoman Divacky   //
50f22ef01cSRoman Divacky   std::vector<BasicBlock*> ReturningBlocks;
51f22ef01cSRoman Divacky   std::vector<BasicBlock*> UnreachableBlocks;
527d523365SDimitry Andric   for (BasicBlock &I : F)
537d523365SDimitry Andric     if (isa<ReturnInst>(I.getTerminator()))
547d523365SDimitry Andric       ReturningBlocks.push_back(&I);
557d523365SDimitry Andric     else if (isa<UnreachableInst>(I.getTerminator()))
567d523365SDimitry Andric       UnreachableBlocks.push_back(&I);
57f22ef01cSRoman Divacky 
58f22ef01cSRoman Divacky   // Then unreachable blocks.
59f22ef01cSRoman Divacky   if (UnreachableBlocks.empty()) {
6091bc56edSDimitry Andric     UnreachableBlock = nullptr;
61f22ef01cSRoman Divacky   } else if (UnreachableBlocks.size() == 1) {
62f22ef01cSRoman Divacky     UnreachableBlock = UnreachableBlocks.front();
63f22ef01cSRoman Divacky   } else {
64f22ef01cSRoman Divacky     UnreachableBlock = BasicBlock::Create(F.getContext(),
65f22ef01cSRoman Divacky                                           "UnifiedUnreachableBlock", &F);
66f22ef01cSRoman Divacky     new UnreachableInst(F.getContext(), UnreachableBlock);
67f22ef01cSRoman Divacky 
683ca95b02SDimitry Andric     for (BasicBlock *BB : UnreachableBlocks) {
69f22ef01cSRoman Divacky       BB->getInstList().pop_back();  // Remove the unreachable inst.
70f22ef01cSRoman Divacky       BranchInst::Create(UnreachableBlock, BB);
71f22ef01cSRoman Divacky     }
72f22ef01cSRoman Divacky   }
73f22ef01cSRoman Divacky 
74f22ef01cSRoman Divacky   // Now handle return blocks.
75f22ef01cSRoman Divacky   if (ReturningBlocks.empty()) {
7691bc56edSDimitry Andric     ReturnBlock = nullptr;
77f22ef01cSRoman Divacky     return false;                          // No blocks return
78f22ef01cSRoman Divacky   } else if (ReturningBlocks.size() == 1) {
79f22ef01cSRoman Divacky     ReturnBlock = ReturningBlocks.front(); // Already has a single return block
80f22ef01cSRoman Divacky     return false;
81f22ef01cSRoman Divacky   }
82f22ef01cSRoman Divacky 
83f22ef01cSRoman Divacky   // Otherwise, we need to insert a new basic block into the function, add a PHI
84f22ef01cSRoman Divacky   // nodes (if the function returns values), and convert all of the return
85f22ef01cSRoman Divacky   // instructions into unconditional branches.
86f22ef01cSRoman Divacky   //
87f22ef01cSRoman Divacky   BasicBlock *NewRetBlock = BasicBlock::Create(F.getContext(),
88f22ef01cSRoman Divacky                                                "UnifiedReturnBlock", &F);
89f22ef01cSRoman Divacky 
9091bc56edSDimitry Andric   PHINode *PN = nullptr;
91f22ef01cSRoman Divacky   if (F.getReturnType()->isVoidTy()) {
9291bc56edSDimitry Andric     ReturnInst::Create(F.getContext(), nullptr, NewRetBlock);
93f22ef01cSRoman Divacky   } else {
94f22ef01cSRoman Divacky     // If the function doesn't return void... add a PHI node to the block...
953b0f4066SDimitry Andric     PN = PHINode::Create(F.getReturnType(), ReturningBlocks.size(),
963b0f4066SDimitry Andric                          "UnifiedRetVal");
97f22ef01cSRoman Divacky     NewRetBlock->getInstList().push_back(PN);
98f22ef01cSRoman Divacky     ReturnInst::Create(F.getContext(), PN, NewRetBlock);
99f22ef01cSRoman Divacky   }
100f22ef01cSRoman Divacky 
101f22ef01cSRoman Divacky   // Loop over all of the blocks, replacing the return instruction with an
102f22ef01cSRoman Divacky   // unconditional branch.
103f22ef01cSRoman Divacky   //
1043ca95b02SDimitry Andric   for (BasicBlock *BB : ReturningBlocks) {
105f22ef01cSRoman Divacky     // Add an incoming element to the PHI node for every return instruction that
106f22ef01cSRoman Divacky     // is merging into this new block...
107f22ef01cSRoman Divacky     if (PN)
108f22ef01cSRoman Divacky       PN->addIncoming(BB->getTerminator()->getOperand(0), BB);
109f22ef01cSRoman Divacky 
110f22ef01cSRoman Divacky     BB->getInstList().pop_back();  // Remove the return insn
111f22ef01cSRoman Divacky     BranchInst::Create(NewRetBlock, BB);
112f22ef01cSRoman Divacky   }
113f22ef01cSRoman Divacky   ReturnBlock = NewRetBlock;
114f22ef01cSRoman Divacky   return true;
115f22ef01cSRoman Divacky }
116