2d2a8a54a0f29ae1ca22008df355ccad4534a5f4
[oota-llvm.git] / lib / Transforms / Utils / LowerSwitch.cpp
1 //===- LowerSwitch.cpp - Eliminate Switch instructions --------------------===//
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 // The LowerSwitch transformation rewrites switch instructions with a sequence
11 // of branches, which allows targets to get away with not implementing the
12 // switch instruction until it is convenient.
13 //
14 //===----------------------------------------------------------------------===//
15
16 #include "llvm/Transforms/Scalar.h"
17 #include "llvm/ADT/STLExtras.h"
18 #include "llvm/IR/Constants.h"
19 #include "llvm/IR/Function.h"
20 #include "llvm/IR/Instructions.h"
21 #include "llvm/IR/LLVMContext.h"
22 #include "llvm/Pass.h"
23 #include "llvm/Support/Compiler.h"
24 #include "llvm/Support/Debug.h"
25 #include "llvm/Support/raw_ostream.h"
26 #include "llvm/Transforms/Utils/UnifyFunctionExitNodes.h"
27 #include <algorithm>
28 using namespace llvm;
29
30 namespace {
31   /// LowerSwitch Pass - Replace all SwitchInst instructions with chained branch
32   /// instructions.
33   class LowerSwitch : public FunctionPass {
34   public:
35     static char ID; // Pass identification, replacement for typeid
36     LowerSwitch() : FunctionPass(ID) {
37       initializeLowerSwitchPass(*PassRegistry::getPassRegistry());
38     } 
39
40     virtual bool runOnFunction(Function &F);
41     
42     virtual void getAnalysisUsage(AnalysisUsage &AU) const {
43       // This is a cluster of orthogonal Transforms
44       AU.addPreserved<UnifyFunctionExitNodes>();
45       AU.addPreserved("mem2reg");
46       AU.addPreservedID(LowerInvokePassID);
47     }
48
49     struct CaseRange {
50       Constant* Low;
51       Constant* High;
52       BasicBlock* BB;
53
54       CaseRange(Constant *low = 0, Constant *high = 0, BasicBlock *bb = 0) :
55         Low(low), High(high), BB(bb) { }
56     };
57
58     typedef std::vector<CaseRange>           CaseVector;
59     typedef std::vector<CaseRange>::iterator CaseItr;
60   private:
61     void processSwitchInst(SwitchInst *SI);
62
63     BasicBlock* switchConvert(CaseItr Begin, CaseItr End, Value* Val,
64                               BasicBlock* OrigBlock, BasicBlock* Default);
65     BasicBlock* newLeafBlock(CaseRange& Leaf, Value* Val,
66                              BasicBlock* OrigBlock, BasicBlock* Default);
67     unsigned Clusterify(CaseVector& Cases, SwitchInst *SI);
68   };
69
70   /// The comparison function for sorting the switch case values in the vector.
71   /// WARNING: Case ranges should be disjoint!
72   struct CaseCmp {
73     bool operator () (const LowerSwitch::CaseRange& C1,
74                       const LowerSwitch::CaseRange& C2) {
75
76       const ConstantInt* CI1 = cast<const ConstantInt>(C1.Low);
77       const ConstantInt* CI2 = cast<const ConstantInt>(C2.High);
78       return CI1->getValue().slt(CI2->getValue());
79     }
80   };
81 }
82
83 char LowerSwitch::ID = 0;
84 INITIALIZE_PASS(LowerSwitch, "lowerswitch",
85                 "Lower SwitchInst's to branches", false, false)
86
87 // Publicly exposed interface to pass...
88 char &llvm::LowerSwitchID = LowerSwitch::ID;
89 // createLowerSwitchPass - Interface to this file...
90 FunctionPass *llvm::createLowerSwitchPass() {
91   return new LowerSwitch();
92 }
93
94 bool LowerSwitch::runOnFunction(Function &F) {
95   bool Changed = false;
96
97   for (Function::iterator I = F.begin(), E = F.end(); I != E; ) {
98     BasicBlock *Cur = I++; // Advance over block so we don't traverse new blocks
99
100     if (SwitchInst *SI = dyn_cast<SwitchInst>(Cur->getTerminator())) {
101       Changed = true;
102       processSwitchInst(SI);
103     }
104   }
105
106   return Changed;
107 }
108
109 // operator<< - Used for debugging purposes.
110 //
111 static raw_ostream& operator<<(raw_ostream &O,
112                                const LowerSwitch::CaseVector &C)
113     LLVM_ATTRIBUTE_USED;
114 static raw_ostream& operator<<(raw_ostream &O,
115                                const LowerSwitch::CaseVector &C) {
116   O << "[";
117
118   for (LowerSwitch::CaseVector::const_iterator B = C.begin(),
119          E = C.end(); B != E; ) {
120     O << *B->Low << " -" << *B->High;
121     if (++B != E) O << ", ";
122   }
123
124   return O << "]";
125 }
126
127 // switchConvert - Convert the switch statement into a binary lookup of
128 // the case values. The function recursively builds this tree.
129 //
130 BasicBlock* LowerSwitch::switchConvert(CaseItr Begin, CaseItr End,
131                                        Value* Val, BasicBlock* OrigBlock,
132                                        BasicBlock* Default)
133 {
134   unsigned Size = End - Begin;
135
136   if (Size == 1)
137     return newLeafBlock(*Begin, Val, OrigBlock, Default);
138
139   unsigned Mid = Size / 2;
140   std::vector<CaseRange> LHS(Begin, Begin + Mid);
141   DEBUG(dbgs() << "LHS: " << LHS << "\n");
142   std::vector<CaseRange> RHS(Begin + Mid, End);
143   DEBUG(dbgs() << "RHS: " << RHS << "\n");
144
145   CaseRange& Pivot = *(Begin + Mid);
146   DEBUG(dbgs() << "Pivot ==> " 
147                << cast<ConstantInt>(Pivot.Low)->getValue() << " -"
148                << cast<ConstantInt>(Pivot.High)->getValue() << "\n");
149
150   BasicBlock* LBranch = switchConvert(LHS.begin(), LHS.end(), Val,
151                                       OrigBlock, Default);
152   BasicBlock* RBranch = switchConvert(RHS.begin(), RHS.end(), Val,
153                                       OrigBlock, Default);
154
155   // Create a new node that checks if the value is < pivot. Go to the
156   // left branch if it is and right branch if not.
157   Function* F = OrigBlock->getParent();
158   BasicBlock* NewNode = BasicBlock::Create(Val->getContext(), "NodeBlock");
159   Function::iterator FI = OrigBlock;
160   F->getBasicBlockList().insert(++FI, NewNode);
161
162   ICmpInst* Comp = new ICmpInst(ICmpInst::ICMP_SLT,
163                                 Val, Pivot.Low, "Pivot");
164   NewNode->getInstList().push_back(Comp);
165   BranchInst::Create(LBranch, RBranch, Comp, NewNode);
166   return NewNode;
167 }
168
169 // newLeafBlock - Create a new leaf block for the binary lookup tree. It
170 // checks if the switch's value == the case's value. If not, then it
171 // jumps to the default branch. At this point in the tree, the value
172 // can't be another valid case value, so the jump to the "default" branch
173 // is warranted.
174 //
175 BasicBlock* LowerSwitch::newLeafBlock(CaseRange& Leaf, Value* Val,
176                                       BasicBlock* OrigBlock,
177                                       BasicBlock* Default)
178 {
179   Function* F = OrigBlock->getParent();
180   BasicBlock* NewLeaf = BasicBlock::Create(Val->getContext(), "LeafBlock");
181   Function::iterator FI = OrigBlock;
182   F->getBasicBlockList().insert(++FI, NewLeaf);
183
184   // Emit comparison
185   ICmpInst* Comp = NULL;
186   if (Leaf.Low == Leaf.High) {
187     // Make the seteq instruction...
188     Comp = new ICmpInst(*NewLeaf, ICmpInst::ICMP_EQ, Val,
189                         Leaf.Low, "SwitchLeaf");
190   } else {
191     // Make range comparison
192     if (cast<ConstantInt>(Leaf.Low)->isMinValue(true /*isSigned*/)) {
193       // Val >= Min && Val <= Hi --> Val <= Hi
194       Comp = new ICmpInst(*NewLeaf, ICmpInst::ICMP_SLE, Val, Leaf.High,
195                           "SwitchLeaf");
196     } else if (cast<ConstantInt>(Leaf.Low)->isZero()) {
197       // Val >= 0 && Val <= Hi --> Val <=u Hi
198       Comp = new ICmpInst(*NewLeaf, ICmpInst::ICMP_ULE, Val, Leaf.High,
199                           "SwitchLeaf");      
200     } else {
201       // Emit V-Lo <=u Hi-Lo
202       Constant* NegLo = ConstantExpr::getNeg(Leaf.Low);
203       Instruction* Add = BinaryOperator::CreateAdd(Val, NegLo,
204                                                    Val->getName()+".off",
205                                                    NewLeaf);
206       Constant *UpperBound = ConstantExpr::getAdd(NegLo, Leaf.High);
207       Comp = new ICmpInst(*NewLeaf, ICmpInst::ICMP_ULE, Add, UpperBound,
208                           "SwitchLeaf");
209     }
210   }
211
212   // Make the conditional branch...
213   BasicBlock* Succ = Leaf.BB;
214   BranchInst::Create(Succ, Default, Comp, NewLeaf);
215
216   // If there were any PHI nodes in this successor, rewrite one entry
217   // from OrigBlock to come from NewLeaf.
218   for (BasicBlock::iterator I = Succ->begin(); isa<PHINode>(I); ++I) {
219     PHINode* PN = cast<PHINode>(I);
220     // Remove all but one incoming entries from the cluster
221     uint64_t Range = cast<ConstantInt>(Leaf.High)->getSExtValue() -
222                      cast<ConstantInt>(Leaf.Low)->getSExtValue();    
223     for (uint64_t j = 0; j < Range; ++j) {
224       PN->removeIncomingValue(OrigBlock);
225     }
226     
227     int BlockIdx = PN->getBasicBlockIndex(OrigBlock);
228     assert(BlockIdx != -1 && "Switch didn't go to this successor??");
229     PN->setIncomingBlock((unsigned)BlockIdx, NewLeaf);
230   }
231
232   return NewLeaf;
233 }
234
235 // Clusterify - Transform simple list of Cases into list of CaseRange's
236 unsigned LowerSwitch::Clusterify(CaseVector& Cases, SwitchInst *SI) {
237   unsigned numCmps = 0;
238
239   // Start with "simple" cases
240   for (SwitchInst::CaseIt i = SI->case_begin(), e = SI->case_end(); i != e; ++i)
241     Cases.push_back(CaseRange(i.getCaseValue(), i.getCaseValue(),
242                               i.getCaseSuccessor()));
243   
244   std::sort(Cases.begin(), Cases.end(), CaseCmp());
245
246   // Merge case into clusters
247   if (Cases.size()>=2)
248     for (CaseItr I=Cases.begin(), J=llvm::next(Cases.begin()); J!=Cases.end(); ) {
249       int64_t nextValue = cast<ConstantInt>(J->Low)->getSExtValue();
250       int64_t currentValue = cast<ConstantInt>(I->High)->getSExtValue();
251       BasicBlock* nextBB = J->BB;
252       BasicBlock* currentBB = I->BB;
253
254       // If the two neighboring cases go to the same destination, merge them
255       // into a single case.
256       if ((nextValue-currentValue==1) && (currentBB == nextBB)) {
257         I->High = J->High;
258         J = Cases.erase(J);
259       } else {
260         I = J++;
261       }
262     }
263
264   for (CaseItr I=Cases.begin(), E=Cases.end(); I!=E; ++I, ++numCmps) {
265     if (I->Low != I->High)
266       // A range counts double, since it requires two compares.
267       ++numCmps;
268   }
269
270   return numCmps;
271 }
272
273 // processSwitchInst - Replace the specified switch instruction with a sequence
274 // of chained if-then insts in a balanced binary search.
275 //
276 void LowerSwitch::processSwitchInst(SwitchInst *SI) {
277   BasicBlock *CurBlock = SI->getParent();
278   BasicBlock *OrigBlock = CurBlock;
279   Function *F = CurBlock->getParent();
280   Value *Val = SI->getCondition();  // The value we are switching on...
281   BasicBlock* Default = SI->getDefaultDest();
282
283   // If there is only the default destination, don't bother with the code below.
284   if (!SI->getNumCases()) {
285     BranchInst::Create(SI->getDefaultDest(), CurBlock);
286     CurBlock->getInstList().erase(SI);
287     return;
288   }
289
290   // Create a new, empty default block so that the new hierarchy of
291   // if-then statements go to this and the PHI nodes are happy.
292   BasicBlock* NewDefault = BasicBlock::Create(SI->getContext(), "NewDefault");
293   F->getBasicBlockList().insert(Default, NewDefault);
294
295   BranchInst::Create(Default, NewDefault);
296
297   // If there is an entry in any PHI nodes for the default edge, make sure
298   // to update them as well.
299   for (BasicBlock::iterator I = Default->begin(); isa<PHINode>(I); ++I) {
300     PHINode *PN = cast<PHINode>(I);
301     int BlockIdx = PN->getBasicBlockIndex(OrigBlock);
302     assert(BlockIdx != -1 && "Switch didn't go to this successor??");
303     PN->setIncomingBlock((unsigned)BlockIdx, NewDefault);
304   }
305
306   // Prepare cases vector.
307   CaseVector Cases;
308   unsigned numCmps = Clusterify(Cases, SI);
309
310   DEBUG(dbgs() << "Clusterify finished. Total clusters: " << Cases.size()
311                << ". Total compares: " << numCmps << "\n");
312   DEBUG(dbgs() << "Cases: " << Cases << "\n");
313   (void)numCmps;
314   
315   BasicBlock* SwitchBlock = switchConvert(Cases.begin(), Cases.end(), Val,
316                                           OrigBlock, NewDefault);
317
318   // Branch to our shiny new if-then stuff...
319   BranchInst::Create(SwitchBlock, OrigBlock);
320
321   // We are now done with the switch instruction, delete it.
322   CurBlock->getInstList().erase(SI);
323 }