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/BasicBlock.h"
  12 #include "llvm/Function.h"
  13 #include "llvm/iTerminators.h"
  14 #include "llvm/iPHINode.h"
  15 #include "llvm/Type.h"
  16 using std::vector;
  17
  18 static RegisterOpt<UnifyFunctionExitNodes>
  19 X("mergereturn", "Unify function exit nodes");
  20
  21 // UnifyAllExitNodes - Unify all exit nodes of the CFG by creating a new
  22 // BasicBlock, and converting all returns to unconditional branches to this
  23 // new basic block.  The singular exit node is returned.
  24 //
  25 // If there are no return stmts in the Function, a null pointer is returned.
  26 //
  27 bool UnifyFunctionExitNodes::runOnFunction(Function &F) {
  28   // Loop over all of the blocks in a function, tracking all of the blocks that
  29   // return.
  30   //
  31   vector<BasicBlock*> ReturningBlocks;
  32   for(Function::iterator I = F.begin(), E = F.end(); I != E; ++I)
  33     if (isa<ReturnInst>(I->getTerminator()))
  34       ReturningBlocks.push_back(I);
  35
  36   if (ReturningBlocks.empty()) {
  37     ExitNode = 0;
  38     return false;                          // No blocks return
  39   } else if (ReturningBlocks.size() == 1) {
  40     ExitNode = ReturningBlocks.front();    // Already has a single return block
  41     return false;
  42   }
  43
  44   // Otherwise, we need to insert a new basic block into the function, add a PHI
  45   // node (if the function returns a value), and convert all of the return
  46   // instructions into unconditional branches.
  47   //
  48   BasicBlock *NewRetBlock = new BasicBlock("UnifiedExitNode", &F);
  49
  50   if (F.getReturnType() != Type::VoidTy) {
  51     // If the function doesn't return void... add a PHI node to the block...
  52     PHINode *PN = new PHINode(F.getReturnType(), "UnifiedRetVal",
  53                               NewRetBlock->end());
  54
  55     // Add an incoming element to the PHI node for every return instruction that
  56     // is merging into this new block...
  57     for (vector<BasicBlock*>::iterator I = ReturningBlocks.begin(),
  58                                        E = ReturningBlocks.end(); I != E; ++I)
  59       PN->addIncoming((*I)->getTerminator()->getOperand(0), *I);
  60
  61     // Add a return instruction to return the result of the PHI node...
  62     new ReturnInst(PN, NewRetBlock->end());
  63   } else {
  64     // If it returns void, just add a return void instruction to the block
  65     new ReturnInst(0, NewRetBlock->end());
  66   }
  67
  68   // Loop over all of the blocks, replacing the return instruction with an
  69   // unconditional branch.
  70   //
  71   for (vector<BasicBlock*>::iterator I = ReturningBlocks.begin(),
  72                                      E = ReturningBlocks.end(); I != E; ++I) {
  73     (*I)->getInstList().pop_back();  // Remove the return insn
  74     new BranchInst(NewRetBlock, (*I)->end());
  75   }
  76   ExitNode = NewRetBlock;
  77   return true;
  78 }