lib/Transforms/Utils/UnifyFunctionExitNodes.cpp

   1 //===- UnifyFunctionExitNodes.cpp - Make all functions have a single exit -===//
   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 pass is used to ensure that functions have at most one return
  11 // instruction in them.  Additionally, it keeps track of which node is the new
  12 // exit node of the CFG.  If there are no exit nodes in the CFG, the getExitNode
  13 // method will return a null pointer.
  14 //
  15 //===----------------------------------------------------------------------===//
  16
  17 #include "llvm/Transforms/Utils/UnifyFunctionExitNodes.h"
  18 #include "llvm/Transforms/Scalar.h"
  19 #include "llvm/BasicBlock.h"
  20 #include "llvm/Function.h"
  21 #include "llvm/Instructions.h"
  22 #include "llvm/Type.h"
  23 #include "llvm/ADT/SmallVector.h"
  24 #include "llvm/ADT/StringExtras.h"
  25 using namespace llvm;
  26
  27 char UnifyFunctionExitNodes::ID = 0;
  28 static RegisterPass<UnifyFunctionExitNodes>
  29 X("mergereturn", "Unify function exit nodes");
  30
  31 int UnifyFunctionExitNodes::stub;
  32
  33 Pass *llvm::createUnifyFunctionExitNodesPass() {
  34   return new UnifyFunctionExitNodes();
  35 }
  36
  37 void UnifyFunctionExitNodes::getAnalysisUsage(AnalysisUsage &AU) const{
  38   // We preserve the non-critical-edgeness property
  39   AU.addPreservedID(BreakCriticalEdgesID);
  40   // This is a cluster of orthogonal Transforms
  41   AU.addPreservedID(PromoteMemoryToRegisterID);
  42   AU.addPreservedID(LowerSwitchID);
  43 }
  44
  45 // UnifyAllExitNodes - Unify all exit nodes of the CFG by creating a new
  46 // BasicBlock, and converting all returns to unconditional branches to this
  47 // new basic block.  The singular exit node is returned.
  48 //
  49 // If there are no return stmts in the Function, a null pointer is returned.
  50 //
  51 bool UnifyFunctionExitNodes::runOnFunction(Function &F) {
  52   // Loop over all of the blocks in a function, tracking all of the blocks that
  53   // return.
  54   //
  55   std::vector<BasicBlock*> ReturningBlocks;
  56   std::vector<BasicBlock*> UnwindingBlocks;
  57   std::vector<BasicBlock*> UnreachableBlocks;
  58   for(Function::iterator I = F.begin(), E = F.end(); I != E; ++I)
  59     if (isa<ReturnInst>(I->getTerminator()))
  60       ReturningBlocks.push_back(I);
  61     else if (isa<UnwindInst>(I->getTerminator()))
  62       UnwindingBlocks.push_back(I);
  63     else if (isa<UnreachableInst>(I->getTerminator()))
  64       UnreachableBlocks.push_back(I);
  65
  66   // Handle unwinding blocks first.
  67   if (UnwindingBlocks.empty()) {
  68     UnwindBlock = 0;
  69   } else if (UnwindingBlocks.size() == 1) {
  70     UnwindBlock = UnwindingBlocks.front();
  71   } else {
  72     UnwindBlock = new BasicBlock("UnifiedUnwindBlock", &F);
  73     new UnwindInst(UnwindBlock);
  74
  75     for (std::vector<BasicBlock*>::iterator I = UnwindingBlocks.begin(),
  76            E = UnwindingBlocks.end(); I != E; ++I) {
  77       BasicBlock *BB = *I;
  78       BB->getInstList().pop_back();  // Remove the unwind insn
  79       new BranchInst(UnwindBlock, BB);
  80     }
  81   }
  82
  83   // Then unreachable blocks.
  84   if (UnreachableBlocks.empty()) {
  85     UnreachableBlock = 0;
  86   } else if (UnreachableBlocks.size() == 1) {
  87     UnreachableBlock = UnreachableBlocks.front();
  88   } else {
  89     UnreachableBlock = new BasicBlock("UnifiedUnreachableBlock", &F);
  90     new UnreachableInst(UnreachableBlock);
  91
  92     for (std::vector<BasicBlock*>::iterator I = UnreachableBlocks.begin(),
  93            E = UnreachableBlocks.end(); I != E; ++I) {
  94       BasicBlock *BB = *I;
  95       BB->getInstList().pop_back();  // Remove the unreachable inst.
  96       new BranchInst(UnreachableBlock, BB);
  97     }
  98   }
  99
 100   // Now handle return blocks.
 101   if (ReturningBlocks.empty()) {
 102     ReturnBlock = 0;
 103     return false;                          // No blocks return
 104   } else if (ReturningBlocks.size() == 1) {
 105     ReturnBlock = ReturningBlocks.front(); // Already has a single return block
 106     return false;
 107   }
 108
 109   // Otherwise, we need to insert a new basic block into the function, add a PHI
 110   // nodes (if the function returns values), and convert all of the return
 111   // instructions into unconditional branches.
 112   //
 113   BasicBlock *NewRetBlock = new BasicBlock("UnifiedReturnBlock", &F);
 114
 115   SmallVector<Value *, 4> Phis;
 116   unsigned NumRetVals = ReturningBlocks[0]->getTerminator()->getNumOperands();
 117   if (NumRetVals == 0)
 118     new ReturnInst(NULL, NewRetBlock);
 119   else if (const StructType *STy = dyn_cast<StructType>(F.getReturnType())) {
 120     Instruction *InsertPt = NewRetBlock->getFirstNonPHI();
 121     for (unsigned i = 0; i < NumRetVals; ++i) {
 122       PHINode *PN = new PHINode(STy->getElementType(i), "UnifiedRetVal."
 123                                 + utostr(i), InsertPt);
 124       Phis.push_back(PN);
 125     }
 126     new ReturnInst(&Phis[0], NumRetVals);
 127   }
 128   else {
 129     // If the function doesn't return void... add a PHI node to the block...
 130     PHINode *PN = new PHINode(F.getReturnType(), "UnifiedRetVal");
 131     NewRetBlock->getInstList().push_back(PN);
 132     Phis.push_back(PN);
 133     new ReturnInst(PN, NewRetBlock);
 134   }
 135
 136   // Loop over all of the blocks, replacing the return instruction with an
 137   // unconditional branch.
 138   //
 139   for (std::vector<BasicBlock*>::iterator I = ReturningBlocks.begin(),
 140          E = ReturningBlocks.end(); I != E; ++I) {
 141     BasicBlock *BB = *I;
 142
 143     // Add an incoming element to the PHI node for every return instruction that
 144     // is merging into this new block...
 145     if (!Phis.empty()) {
 146       for (unsigned i = 0; i < NumRetVals; ++i)
 147         cast<PHINode>(Phis[i])->addIncoming(BB->getTerminator()->getOperand(i),
 148                                             BB);
 149     }
 150
 151     BB->getInstList().pop_back();  // Remove the return insn
 152     new BranchInst(NewRetBlock, BB);
 153   }
 154   ReturnBlock = NewRetBlock;
 155   return true;
 156 }