lib/Transforms/Utils/LowerSwitch.cpp

   1 //===- LowerSwitch.cpp - Eliminate Switch instructions --------------------===//
   2 //
   3 // The LowerSwitch transformation rewrites switch statements with a sequence of
   4 // branches, which allows targets to get away with not implementing the switch
   5 // statement until it is convenient.
   6 //
   7 //===----------------------------------------------------------------------===//
   8
   9 #include "llvm/Transforms/Scalar.h"
  10 #include "llvm/Constants.h"
  11 #include "llvm/Function.h"
  12 #include "llvm/iTerminators.h"
  13 #include "llvm/iOperators.h"
  14 #include "llvm/iPHINode.h"
  15 #include "llvm/Pass.h"
  16 #include "Support/Debug.h"
  17 #include "Support/Statistic.h"
  18
  19 namespace {
  20   Statistic<> NumLowered("lowerswitch", "Number of SwitchInst's replaced");
  21
  22   /// LowerSwitch Pass - Replace all SwitchInst instructions with chained branch
  23   /// instructions.  Note that this cannot be a BasicBlock pass because it
  24   /// modifies the CFG!
  25   class LowerSwitch : public FunctionPass {
  26   public:
  27     bool runOnFunction(Function &F);
  28     typedef std::pair<Constant*, BasicBlock*> Case;
  29     typedef std::vector<Case>::iterator       CaseItr;
  30   private:
  31     void processSwitchInst(SwitchInst *SI);
  32
  33     BasicBlock* switchConvert(CaseItr Begin, CaseItr End, Value* Val,
  34                               BasicBlock* OrigBlock, BasicBlock* Default);
  35     BasicBlock* newLeafBlock(Case& Leaf, Value* Val,
  36                              BasicBlock* OrigBlock, BasicBlock* Default);
  37   };
  38
  39   /// The comparison function for sorting the switch case values in the vector.
  40   struct CaseCmp {
  41     bool operator () (const LowerSwitch::Case& C1,
  42                       const LowerSwitch::Case& C2) {
  43       if (const ConstantUInt* U1 = dyn_cast<const ConstantUInt>(C1.first))
  44         return U1->getValue() < cast<const ConstantUInt>(C2.first)->getValue();
  45
  46       const ConstantSInt* S1 = dyn_cast<const ConstantSInt>(C1.first);
  47       return S1->getValue() < cast<const ConstantSInt>(C2.first)->getValue();
  48     }
  49   };
  50
  51   RegisterOpt<LowerSwitch>
  52   X("lowerswitch", "Lower SwitchInst's to branches");
  53 }
  54
  55 // createLowerSwitchPass - Interface to this file...
  56 FunctionPass *createLowerSwitchPass() {
  57   return new LowerSwitch();
  58 }
  59
  60 bool LowerSwitch::runOnFunction(Function &F) {
  61   bool Changed = false;
  62
  63   for (Function::iterator I = F.begin(), E = F.end(); I != E; ) {
  64     BasicBlock *Cur = I++; // Advance over block so we don't traverse new blocks
  65
  66     if (SwitchInst *SI = dyn_cast<SwitchInst>(Cur->getTerminator())) {
  67       Changed = true;
  68       processSwitchInst(SI);
  69     }
  70   }
  71
  72   return Changed;
  73 }
  74
  75 // operator<< - Used for debugging purposes.
  76 //
  77 std::ostream& operator << (std::ostream& O, std::vector<LowerSwitch::Case>& C)
  78 {
  79   O << "[";
  80
  81   for (std::vector<LowerSwitch::Case>::iterator B = C.begin(), E = C.end();
  82       B != E; ) {
  83     O << *B->first;
  84     if (++B != E) O << ", ";
  85   }
  86
  87   return O << "]";
  88 }
  89
  90 // switchConvert - Convert the switch statement into a binary lookup of
  91 // the case values. The function recursively builds this tree.
  92 //
  93 BasicBlock* LowerSwitch::switchConvert(CaseItr Begin, CaseItr End,
  94                                        Value* Val, BasicBlock* OrigBlock,
  95                                        BasicBlock* Default)
  96 {
  97   unsigned Size = End - Begin;
  98
  99   if (Size == 1)
 100     return newLeafBlock(*Begin, Val, OrigBlock, Default);
 101
 102   unsigned Mid = Size / 2;
 103   std::vector<Case> LHS(Begin, Begin + Mid);
 104   DEBUG(std::cerr << "LHS: " << LHS << "\n");
 105   std::vector<Case> RHS(Begin + Mid, End);
 106   DEBUG(std::cerr << "RHS: " << RHS << "\n");
 107
 108   Case& Pivot = *(Begin + Mid);
 109   DEBUG(std::cerr << "Pivot ==> "
 110                   << cast<ConstantUInt>(Pivot.first)->getValue() << "\n");
 111
 112   BasicBlock* LBranch = switchConvert(LHS.begin(), LHS.end(), Val,
 113                                       OrigBlock, Default);
 114   BasicBlock* RBranch = switchConvert(RHS.begin(), RHS.end(), Val,
 115                                       OrigBlock, Default);
 116
 117   // Create a new node that checks if the value is < pivot. Go to the
 118   // left branch if it is and right branch if not.
 119   Function* F = OrigBlock->getParent();
 120   BasicBlock* NewNode = new BasicBlock("NodeBlock");
 121   F->getBasicBlockList().insert(OrigBlock->getNext(), NewNode);
 122
 123   SetCondInst* Comp = new SetCondInst(Instruction::SetLT, Val, Pivot.first,
 124                                       "Pivot");
 125   NewNode->getInstList().push_back(Comp);
 126   BranchInst* Br = new BranchInst(LBranch, RBranch, Comp);
 127   NewNode->getInstList().push_back(Br);
 128   return NewNode;
 129 }
 130
 131 // newLeafBlock - Create a new leaf block for the binary lookup tree. It
 132 // checks if the switch's value == the case's value. If not, then it
 133 // jumps to the default branch. At this point in the tree, the value
 134 // can't be another valid case value, so the jump to the "default" branch
 135 // is warranted.
 136 //
 137 BasicBlock* LowerSwitch::newLeafBlock(Case& Leaf, Value* Val,
 138                                       BasicBlock* OrigBlock,
 139                                       BasicBlock* Default)
 140 {
 141   Function* F = OrigBlock->getParent();
 142   BasicBlock* NewLeaf = new BasicBlock("LeafBlock");
 143   F->getBasicBlockList().insert(OrigBlock->getNext(), NewLeaf);
 144
 145   // Make the seteq instruction...
 146   SetCondInst* Comp = new SetCondInst(Instruction::SetEQ, Val,
 147                                       Leaf.first, "SwitchLeaf");
 148   NewLeaf->getInstList().push_back(Comp);
 149
 150   // Make the conditional branch...
 151   BasicBlock* Succ = Leaf.second;
 152   Instruction* Br = new BranchInst(Succ, Default, Comp);
 153   NewLeaf->getInstList().push_back(Br);
 154
 155   // If there were any PHI nodes in this successor, rewrite one entry
 156   // from OrigBlock to come from NewLeaf.
 157   for (BasicBlock::iterator I = Succ->begin();
 158        PHINode* PN = dyn_cast<PHINode>(I); ++I) {
 159     int BlockIdx = PN->getBasicBlockIndex(OrigBlock);
 160     assert(BlockIdx != -1 && "Switch didn't go to this successor??");
 161     PN->setIncomingBlock((unsigned)BlockIdx, NewLeaf);
 162   }
 163
 164   return NewLeaf;
 165 }
 166
 167 // processSwitchInst - Replace the specified switch instruction with a sequence
 168 // of chained if-then insts in a balanced binary search.
 169 //
 170 void LowerSwitch::processSwitchInst(SwitchInst *SI) {
 171   BasicBlock *CurBlock = SI->getParent();
 172   BasicBlock *OrigBlock = CurBlock;
 173   Function *F = CurBlock->getParent();
 174   Value *Val = SI->getOperand(0);  // The value we are switching on...
 175   BasicBlock* Default = SI->getDefaultDest();
 176
 177   // Unlink the switch instruction from it's block.
 178   CurBlock->getInstList().remove(SI);
 179
 180   // If there is only the default destination, don't bother with the code below.
 181   if (SI->getNumOperands() == 2) {
 182     CurBlock->getInstList().push_back(new BranchInst(SI->getDefaultDest()));
 183     delete SI;
 184     return;
 185   }
 186
 187   // Create a new, empty default block so that the new hierarchy of
 188   // if-then statements go to this and the PHI nodes are happy.
 189   BasicBlock* NewDefault = new BasicBlock("NewDefault");
 190   F->getBasicBlockList().insert(Default, NewDefault);
 191
 192   NewDefault->getInstList().push_back(new BranchInst(Default));
 193
 194   // If there is an entry in any PHI nodes for the default edge, make sure
 195   // to update them as well.
 196   for (BasicBlock::iterator I = Default->begin();
 197        PHINode *PN = dyn_cast<PHINode>(I); ++I) {
 198     int BlockIdx = PN->getBasicBlockIndex(OrigBlock);
 199     assert(BlockIdx != -1 && "Switch didn't go to this successor??");
 200     PN->setIncomingBlock((unsigned)BlockIdx, NewDefault);
 201   }
 202
 203   std::vector<Case> Cases;
 204
 205   // Expand comparisons for all of the non-default cases...
 206   for (unsigned i = 1; i < SI->getNumSuccessors(); ++i)
 207     Cases.push_back(Case(SI->getSuccessorValue(i), SI->getSuccessor(i)));
 208
 209   std::sort(Cases.begin(), Cases.end(), CaseCmp());
 210   DEBUG(std::cerr << "Cases: " << Cases << "\n");
 211   BasicBlock* SwitchBlock = switchConvert(Cases.begin(), Cases.end(), Val,
 212                                           OrigBlock, NewDefault);
 213
 214   // Branch to our shiny new if-then stuff...
 215   OrigBlock->getInstList().push_back(new BranchInst(SwitchBlock));
 216
 217   // We are now done with the switch instruction, delete it.
 218   delete SI;
 219 }