#ifndef LLVM_TRANSFORMS_LOOPCHAINOUTLINE_H
#define LLVM_TRANSFORMS_LOOPCHAINOUTLINE_H

#include "llvm/ADT/SmallSet.h"
#include "llvm/Analysis/DomTreeUpdater.h"
#include "llvm/Analysis/LoopChainDetect.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/Analysis/PostDominators.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/InstVisitor.h"
#include "llvm/IR/PassManager.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "llvm/Transforms/Utils/Cloning.h"
#include "llvm/Transforms/Utils/CodeExtractor.h"
#include "llvm/Transforms/Utils/CodeMoverUtils.h"
#include "llvm/Transforms/Utils/ValueMapper.h"

namespace llvm {

class LoopChainOutlineOptions {

public:
  unsigned MinLCSize;
  unsigned MaxLCSize;
  bool ForceOutlineInMain = false;
  std::string TargetFunction;

  LoopChainOutlineOptions() {
    MinLCSize = 1;
    MaxLCSize = 1000;
    TargetFunction = "MAIN_";
  }

  LoopChainOutlineOptions(const std::string &TF, unsigned MinSize,
                          unsigned MaxSize, bool ForceOutInMain = false) : LoopChainOutlineOptions() {
    setTargetFunction(TF);
    setMinLCSize(MinSize);
    setMaxLCSize(MaxSize);
    ForceOutlineInMain = ForceOutInMain;
  }

  LoopChainOutlineOptions &setMinLCSize(unsigned Sz) {
    if (Sz > 0)
      MinLCSize = Sz;
    if (MinLCSize > MaxLCSize)
      MaxLCSize = MinLCSize + 1;
    return *this;
  }

  LoopChainOutlineOptions &setMaxLCSize(unsigned Sz) {
    if (Sz > 0)
      MaxLCSize = Sz;
    if (MaxLCSize < MinLCSize)
      MinLCSize = MaxLCSize - 1;
    return *this;
  }

  LoopChainOutlineOptions &setTargetFunction(const std::string &TF) {
    if (!TF.empty())
      TargetFunction = TF;
    return *this;
  }
};

class FunctionCallUpdater : public InstVisitor<FunctionCallUpdater> {

public:
  Module *M;
  SmallVector<CallInst *> UpdatedCalls;
  SmallSet<CallInst *, 32> Visited;

  FunctionCallUpdater(Module *m) : M(m){};

  void visitCallInst(CallInst &CI);
  static void replaceCalleeWithThis(Function *ReplaceWith);
  static bool guardFunctionCallingOutlined(Function *OrigF, Function *ClonedF,
                                           DominatorTree &DT,
                                           PostDominatorTree &PDT,
                                           LoopInfo &LI,
                                           Function *FOutlinedCloned);
  static void replaceAllFunctionCalls(Function *Func,
                                      SmallSet<CallInst *, 32> &Visited);
  static MDNode* setMetadata(Instruction *I, const std::string& Data,
                          const std::string& Name);
};

class LoopChainOutlinePass : public PassInfoMixin<LoopChainOutlinePass> {

public:
  LoopChainOutlinePass(LoopChainOutlineOptions Opts = {})
      : Opts(Opts), OutlinedCount(0) {}
  LoopChainOutlineOptions Opts;
  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);

protected:
  void outlineLoopChainCandidates(LoopChainCandidates Candidates);
  Function *outlineLoopBody(Loop *LCC, DominatorTree &DT,
                            PostDominatorTree &PDT, LoopInfo &LI);
  Function *guardOutlinedFunction(Function *F_outlined, DominatorTree &DT,
                             PostDominatorTree &PDT, LoopInfo &LI);
  bool guardFunctionCallingOutlined(Function *OrigF, Function *ClonedF,
                              DominatorTree &DT,
                              PostDominatorTree &PDT,
                              LoopInfo &LI,
                              Function *FOutlinedCloned);
  bool isMain(Function * F);
  Function *cloneOutlined(Function *FOutlined);

private:
  SmallVector<BasicBlock *> isolateLoopBody(Loop *Target, DominatorTree *DT,
                                            PostDominatorTree *PDT,
                                            LoopInfo *LI);
  bool copyAllInputInstructions(Loop *Target, BranchInst *Br,
                                ValueToValueMapTy &VMap, bool DryRun = false);
  Function* getOrigFunc(Function *F);
  bool verifyFunctionSignatures(Function *Original, Function *Cloned);
  void emitOptimizationRemark(Function *Func, BasicBlock *BB,
                              const std::string& Status,
                              const std::string& Reason = "");

  FunctionAnalysisManager *FAM;
  DenseMap<Function *, PostDominatorTree *> PDTs;
  unsigned OutlinedCount;
  bool ForceOutlineInMain = false;
};

class LoopChainOutlineCleanupPass
    : public PassInfoMixin<LoopChainOutlineCleanupPass> {
  bool executeClonedOutlined = false;
  class LoopChainOutlineCleanupPassVisitor
      : public InstVisitor<LoopChainOutlineCleanupPassVisitor> {
  bool executeClonedOutlined = false;
  public:
    LoopChainOutlineCleanupPassVisitor(bool executeClonedOutlined = false) 
      : executeClonedOutlined(executeClonedOutlined) {}  
    void visitCallInst(CallInst &CI);
  };

public:
  LoopChainOutlineCleanupPass(bool executeClonedOutlined = false) : executeClonedOutlined(executeClonedOutlined) {}
  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
  bool runOnFunction(Function &F, FunctionAnalysisManager &AM);
};
} // namespace llvm
#endif // LLVM_TRANSFORMS_LOOPCHAINOUTLINE_H