Patch for function cloning to inline all blocks whose address is taken
[oota-llvm.git] / lib / Transforms / Utils / CloneFunction.cpp
index e56260677edeea68ffd0993c3cad4eb940a776f0..5c2bbd19564459135783093fabd526527e4d1ae7 100644 (file)
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Utils/Cloning.h"
-#include "llvm/Constants.h"
-#include "llvm/DerivedTypes.h"
-#include "llvm/Instructions.h"
-#include "llvm/IntrinsicInst.h"
-#include "llvm/GlobalVariable.h"
-#include "llvm/Function.h"
-#include "llvm/LLVMContext.h"
-#include "llvm/Support/CFG.h"
-#include "llvm/Transforms/Utils/ValueMapper.h"
-#include "llvm/Analysis/ConstantFolding.h"
-#include "llvm/Analysis/DebugInfo.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/Analysis/ConstantFolding.h"
+#include "llvm/Analysis/InstructionSimplify.h"
+#include "llvm/IR/CFG.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/DebugInfo.h"
+#include "llvm/IR/DerivedTypes.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/GlobalVariable.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Metadata.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/Local.h"
+#include "llvm/Transforms/Utils/ValueMapper.h"
 #include <map>
+#include <set>
 using namespace llvm;
 
 // CloneBasicBlock - See comments in Cloning.h
 BasicBlock *llvm::CloneBasicBlock(const BasicBlock *BB,
-                                  DenseMap<const Value*, Value*> &ValueMap,
-                                  const char *NameSuffix, Function *F,
+                                  ValueToValueMapTy &VMap,
+                                  const Twine &NameSuffix, Function *F,
                                   ClonedCodeInfo *CodeInfo) {
   BasicBlock *NewBB = BasicBlock::Create(BB->getContext(), "", F);
   if (BB->hasName()) NewBB->setName(BB->getName()+NameSuffix);
@@ -46,7 +52,7 @@ BasicBlock *llvm::CloneBasicBlock(const BasicBlock *BB,
     if (II->hasName())
       NewInst->setName(II->getName()+NameSuffix);
     NewBB->getInstList().push_back(NewInst);
-    ValueMap[II] = NewInst;                // Add instruction map to value.
+    VMap[II] = NewInst;                // Add instruction map to value.
     
     hasCalls |= (isa<CallInst>(II) && !isa<DbgInfoIntrinsic>(II));
     if (const AllocaInst *AI = dyn_cast<AllocaInst>(II)) {
@@ -59,7 +65,6 @@ BasicBlock *llvm::CloneBasicBlock(const BasicBlock *BB,
   
   if (CodeInfo) {
     CodeInfo->ContainsCalls          |= hasCalls;
-    CodeInfo->ContainsUnwinds        |= isa<UnwindInst>(BB->getTerminator());
     CodeInfo->ContainsDynamicAllocas |= hasDynamicAllocas;
     CodeInfo->ContainsDynamicAllocas |= hasStaticAllocas && 
                                         BB != &BB->getParent()->getEntryBlock();
@@ -68,38 +73,45 @@ BasicBlock *llvm::CloneBasicBlock(const BasicBlock *BB,
 }
 
 // Clone OldFunc into NewFunc, transforming the old arguments into references to
-// ArgMap values.
+// VMap values.
 //
 void llvm::CloneFunctionInto(Function *NewFunc, const Function *OldFunc,
-                             DenseMap<const Value*, Value*> &ValueMap,
+                             ValueToValueMapTy &VMap,
+                             bool ModuleLevelChanges,
                              SmallVectorImpl<ReturnInst*> &Returns,
-                             const char *NameSuffix, ClonedCodeInfo *CodeInfo) {
+                             const char *NameSuffix, ClonedCodeInfo *CodeInfo,
+                             ValueMapTypeRemapper *TypeMapper,
+                             ValueMaterializer *Materializer) {
   assert(NameSuffix && "NameSuffix cannot be null!");
 
 #ifndef NDEBUG
   for (Function::const_arg_iterator I = OldFunc->arg_begin(), 
        E = OldFunc->arg_end(); I != E; ++I)
-    assert(ValueMap.count(I) && "No mapping from source argument specified!");
+    assert(VMap.count(I) && "No mapping from source argument specified!");
 #endif
 
-  // Clone any attributes.
-  if (NewFunc->arg_size() == OldFunc->arg_size())
-    NewFunc->copyAttributesFrom(OldFunc);
-  else {
-    //Some arguments were deleted with the ValueMap. Copy arguments one by one
-    for (Function::const_arg_iterator I = OldFunc->arg_begin(), 
-           E = OldFunc->arg_end(); I != E; ++I)
-      if (Argument* Anew = dyn_cast<Argument>(ValueMap[I]))
-        Anew->addAttr( OldFunc->getAttributes()
-                       .getParamAttributes(I->getArgNo() + 1));
-    NewFunc->setAttributes(NewFunc->getAttributes()
-                           .addAttr(0, OldFunc->getAttributes()
-                                     .getRetAttributes()));
-    NewFunc->setAttributes(NewFunc->getAttributes()
-                           .addAttr(~0, OldFunc->getAttributes()
-                                     .getFnAttributes()));
+  // Copy all attributes other than those stored in the AttributeSet.  We need
+  // to remap the parameter indices of the AttributeSet.
+  AttributeSet NewAttrs = NewFunc->getAttributes();
+  NewFunc->copyAttributesFrom(OldFunc);
+  NewFunc->setAttributes(NewAttrs);
+
+  AttributeSet OldAttrs = OldFunc->getAttributes();
+  // Clone any argument attributes that are present in the VMap.
+  for (const Argument &OldArg : OldFunc->args())
+    if (Argument *NewArg = dyn_cast<Argument>(VMap[&OldArg])) {
+      AttributeSet attrs =
+          OldAttrs.getParamAttributes(OldArg.getArgNo() + 1);
+      if (attrs.getNumSlots() > 0)
+        NewArg->addAttr(attrs);
+    }
 
-  }
+  NewFunc->setAttributes(
+      NewFunc->getAttributes()
+          .addAttributes(NewFunc->getContext(), AttributeSet::ReturnIndex,
+                         OldAttrs.getRetAttributes())
+          .addAttributes(NewFunc->getContext(), AttributeSet::FunctionIndex,
+                         OldAttrs.getFnAttributes()));
 
   // Loop over all of the basic blocks in the function, cloning them as
   // appropriate.  Note that we save BE this way in order to handle cloning of
@@ -110,43 +122,106 @@ void llvm::CloneFunctionInto(Function *NewFunc, const Function *OldFunc,
     const BasicBlock &BB = *BI;
 
     // Create a new basic block and copy instructions into it!
-    BasicBlock *CBB = CloneBasicBlock(&BB, ValueMap, NameSuffix, NewFunc,
-                                      CodeInfo);
-    ValueMap[&BB] = CBB;                       // Add basic block mapping.
+    BasicBlock *CBB = CloneBasicBlock(&BB, VMap, NameSuffix, NewFunc, CodeInfo);
+
+    // Add basic block mapping.
+    VMap[&BB] = CBB;
+
+    // It is only legal to clone a function if a block address within that
+    // function is never referenced outside of the function.  Given that, we
+    // want to map block addresses from the old function to block addresses in
+    // the clone. (This is different from the generic ValueMapper
+    // implementation, which generates an invalid blockaddress when
+    // cloning a function.)
+    if (BB.hasAddressTaken()) {
+      Constant *OldBBAddr = BlockAddress::get(const_cast<Function*>(OldFunc),
+                                              const_cast<BasicBlock*>(&BB));
+      VMap[OldBBAddr] = BlockAddress::get(NewFunc, CBB);                                         
+    }
 
+    // Note return instructions for the caller.
     if (ReturnInst *RI = dyn_cast<ReturnInst>(CBB->getTerminator()))
       Returns.push_back(RI);
   }
 
   // Loop over all of the instructions in the function, fixing up operand
-  // references as we go.  This uses ValueMap to do all the hard work.
-  //
-  for (Function::iterator BB = cast<BasicBlock>(ValueMap[OldFunc->begin()]),
+  // references as we go.  This uses VMap to do all the hard work.
+  for (Function::iterator BB = cast<BasicBlock>(VMap[OldFunc->begin()]),
          BE = NewFunc->end(); BB != BE; ++BB)
     // Loop over all instructions, fixing each one as we find it...
     for (BasicBlock::iterator II = BB->begin(); II != BB->end(); ++II)
-      RemapInstruction(II, ValueMap);
+      RemapInstruction(II, VMap,
+                       ModuleLevelChanges ? RF_None : RF_NoModuleLevelChanges,
+                       TypeMapper, Materializer);
+}
+
+// Find the MDNode which corresponds to the DISubprogram data that described F.
+static MDNode* FindSubprogram(const Function *F, DebugInfoFinder &Finder) {
+  for (DISubprogram Subprogram : Finder.subprograms()) {
+    if (Subprogram.describes(F)) return Subprogram;
+  }
+  return nullptr;
+}
+
+// Add an operand to an existing MDNode. The new operand will be added at the
+// back of the operand list.
+static void AddOperand(MDNode *Node, Value *Operand) {
+  SmallVector<Value*, 16> Operands;
+  for (unsigned i = 0; i < Node->getNumOperands(); i++) {
+    Operands.push_back(Node->getOperand(i));
+  }
+  Operands.push_back(Operand);
+  MDNode *NewNode = MDNode::get(Node->getContext(), Operands);
+  Node->replaceAllUsesWith(NewNode);
+}
+
+// Clone the module-level debug info associated with OldFunc. The cloned data
+// will point to NewFunc instead.
+static void CloneDebugInfoMetadata(Function *NewFunc, const Function *OldFunc,
+                            ValueToValueMapTy &VMap) {
+  DebugInfoFinder Finder;
+  Finder.processModule(*OldFunc->getParent());
+
+  const MDNode *OldSubprogramMDNode = FindSubprogram(OldFunc, Finder);
+  if (!OldSubprogramMDNode) return;
+
+  // Ensure that OldFunc appears in the map.
+  // (if it's already there it must point to NewFunc anyway)
+  VMap[OldFunc] = NewFunc;
+  DISubprogram NewSubprogram(MapValue(OldSubprogramMDNode, VMap));
+
+  for (DICompileUnit CU : Finder.compile_units()) {
+    DIArray Subprograms(CU.getSubprograms());
+
+    // If the compile unit's function list contains the old function, it should
+    // also contain the new one.
+    for (unsigned i = 0; i < Subprograms.getNumElements(); i++) {
+      if ((MDNode*)Subprograms.getElement(i) == OldSubprogramMDNode) {
+        AddOperand(Subprograms, NewSubprogram);
+      }
+    }
+  }
 }
 
 /// CloneFunction - Return a copy of the specified function, but without
 /// embedding the function into another module.  Also, any references specified
-/// in the ValueMap are changed to refer to their mapped value instead of the
-/// original one.  If any of the arguments to the function are in the ValueMap,
-/// the arguments are deleted from the resultant function.  The ValueMap is
+/// in the VMap are changed to refer to their mapped value instead of the
+/// original one.  If any of the arguments to the function are in the VMap,
+/// the arguments are deleted from the resultant function.  The VMap is
 /// updated to include mappings from all of the instructions and basicblocks in
 /// the function from their old to new values.
 ///
-Function *llvm::CloneFunction(const Function *F,
-                              DenseMap<const Value*, Value*> &ValueMap,
+Function *llvm::CloneFunction(const Function *F, ValueToValueMapTy &VMap,
+                              bool ModuleLevelChanges,
                               ClonedCodeInfo *CodeInfo) {
-  std::vector<const Type*> ArgTypes;
+  std::vector<Type*> ArgTypes;
 
   // The user might be deleting arguments to the function by specifying them in
-  // the ValueMap.  If so, we need to not add the arguments to the arg ty vector
+  // the VMap.  If so, we need to not add the arguments to the arg ty vector
   //
   for (Function::const_arg_iterator I = F->arg_begin(), E = F->arg_end();
        I != E; ++I)
-    if (ValueMap.count(I) == 0)  // Haven't mapped the argument to anything yet?
+    if (VMap.count(I) == 0)  // Haven't mapped the argument to anything yet?
       ArgTypes.push_back(I->getType());
 
   // Create a new function type...
@@ -160,13 +235,16 @@ Function *llvm::CloneFunction(const Function *F,
   Function::arg_iterator DestI = NewF->arg_begin();
   for (Function::const_arg_iterator I = F->arg_begin(), E = F->arg_end();
        I != E; ++I)
-    if (ValueMap.count(I) == 0) {   // Is this argument preserved?
+    if (VMap.count(I) == 0) {   // Is this argument preserved?
       DestI->setName(I->getName()); // Copy the name over...
-      ValueMap[I] = DestI++;        // Add mapping to ValueMap
+      VMap[I] = DestI++;        // Add mapping to VMap
     }
 
+  if (ModuleLevelChanges)
+    CloneDebugInfoMetadata(NewF, F, VMap);
+
   SmallVector<ReturnInst*, 8> Returns;  // Ignore returns cloned.
-  CloneFunctionInto(NewF, F, ValueMap, Returns, "", CodeInfo);
+  CloneFunctionInto(NewF, F, VMap, ModuleLevelChanges, Returns, "", CodeInfo);
   return NewF;
 }
 
@@ -178,48 +256,53 @@ namespace {
   struct PruningFunctionCloner {
     Function *NewFunc;
     const Function *OldFunc;
-    DenseMap<const Value*, Value*> &ValueMap;
-    SmallVectorImpl<ReturnInst*> &Returns;
+    ValueToValueMapTy &VMap;
+    bool ModuleLevelChanges;
     const char *NameSuffix;
     ClonedCodeInfo *CodeInfo;
-    const TargetData *TD;
-    Value *DbgFnStart;
+    const DataLayout *DL;
   public:
     PruningFunctionCloner(Function *newFunc, const Function *oldFunc,
-                          DenseMap<const Value*, Value*> &valueMap,
-                          SmallVectorImpl<ReturnInst*> &returns,
+                          ValueToValueMapTy &valueMap,
+                          bool moduleLevelChanges,
                           const char *nameSuffix, 
                           ClonedCodeInfo *codeInfo,
-                          const TargetData *td)
-    : NewFunc(newFunc), OldFunc(oldFunc), ValueMap(valueMap), Returns(returns),
-      NameSuffix(nameSuffix), CodeInfo(codeInfo), TD(td), DbgFnStart(NULL) {
+                          const DataLayout *DL)
+    : NewFunc(newFunc), OldFunc(oldFunc),
+      VMap(valueMap), ModuleLevelChanges(moduleLevelChanges),
+      NameSuffix(nameSuffix), CodeInfo(codeInfo), DL(DL) {
     }
 
-    /// CloneBlock - The specified block is found to be reachable, clone it and
-    /// anything that it can reach.
+    /// CloneBlock - The specified block is found to be reachable, so clone it
+    /// into NewBB.
+    /// ToClone is the vector of actually cloned blocks.
+    /// OrigBBs is the set of all blocks reacheable from the entry block.
+    /// It contains the block candidates and makes sure each block
+    /// is cloned at most once.
     void CloneBlock(const BasicBlock *BB,
-                    std::vector<const BasicBlock*> &ToClone);
-    
-  public:
-    /// ConstantFoldMappedInstruction - Constant fold the specified instruction,
-    /// mapping its operands through ValueMap if they are available.
-    Constant *ConstantFoldMappedInstruction(const Instruction *I);
+                    BasicBlock *NewBB,
+                    std::vector<const BasicBlock *> &ToClone,
+                    std::set<const BasicBlock *> &OrigBBs);
   };
 }
 
-/// CloneBlock - The specified block is found to be reachable, clone it and
-/// anything that it can reach.
+/// CloneBlock - The specified block is found to be reachable, so clone it
+/// into NewBB.
+/// ToClone is the vector of actually cloned blocks.
+/// OrigBBs is the set of all blocks reacheable from the entry block.
+/// It contains the block candidates and makes sure each block
+/// is cloned at most once.
 void PruningFunctionCloner::CloneBlock(const BasicBlock *BB,
-                                       std::vector<const BasicBlock*> &ToClone){
-  Value *&BBEntry = ValueMap[BB];
-
-  // Have we already cloned this block?
-  if (BBEntry) return;
+                                       BasicBlock *NewBB,
+                                       std::vector<const BasicBlock *> &ToClone,
+                                       std::set<const BasicBlock *> &OrigBBs) {
   
+  // Remove BB from list of blocks to clone.
+  // When it was not in the list, it has been cloned already, so
+  // don't clone again.
+  if (!OrigBBs.erase(BB)) return;
+
   // Nope, clone it now.
-  BasicBlock *NewBB;
-  BBEntry = NewBB = BasicBlock::Create(BB->getContext());
-  if (BB->hasName()) NewBB->setName(BB->getName()+NameSuffix);
 
   bool hasCalls = false, hasDynamicAllocas = false, hasStaticAllocas = false;
   
@@ -227,32 +310,33 @@ void PruningFunctionCloner::CloneBlock(const BasicBlock *BB,
   // loop doesn't include the terminator.
   for (BasicBlock::const_iterator II = BB->begin(), IE = --BB->end();
        II != IE; ++II) {
-    // If this instruction constant folds, don't bother cloning the instruction,
-    // instead, just add the constant to the value map.
-    if (Constant *C = ConstantFoldMappedInstruction(II)) {
-      ValueMap[II] = C;
-      continue;
-    }
+    Instruction *NewInst = II->clone();
 
-    // Do not clone llvm.dbg.region.end. It will be adjusted by the inliner.
-    if (const DbgFuncStartInst *DFSI = dyn_cast<DbgFuncStartInst>(II)) {
-      if (DbgFnStart == NULL) {
-        DISubprogram SP(DFSI->getSubprogram());
-        if (SP.describes(BB->getParent()))
-          DbgFnStart = DFSI->getSubprogram();
-      }
-    } 
-    if (const DbgRegionEndInst *DREIS = dyn_cast<DbgRegionEndInst>(II)) {
-      if (DREIS->getContext() == DbgFnStart)
+    // Eagerly remap operands to the newly cloned instruction, except for PHI
+    // nodes for which we defer processing until we update the CFG.
+    if (!isa<PHINode>(NewInst)) {
+      RemapInstruction(NewInst, VMap,
+                       ModuleLevelChanges ? RF_None : RF_NoModuleLevelChanges);
+
+      // If we can simplify this instruction to some other value, simply add
+      // a mapping to that value rather than inserting a new instruction into
+      // the basic block.
+      if (Value *V = SimplifyInstruction(NewInst, DL)) {
+        // On the off-chance that this simplifies to an instruction in the old
+        // function, map it back into the new function.
+        if (Value *MappedV = VMap.lookup(V))
+          V = MappedV;
+
+        VMap[II] = V;
+        delete NewInst;
         continue;
+      }
     }
-      
-    Instruction *NewInst = II->clone();
+
     if (II->hasName())
       NewInst->setName(II->getName()+NameSuffix);
+    VMap[II] = NewInst;                // Add instruction map to value.
     NewBB->getInstList().push_back(NewInst);
-    ValueMap[II] = NewInst;                // Add instruction map to value.
-    
     hasCalls |= (isa<CallInst>(II) && !isa<DbgInfoIntrinsic>(II));
     if (const AllocaInst *AI = dyn_cast<AllocaInst>(II)) {
       if (isa<ConstantInt>(AI->getArraySize()))
@@ -270,13 +354,15 @@ void PruningFunctionCloner::CloneBlock(const BasicBlock *BB,
       // If the condition was a known constant in the callee...
       ConstantInt *Cond = dyn_cast<ConstantInt>(BI->getCondition());
       // Or is a known constant in the caller...
-      if (Cond == 0)  
-        Cond = dyn_cast_or_null<ConstantInt>(ValueMap[BI->getCondition()]);
+      if (!Cond) {
+        Value *V = VMap[BI->getCondition()];
+        Cond = dyn_cast_or_null<ConstantInt>(V);
+      }
 
       // Constant fold to uncond branch!
       if (Cond) {
         BasicBlock *Dest = BI->getSuccessor(!Cond->getZExtValue());
-        ValueMap[OldTI] = BranchInst::Create(Dest, NewBB);
+        VMap[OldTI] = BranchInst::Create(Dest, NewBB);
         ToClone.push_back(Dest);
         TerminatorDone = true;
       }
@@ -284,11 +370,14 @@ void PruningFunctionCloner::CloneBlock(const BasicBlock *BB,
   } else if (const SwitchInst *SI = dyn_cast<SwitchInst>(OldTI)) {
     // If switching on a value known constant in the caller.
     ConstantInt *Cond = dyn_cast<ConstantInt>(SI->getCondition());
-    if (Cond == 0)  // Or known constant after constant prop in the callee...
-      Cond = dyn_cast_or_null<ConstantInt>(ValueMap[SI->getCondition()]);
+    if (!Cond) { // Or known constant after constant prop in the callee...
+      Value *V = VMap[SI->getCondition()];
+      Cond = dyn_cast_or_null<ConstantInt>(V);
+    }
     if (Cond) {     // Constant fold to uncond branch!
-      BasicBlock *Dest = SI->getSuccessor(SI->findCaseValue(Cond));
-      ValueMap[OldTI] = BranchInst::Create(Dest, NewBB);
+      SwitchInst::ConstCaseIt Case = SI->findCaseValue(Cond);
+      BasicBlock *Dest = const_cast<BasicBlock*>(Case.getCaseSuccessor());
+      VMap[OldTI] = BranchInst::Create(Dest, NewBB);
       ToClone.push_back(Dest);
       TerminatorDone = true;
     }
@@ -299,7 +388,7 @@ void PruningFunctionCloner::CloneBlock(const BasicBlock *BB,
     if (OldTI->hasName())
       NewInst->setName(OldTI->getName()+NameSuffix);
     NewBB->getInstList().push_back(NewInst);
-    ValueMap[OldTI] = NewInst;             // Add instruction map to value.
+    VMap[OldTI] = NewInst;             // Add instruction map to value.
     
     // Recursively clone any reachable successor blocks.
     const TerminatorInst *TI = BB->getTerminator();
@@ -309,63 +398,10 @@ void PruningFunctionCloner::CloneBlock(const BasicBlock *BB,
   
   if (CodeInfo) {
     CodeInfo->ContainsCalls          |= hasCalls;
-    CodeInfo->ContainsUnwinds        |= isa<UnwindInst>(OldTI);
     CodeInfo->ContainsDynamicAllocas |= hasDynamicAllocas;
     CodeInfo->ContainsDynamicAllocas |= hasStaticAllocas && 
       BB != &BB->getParent()->front();
   }
-  
-  if (ReturnInst *RI = dyn_cast<ReturnInst>(NewBB->getTerminator()))
-    Returns.push_back(RI);
-}
-
-/// ConstantFoldMappedInstruction - Constant fold the specified instruction,
-/// mapping its operands through ValueMap if they are available.
-Constant *PruningFunctionCloner::
-ConstantFoldMappedInstruction(const Instruction *I) {
-  SmallVector<Constant*, 8> Ops;
-  for (unsigned i = 0, e = I->getNumOperands(); i != e; ++i)
-    if (Constant *Op = dyn_cast_or_null<Constant>(MapValue(I->getOperand(i),
-                                                           ValueMap)))
-      Ops.push_back(Op);
-    else
-      return 0;  // All operands not constant!
-
-  if (const CmpInst *CI = dyn_cast<CmpInst>(I))
-    return ConstantFoldCompareInstOperands(CI->getPredicate(), Ops[0], Ops[1],
-                                           TD);
-
-  if (const LoadInst *LI = dyn_cast<LoadInst>(I))
-    if (ConstantExpr *CE = dyn_cast<ConstantExpr>(Ops[0]))
-      if (!LI->isVolatile() && CE->getOpcode() == Instruction::GetElementPtr)
-        if (GlobalVariable *GV = dyn_cast<GlobalVariable>(CE->getOperand(0)))
-          if (GV->isConstant() && GV->hasDefinitiveInitializer())
-            return ConstantFoldLoadThroughGEPConstantExpr(GV->getInitializer(),
-                                                          CE);
-
-  return ConstantFoldInstOperands(I->getOpcode(), I->getType(), &Ops[0],
-                                  Ops.size(), TD);
-}
-
-static MDNode *UpdateInlinedAtInfo(MDNode *InsnMD, MDNode *TheCallMD,
-                                   LLVMContext &Context) {
-  DILocation ILoc(InsnMD);
-  if (ILoc.isNull()) return InsnMD;
-
-  DILocation CallLoc(TheCallMD);
-  if (CallLoc.isNull()) return InsnMD;
-
-  DILocation OrigLocation = ILoc.getOrigLocation();
-  MDNode *NewLoc = TheCallMD;
-  if (!OrigLocation.isNull())
-    NewLoc = UpdateInlinedAtInfo(OrigLocation.getNode(), TheCallMD, Context);
-
-  SmallVector<Value *, 4> MDVs;
-  MDVs.push_back(InsnMD->getElement(0)); // Line
-  MDVs.push_back(InsnMD->getElement(1)); // Col
-  MDVs.push_back(InsnMD->getElement(2)); // Scope
-  MDVs.push_back(NewLoc);
-  return MDNode::get(Context, MDVs.data(), MDVs.size());
 }
 
 /// CloneAndPruneFunctionInto - This works exactly like CloneFunctionInto,
@@ -376,32 +412,109 @@ static MDNode *UpdateInlinedAtInfo(MDNode *InsnMD, MDNode *TheCallMD,
 /// dead.  Since this doesn't produce an exact copy of the input, it can't be
 /// used for things like CloneFunction or CloneModule.
 void llvm::CloneAndPruneFunctionInto(Function *NewFunc, const Function *OldFunc,
-                                     DenseMap<const Value*, Value*> &ValueMap,
+                                     ValueToValueMapTy &VMap,
+                                     bool ModuleLevelChanges,
                                      SmallVectorImpl<ReturnInst*> &Returns,
                                      const char *NameSuffix, 
                                      ClonedCodeInfo *CodeInfo,
-                                     const TargetData *TD,
+                                     const DataLayout *DL,
                                      Instruction *TheCall) {
   assert(NameSuffix && "NameSuffix cannot be null!");
-  
+
 #ifndef NDEBUG
   for (Function::const_arg_iterator II = OldFunc->arg_begin(), 
        E = OldFunc->arg_end(); II != E; ++II)
-    assert(ValueMap.count(II) && "No mapping from source argument specified!");
+    assert(VMap.count(II) && "No mapping from source argument specified!");
 #endif
 
-  PruningFunctionCloner PFC(NewFunc, OldFunc, ValueMap, Returns,
-                            NameSuffix, CodeInfo, TD);
+  PruningFunctionCloner PFC(NewFunc, OldFunc, VMap, ModuleLevelChanges,
+                            NameSuffix, CodeInfo, DL);
 
-  // Clone the entry block, and anything recursively reachable from it.
+  // Since all BB address references need to be known before block-by-block
+  // processing, we need to create all reachable blocks before processing
+  // them for instruction cloning and pruning. Some of these blocks may
+  // be removed due to later pruning.
   std::vector<const BasicBlock*> CloneWorklist;
+  //
+  // OrigBBs consists of all blocks reachable from the entry
+  // block.
+  // This list will be pruned down by the CloneFunction() due to two
+  // two optimizations:
+  // First, when a conditional branch target is known at compile-time,
+  // only the actual branch destination block needs to be cloned.
+  // Second, when a switch statement target is known at compile-time,
+  // only the actual case statement needs to be cloned.
+  std::set<const BasicBlock *> OrigBBs;
+
   CloneWorklist.push_back(&OldFunc->getEntryBlock());
   while (!CloneWorklist.empty()) {
     const BasicBlock *BB = CloneWorklist.back();
     CloneWorklist.pop_back();
-    PFC.CloneBlock(BB, CloneWorklist);
+
+    // Don't revisit blocks.
+    if (VMap.count(BB))
+      continue;
+
+    BasicBlock *NewBB = BasicBlock::Create(BB->getContext());
+    if (BB->hasName())
+      NewBB->setName(BB->getName() + NameSuffix);
+
+    // It is legal to clone a function when a block address within that
+    // function is never escapes outside of the function.  Given that, we
+    // want to map block addresses from the old function to block addresses in
+    // the clone. (This is different from the generic ValueMapper
+    // implementation, which generates an invalid block address when
+    // cloning a function.)
+    // Note the current escape address does not catch all legal cases: even
+    // when all block addresses taken are local and the function has the
+    // always_inline attribute due to the indirect branch inlining is
+    // suppressed.
+    // Note that we don't need to fix the mapping for unreachable blocks;
+    // the default mapping there is safe.
+    if (BB->hasAddressTaken()) {
+      Constant *OldBBAddr = BlockAddress::get(const_cast<Function*>(OldFunc),
+                                              const_cast<BasicBlock*>(BB));
+      VMap[OldBBAddr] = BlockAddress::get(NewFunc, NewBB);
+    }
+
+    OrigBBs.insert(BB);
+    VMap[BB] = NewBB;
+    // Iterate over all possible successors and add them to the CloneWorklist.
+    const TerminatorInst *Term = BB->getTerminator();
+    for (unsigned i = 0, e = Term->getNumSuccessors(); i != e; ++i) {
+      BasicBlock *Succ = Term->getSuccessor(i);
+      CloneWorklist.push_back(Succ);
+    }
   }
-  
+
+  // Now, fill only the reachable blocks with the cloned contents
+  // of the originals.
+  assert(CloneWorklist.empty() && "Dirty worklist before re-use\n");
+  CloneWorklist.push_back(&OldFunc->getEntryBlock());
+  while (!CloneWorklist.empty()) {
+    const BasicBlock *BB = CloneWorklist.back();
+    CloneWorklist.pop_back();
+    PFC.CloneBlock(BB, cast<BasicBlock>(VMap[BB]), CloneWorklist,
+                   OrigBBs);
+  }
+
+  // FIXME: Delete BB's that were created but have been pruned.
+  // Actual cloning may have found pruning opportunities since
+  // branch or switch statement target may have been known at compile-time.
+  // Alternatively we could write a routine CloneFunction and add a) a
+  // parameter to actually do the cloning and b) a return parameter that
+  // gives a list of blocks that need to be cloned also. Then we could
+  // call CloneFunction when we collect the blocks to clone, but suppress
+  // cloning. And then actually *do* the cloning in the while loop above. Then
+  // the cleanup here would become redundant, and so would be the OrigBBs.
+  for (std::set<const BasicBlock *>::iterator Oi = OrigBBs.begin(),
+       Oe = OrigBBs.end(); Oi != Oe; ++Oi) {
+    const BasicBlock *Orig = *Oi;
+    BasicBlock *NewBB = cast<BasicBlock>(VMap[Orig]);
+    delete NewBB;
+    VMap[Orig] = 0;
+  }
+
   // Loop over all of the basic blocks in the old function.  If the block was
   // reachable, we have cloned it and the old block is now in the value map:
   // insert it into the new function in the right order.  If not, ignore it.
@@ -410,59 +523,26 @@ void llvm::CloneAndPruneFunctionInto(Function *NewFunc, const Function *OldFunc,
   SmallVector<const PHINode*, 16> PHIToResolve;
   for (Function::const_iterator BI = OldFunc->begin(), BE = OldFunc->end();
        BI != BE; ++BI) {
-    BasicBlock *NewBB = cast_or_null<BasicBlock>(ValueMap[BI]);
-    if (NewBB == 0) continue;  // Dead block.
+    Value *V = VMap[BI];
+    BasicBlock *NewBB = cast_or_null<BasicBlock>(V);
+    if (!NewBB)
+      continue; // Dead block.
 
     // Add the new block to the new function.
     NewFunc->getBasicBlockList().push_back(NewBB);
-    
-    // Loop over all of the instructions in the block, fixing up operand
-    // references as we go.  This uses ValueMap to do all the hard work.
-    //
-    BasicBlock::iterator I = NewBB->begin();
-
-    LLVMContext &Context = OldFunc->getContext();
-    unsigned DbgKind = Context.getMetadata().getMDKind("dbg");
-    MDNode *TheCallMD = NULL;
-    SmallVector<Value *, 4> MDVs;
-    if (TheCall && TheCall->hasMetadata()) 
-      TheCallMD = Context.getMetadata().getMD(DbgKind, TheCall);
-    
+
     // Handle PHI nodes specially, as we have to remove references to dead
     // blocks.
-    if (PHINode *PN = dyn_cast<PHINode>(I)) {
-      // Skip over all PHI nodes, remembering them for later.
-      BasicBlock::const_iterator OldI = BI->begin();
-      for (; (PN = dyn_cast<PHINode>(I)); ++I, ++OldI) {
-        if (I->hasMetadata())
-          if (TheCallMD) {
-            if (MDNode *IMD = Context.getMetadata().getMD(DbgKind, I)) {
-              MDNode *NewMD = UpdateInlinedAtInfo(IMD, TheCallMD, Context);
-              Context.getMetadata().addMD(DbgKind, NewMD, I);
-            }
-          } else 
-            // The cloned instruction has dbg info but the call instruction
-            // does not have dbg info. Remove dbg info from cloned instruction.
-            Context.getMetadata().removeMD(DbgKind, I);
-        PHIToResolve.push_back(cast<PHINode>(OldI));
-      }
-    }
-    
-    // Otherwise, remap the rest of the instructions normally.
-    for (; I != NewBB->end(); ++I) {
-      if (I->hasMetadata())
-        if (TheCallMD) {
-          if (MDNode *IMD = Context.getMetadata().getMD(DbgKind, I)) {
-            MDNode *NewMD = UpdateInlinedAtInfo(IMD, TheCallMD, Context);
-            Context.getMetadata().addMD(DbgKind, NewMD, I);
-          }
-        } else 
-          // The cloned instruction has dbg info but the call instruction
-          // does not have dbg info. Remove dbg info from cloned instruction.
-          Context.getMetadata().removeMD(DbgKind, I);
-
-      RemapInstruction(I, ValueMap);
-    }
+    for (BasicBlock::const_iterator I = BI->begin(), E = BI->end(); I != E; ++I)
+      if (const PHINode *PN = dyn_cast<PHINode>(I))
+        PHIToResolve.push_back(PN);
+      else
+        break;
+
+    // Finally, remap the terminator instructions, as those can't be remapped
+    // until all BBs are mapped.
+    RemapInstruction(NewBB->getTerminator(), VMap,
+                     ModuleLevelChanges ? RF_None : RF_NoModuleLevelChanges);
   }
   
   // Defer PHI resolution until rest of function is resolved, PHI resolution
@@ -471,19 +551,20 @@ void llvm::CloneAndPruneFunctionInto(Function *NewFunc, const Function *OldFunc,
     const PHINode *OPN = PHIToResolve[phino];
     unsigned NumPreds = OPN->getNumIncomingValues();
     const BasicBlock *OldBB = OPN->getParent();
-    BasicBlock *NewBB = cast<BasicBlock>(ValueMap[OldBB]);
+    BasicBlock *NewBB = cast<BasicBlock>(VMap[OldBB]);
 
     // Map operands for blocks that are live and remove operands for blocks
     // that are dead.
     for (; phino != PHIToResolve.size() &&
          PHIToResolve[phino]->getParent() == OldBB; ++phino) {
       OPN = PHIToResolve[phino];
-      PHINode *PN = cast<PHINode>(ValueMap[OPN]);
+      PHINode *PN = cast<PHINode>(VMap[OPN]);
       for (unsigned pred = 0, e = NumPreds; pred != e; ++pred) {
-        if (BasicBlock *MappedBlock = 
-            cast_or_null<BasicBlock>(ValueMap[PN->getIncomingBlock(pred)])) {
+        Value *V = VMap[PN->getIncomingBlock(pred)];
+        if (BasicBlock *MappedBlock = cast_or_null<BasicBlock>(V)) {
           Value *InVal = MapValue(PN->getIncomingValue(pred),
-                                  ValueMap);
+                                  VMap, 
+                        ModuleLevelChanges ? RF_None : RF_NoModuleLevelChanges);
           assert(InVal && "Unknown input value?");
           PN->setIncomingValue(pred, InVal);
           PN->setIncomingBlock(pred, MappedBlock);
@@ -537,50 +618,83 @@ void llvm::CloneAndPruneFunctionInto(Function *NewFunc, const Function *OldFunc,
       while ((PN = dyn_cast<PHINode>(I++))) {
         Value *NV = UndefValue::get(PN->getType());
         PN->replaceAllUsesWith(NV);
-        assert(ValueMap[OldI] == PN && "ValueMap mismatch");
-        ValueMap[OldI] = NV;
+        assert(VMap[OldI] == PN && "VMap mismatch");
+        VMap[OldI] = NV;
         PN->eraseFromParent();
         ++OldI;
       }
     }
-    // NOTE: We cannot eliminate single entry phi nodes here, because of
-    // ValueMap.  Single entry phi nodes can have multiple ValueMap entries
-    // pointing at them.  Thus, deleting one would require scanning the ValueMap
-    // to update any entries in it that would require that.  This would be
-    // really slow.
   }
-  
+
+  // Make a second pass over the PHINodes now that all of them have been
+  // remapped into the new function, simplifying the PHINode and performing any
+  // recursive simplifications exposed. This will transparently update the
+  // WeakVH in the VMap. Notably, we rely on that so that if we coalesce
+  // two PHINodes, the iteration over the old PHIs remains valid, and the
+  // mapping will just map us to the new node (which may not even be a PHI
+  // node).
+  for (unsigned Idx = 0, Size = PHIToResolve.size(); Idx != Size; ++Idx)
+    if (PHINode *PN = dyn_cast<PHINode>(VMap[PHIToResolve[Idx]]))
+      recursivelySimplifyInstruction(PN, DL);
+
   // Now that the inlined function body has been fully constructed, go through
   // and zap unconditional fall-through branches.  This happen all the time when
   // specializing code: code specialization turns conditional branches into
   // uncond branches, and this code folds them.
-  Function::iterator I = cast<BasicBlock>(ValueMap[&OldFunc->getEntryBlock()]);
+  Function::iterator Begin = cast<BasicBlock>(VMap[&OldFunc->getEntryBlock()]);
+  Function::iterator I = Begin;
   while (I != NewFunc->end()) {
+    // Check if this block has become dead during inlining or other
+    // simplifications. Note that the first block will appear dead, as it has
+    // not yet been wired up properly.
+    if (I != Begin && (pred_begin(I) == pred_end(I) ||
+                       I->getSinglePredecessor() == I)) {
+      BasicBlock *DeadBB = I++;
+      DeleteDeadBlock(DeadBB);
+      continue;
+    }
+
+    // We need to simplify conditional branches and switches with a constant
+    // operand. We try to prune these out when cloning, but if the
+    // simplification required looking through PHI nodes, those are only
+    // available after forming the full basic block. That may leave some here,
+    // and we still want to prune the dead code as early as possible.
+    ConstantFoldTerminator(I);
+
     BranchInst *BI = dyn_cast<BranchInst>(I->getTerminator());
     if (!BI || BI->isConditional()) { ++I; continue; }
     
-    // Note that we can't eliminate uncond branches if the destination has
-    // single-entry PHI nodes.  Eliminating the single-entry phi nodes would
-    // require scanning the ValueMap to update any entries that point to the phi
-    // node.
     BasicBlock *Dest = BI->getSuccessor(0);
-    if (!Dest->getSinglePredecessor() || isa<PHINode>(Dest->begin())) {
+    if (!Dest->getSinglePredecessor()) {
       ++I; continue;
     }
-    
+
+    // We shouldn't be able to get single-entry PHI nodes here, as instsimplify
+    // above should have zapped all of them..
+    assert(!isa<PHINode>(Dest->begin()));
+
     // We know all single-entry PHI nodes in the inlined function have been
     // removed, so we just need to splice the blocks.
     BI->eraseFromParent();
     
-    // Move all the instructions in the succ to the pred.
-    I->getInstList().splice(I->end(), Dest->getInstList());
-    
     // Make all PHI nodes that referred to Dest now refer to I as their source.
     Dest->replaceAllUsesWith(I);
 
+    // Move all the instructions in the succ to the pred.
+    I->getInstList().splice(I->end(), Dest->getInstList());
+    
     // Remove the dest block.
     Dest->eraseFromParent();
     
     // Do not increment I, iteratively merge all things this block branches to.
   }
+
+  // Make a final pass over the basic blocks from theh old function to gather
+  // any return instructions which survived folding. We have to do this here
+  // because we can iteratively remove and merge returns above.
+  for (Function::iterator I = cast<BasicBlock>(VMap[&OldFunc->getEntryBlock()]),
+                          E = NewFunc->end();
+       I != E; ++I)
+    if (ReturnInst *RI = dyn_cast<ReturnInst>(I->getTerminator()))
+      Returns.push_back(RI);
 }