lib/Transforms/Utils/UnifyFunctionExitNodes.cpp

   1 //===- UnifyFunctionExitNodes.cpp - Make all functions have a single exit -===//
   2 //
   3 // This pass is used to ensure that functions have at most one return
   4 // instruction in them.  Additionally, it keeps track of which node is the new
   5 // exit node of the CFG.  If there are no exit nodes in the CFG, the getExitNode
   6 // method will return a null pointer.
   7 //
   8 //===----------------------------------------------------------------------===//
   9
  10 #include "llvm/Transforms/Utils/UnifyFunctionExitNodes.h"
  11 #include "llvm/Transforms/Scalar.h"
  12 #include "llvm/BasicBlock.h"
  13 #include "llvm/Function.h"
  14 #include "llvm/iTerminators.h"
  15 #include "llvm/iPHINode.h"
  16 #include "llvm/Type.h"
  17
  18 static RegisterOpt<UnifyFunctionExitNodes>
  19 X("mergereturn", "Unify function exit nodes");
  20
  21 Pass *createUnifyFunctionExitNodesPass() {
  22   return new UnifyFunctionExitNodes();
  23 }
  24
  25 void UnifyFunctionExitNodes::getAnalysisUsage(AnalysisUsage &AU) const{
  26   // We preserve the non-critical-edgeness property
  27   AU.addPreservedID(BreakCriticalEdgesID);
  28 }
  29
  30 // UnifyAllExitNodes - Unify all exit nodes of the CFG by creating a new
  31 // BasicBlock, and converting all returns to unconditional branches to this
  32 // new basic block.  The singular exit node is returned.
  33 //
  34 // If there are no return stmts in the Function, a null pointer is returned.
  35 //
  36 bool UnifyFunctionExitNodes::runOnFunction(Function &F) {
  37   // Loop over all of the blocks in a function, tracking all of the blocks that
  38   // return.
  39   //
  40   std::vector<BasicBlock*> ReturningBlocks;
  41   std::vector<BasicBlock*> UnwindingBlocks;
  42   for(Function::iterator I = F.begin(), E = F.end(); I != E; ++I)
  43     if (isa<ReturnInst>(I->getTerminator()))
  44       ReturningBlocks.push_back(I);
  45     else if (isa<UnwindInst>(I->getTerminator()))
  46       UnwindingBlocks.push_back(I);
  47
  48   // Handle unwinding blocks first...
  49   if (UnwindingBlocks.empty()) {
  50     UnwindBlock = 0;
  51   } else if (UnwindingBlocks.size() == 1) {
  52     UnwindBlock = UnwindingBlocks.front();
  53   } else {
  54     UnwindBlock = new BasicBlock("UnifiedUnwindBlock", &F);
  55     UnwindBlock->getInstList().push_back(new UnwindInst());
  56
  57     for (std::vector<BasicBlock*>::iterator I = UnwindingBlocks.begin(),
  58            E = UnwindingBlocks.end(); I != E; ++I) {
  59       BasicBlock *BB = *I;
  60       BB->getInstList().pop_back();  // Remove the return insn
  61       BB->getInstList().push_back(new BranchInst(UnwindBlock));
  62     }
  63   }
  64
  65   // Now handle return blocks...
  66   if (ReturningBlocks.empty()) {
  67     ReturnBlock = 0;
  68     return false;                          // No blocks return
  69   } else if (ReturningBlocks.size() == 1) {
  70     ReturnBlock = ReturningBlocks.front(); // Already has a single return block
  71     return false;
  72   }
  73
  74   // Otherwise, we need to insert a new basic block into the function, add a PHI
  75   // node (if the function returns a value), and convert all of the return
  76   // instructions into unconditional branches.
  77   //
  78   BasicBlock *NewRetBlock = new BasicBlock("UnifiedReturnBlock", &F);
  79
  80   PHINode *PN = 0;
  81   if (F.getReturnType() != Type::VoidTy) {
  82     // If the function doesn't return void... add a PHI node to the block...
  83     PN = new PHINode(F.getReturnType(), "UnifiedRetVal");
  84     NewRetBlock->getInstList().push_back(PN);
  85     NewRetBlock->getInstList().push_back(new ReturnInst(PN));
  86   } else {
  87     // If it returns void, just add a return void instruction to the block
  88     NewRetBlock->getInstList().push_back(new ReturnInst());
  89   }
  90
  91   // Loop over all of the blocks, replacing the return instruction with an
  92   // unconditional branch.
  93   //
  94   for (std::vector<BasicBlock*>::iterator I = ReturningBlocks.begin(),
  95          E = ReturningBlocks.end(); I != E; ++I) {
  96     BasicBlock *BB = *I;
  97
  98     // Add an incoming element to the PHI node for every return instruction that
  99     // is merging into this new block...
 100     if (PN) PN->addIncoming(BB->getTerminator()->getOperand(0), BB);
 101
 102     BB->getInstList().pop_back();  // Remove the return insn
 103     BB->getInstList().push_back(new BranchInst(NewRetBlock));
 104   }
 105   ReturnBlock = NewRetBlock;
 106   return true;
 107 }