1 files changed, 114 insertions, 4 deletions

diff --git a/lib/Transforms/Scalar/CodeGenPrepare.cpp b/lib/Transforms/Scalar/CodeGenPrepare.cpp
index 65249153bd..495cdc6321 100644
--- a/lib/Transforms/Scalar/CodeGenPrepare.cpp
+++ b/lib/Transforms/Scalar/CodeGenPrepare.cpp
@@ -18,6 +18,7 @@
 #include "llvm/Constants.h"
 #include "llvm/DerivedTypes.h"
 #include "llvm/Function.h"
+#include "llvm/GlobalVariable.h"
 #include "llvm/IRBuilder.h"
 #include "llvm/InlineAsm.h"
 #include "llvm/Instructions.h"
@@ -126,6 +127,7 @@ namespace {
     bool OptimizeSelectInst(SelectInst *SI);
     bool DupRetToEnableTailCallOpts(ReturnInst *RI);
     bool PlaceDbgValues(Function &F);
+    bool ConvertLoadToSwitch(LoadInst *LI);
   };
 }
 
@@ -169,7 +171,7 @@ bool CodeGenPrepare::runOnFunction(Function &F) {
   bool MadeChange = true;
   while (MadeChange) {
     MadeChange = false;
-    for (Function::iterator I = F.begin(), E = F.end(); I != E; ) {
+    for (Function::iterator I = F.begin(); I != F.end(); ) {
       BasicBlock *BB = I++;
       MadeChange |= OptimizeBlock(*BB);
     }
@@ -1283,9 +1285,11 @@ bool CodeGenPrepare::OptimizeInst(Instruction *I) {
     return OptimizeCmpExpression(CI);
 
   if (LoadInst *LI = dyn_cast<LoadInst>(I)) {
+    bool Changed = false;
     if (TLI)
-      return OptimizeMemoryInst(I, I->getOperand(0), LI->getType());
-    return false;
+      Changed |= OptimizeMemoryInst(I, I->getOperand(0), LI->getType());
+    Changed |= ConvertLoadToSwitch(LI);
+    return Changed;
   }
 
   if (StoreInst *SI = dyn_cast<StoreInst>(I)) {
@@ -1329,7 +1333,7 @@ bool CodeGenPrepare::OptimizeBlock(BasicBlock &BB) {
   bool MadeChange = false;
 
   CurInstIterator = BB.begin();
-  for (BasicBlock::iterator E = BB.end(); CurInstIterator != E; )
+  while (CurInstIterator != BB.end())
     MadeChange |= OptimizeInst(CurInstIterator++);
 
   return MadeChange;
@@ -1365,3 +1369,109 @@ bool CodeGenPrepare::PlaceDbgValues(Function &F) {
   }
   return MadeChange;
 }
+
+static bool TargetSupportsJumpTables(const TargetLowering &TLI) {
+  return TLI.supportJumpTables() &&
+          (TLI.isOperationLegalOrCustom(ISD::BR_JT, MVT::Other) ||
+           TLI.isOperationLegalOrCustom(ISD::BRIND, MVT::Other));
+}
+
+/// ConvertLoadToSwitch - Convert loads from constant lookup tables into
+/// switches. This undos the switch-to-lookup table transformation in
+/// SimplifyCFG for targets where that is inprofitable.
+bool CodeGenPrepare::ConvertLoadToSwitch(LoadInst *LI) {
+  // This only applies to targets that don't support jump tables.
+  if (!TLI || TargetSupportsJumpTables(*TLI))
+    return false;
+
+  // FIXME: In the future, it would be desirable to have enough target
+  // information in SimplifyCFG, so we could decide at that stage whether to
+  // transform the switch to a lookup table or not, and this
+  // reverse-transformation could be removed.
+
+  GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(LI->getPointerOperand());
+  if (!GEP || !GEP->isInBounds() || GEP->getPointerAddressSpace())
+    return false;
+  if (GEP->getNumIndices() != 2)
+    return false;
+  Value *FirstIndex = GEP->idx_begin()[0];
+  ConstantInt *FirstIndexInt = dyn_cast<ConstantInt>(FirstIndex);
+  if (!FirstIndexInt || !FirstIndexInt->isZero())
+    return false;
+
+  Value *TableIndex = GEP->idx_begin()[1];
+  IntegerType *TableIndexTy = cast<IntegerType>(TableIndex->getType());
+
+  GlobalVariable *GV = dyn_cast<GlobalVariable>(GEP->getPointerOperand());
+  if (!GV || !GV->isConstant() || !GV->hasDefinitiveInitializer())
+    return false;
+
+  Constant *Arr = GV->getInitializer();
+  uint64_t NumElements;
+  if (ConstantArray *CA = dyn_cast<ConstantArray>(Arr))
+    NumElements = CA->getType()->getNumElements();
+  else if (ConstantDataArray *CDA = dyn_cast<ConstantDataArray>(Arr))
+    NumElements = CDA->getNumElements();
+  else
+    return false;
+  if (NumElements < 2)
+    return false;
+
+  // Split the block.
+  BasicBlock *OriginalBB = LI->getParent();
+  BasicBlock *PostSwitchBB = OriginalBB->splitBasicBlock(LI);
+
+  // Replace OriginalBB's terminator with a switch.
+  IRBuilder<> Builder(OriginalBB->getTerminator());
+  SwitchInst *Switch = Builder.CreateSwitch(TableIndex, PostSwitchBB,
+                                            NumElements - 1);
+  OriginalBB->getTerminator()->eraseFromParent();
+
+  // Count the frequency of each value to decide which to use as default.
+  SmallDenseMap<Constant*, uint64_t> ValueFreq;
+  for (uint64_t I = 0; I < NumElements; ++I)
+    ++ValueFreq[Arr->getAggregateElement(I)];
+  uint64_t MaxCount = 0;
+  Constant *DefaultValue = NULL;
+  for (SmallDenseMap<Constant*, uint64_t>::iterator I = ValueFreq.begin(),
+       E = ValueFreq.end(); I != E; ++I) {
+    if (I->second > MaxCount) {
+      MaxCount = I->second;
+      DefaultValue = I->first;
+    }
+  }
+  assert(DefaultValue && "No values in the array?");
+
+  // Create the phi node in PostSwitchBB, which will replace the load.
+  Builder.SetInsertPoint(PostSwitchBB->begin());
+  PHINode *PHI = Builder.CreatePHI(LI->getType(), NumElements);
+  PHI->addIncoming(DefaultValue, OriginalBB);
+
+  // Build basic blocks to target with the switch.
+  for (uint64_t I = 0; I < NumElements; ++I) {
+    Constant *C = Arr->getAggregateElement(I);
+    if (C == DefaultValue) continue; // Already covered by the default case.
+
+    BasicBlock *BB = BasicBlock::Create(PostSwitchBB->getContext(),
+                                        "lookup.bb",
+                                        PostSwitchBB->getParent(),
+                                        PostSwitchBB);
+    Switch->addCase(ConstantInt::get(TableIndexTy, I), BB);
+    Builder.SetInsertPoint(BB);
+    Builder.CreateBr(PostSwitchBB);
+    PHI->addIncoming(C, BB);
+  }
+
+  // Remove the load.
+  LI->replaceAllUsesWith(PHI);
+  LI->eraseFromParent();
+
+  // Clean up.
+  if (GEP->use_empty())
+    GEP->eraseFromParent();
+  if (GV->hasUnnamedAddr() && GV->hasPrivateLinkage() && GV->use_empty())
+    GV->eraseFromParent();
+
+  CurInstIterator = Switch;
+  return true;
+}