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