1 files changed, 22 insertions, 10 deletions

diff --git a/lib/CodeGen/BranchFolding.h b/lib/CodeGen/BranchFolding.h
index d759d53e27f2f..36a5a2e2c97c6 100644
--- a/lib/CodeGen/BranchFolding.h
+++ b/lib/CodeGen/BranchFolding.h
@@ -11,6 +11,7 @@
 #define LLVM_LIB_CODEGEN_BRANCHFOLDING_H
 
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/CodeGen/LivePhysRegs.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/Support/BlockFrequency.h"
 #include <vector>
@@ -20,20 +21,23 @@ namespace llvm {
   class MachineBranchProbabilityInfo;
   class MachineFunction;
   class MachineModuleInfo;
-  class RegScavenger;
+  class MachineLoopInfo;
   class TargetInstrInfo;
   class TargetRegisterInfo;
 
   class LLVM_LIBRARY_VISIBILITY BranchFolder {
   public:
+    class MBFIWrapper;
+
     explicit BranchFolder(bool defaultEnableTailMerge, bool CommonHoist,
-                          const MachineBlockFrequencyInfo &MBFI,
+                          MBFIWrapper &MBFI,
                           const MachineBranchProbabilityInfo &MBPI);
 
-    bool OptimizeFunction(MachineFunction &MF,
-                          const TargetInstrInfo *tii,
-                          const TargetRegisterInfo *tri,
-                          MachineModuleInfo *mmi);
+    bool OptimizeFunction(MachineFunction &MF, const TargetInstrInfo *tii,
+                          const TargetRegisterInfo *tri, MachineModuleInfo *mmi,
+                          MachineLoopInfo *mli = nullptr,
+                          bool AfterPlacement = false);
+
   private:
     class MergePotentialsElt {
       unsigned Hash;
@@ -91,13 +95,17 @@ namespace llvm {
     };
     std::vector<SameTailElt> SameTails;
 
+    bool AfterBlockPlacement;
     bool EnableTailMerge;
     bool EnableHoistCommonCode;
+    bool UpdateLiveIns;
     const TargetInstrInfo *TII;
     const TargetRegisterInfo *TRI;
     MachineModuleInfo *MMI;
-    RegScavenger *RS;
+    MachineLoopInfo *MLI;
+    LivePhysRegs LiveRegs;
 
+  public:
     /// \brief This class keeps track of branch frequencies of newly created
     /// blocks and tail-merged blocks.
     class MBFIWrapper {
@@ -105,21 +113,25 @@ namespace llvm {
       MBFIWrapper(const MachineBlockFrequencyInfo &I) : MBFI(I) {}
       BlockFrequency getBlockFreq(const MachineBasicBlock *MBB) const;
       void setBlockFreq(const MachineBasicBlock *MBB, BlockFrequency F);
+      raw_ostream &printBlockFreq(raw_ostream &OS,
+                                  const MachineBasicBlock *MBB) const;
+      raw_ostream &printBlockFreq(raw_ostream &OS,
+                                  const BlockFrequency Freq) const;
 
     private:
       const MachineBlockFrequencyInfo &MBFI;
       DenseMap<const MachineBasicBlock *, BlockFrequency> MergedBBFreq;
     };
 
-    MBFIWrapper MBBFreqInfo;
+  private:
+    MBFIWrapper &MBBFreqInfo;
     const MachineBranchProbabilityInfo &MBPI;
 
     bool TailMergeBlocks(MachineFunction &MF);
     bool TryTailMergeBlocks(MachineBasicBlock* SuccBB,
                        MachineBasicBlock* PredBB);
     void setCommonTailEdgeWeights(MachineBasicBlock &TailMBB);
-    void MaintainLiveIns(MachineBasicBlock *CurMBB,
-                         MachineBasicBlock *NewMBB);
+    void computeLiveIns(MachineBasicBlock &MBB);
     void ReplaceTailWithBranchTo(MachineBasicBlock::iterator OldInst,
                                  MachineBasicBlock *NewDest);
     MachineBasicBlock *SplitMBBAt(MachineBasicBlock &CurMBB,