#ifndef LLVM_ANALYSIS_LOOPCHAINDETECT_H
#define LLVM_ANALYSIS_LOOPCHAINDETECT_H

#include "llvm/ADT/DirectedGraph.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/IntrinsicInst.h"
#include "llvm/IR/PassManager.h"

namespace llvm {

/// \return string containing a file name and a line # for the given loop.
static std::string getDebugLocString(const Loop *L) {
  std::string Result;
  if (L) {
    raw_string_ostream OS(Result);
    if (const DebugLoc LoopDbgLoc = L->getStartLoc())
      LoopDbgLoc.print(OS);
    else // Just print the module name.
      OS << L->getHeader()->getParent()->getParent()->getModuleIdentifier();
    OS.flush();
  }
  return Result;
}

static bool isIntrinsicCall(CallBase *CI) {
  // Catch llvm intrinsics
  if (isa<IntrinsicInst>(CI)) {
    return true;
  }
  // Catch flang intrinsics
  StringRef OpName = CI->getCalledOperand()->getName();
  if (OpName.starts_with("f90") || OpName.contains("__mth_")) {
    return true;
  }
  return false;
}

using CandidateTy = std::pair<BasicBlock *, unsigned>;
class LoopNode;
class LoopEdge;

class LoopNode : public DGNode<LoopNode, LoopEdge> {

public:
  LoopNode(Loop &LoopRep) : LoopRep(LoopRep) {}
  Loop &getLoop() const { return LoopRep; }
  std::string getName() const;

protected:
  Loop &LoopRep;
};

class LoopEdge : public DGEdge<LoopNode, LoopEdge> {
public:
  LoopEdge(LoopNode &N) : DGEdge(N) {}
};

template <typename NodeType> class LoopGraphInfo {
public:
  LoopGraphInfo() = delete;
  LoopGraphInfo(const LoopGraphInfo &G) = delete;
  LoopGraphInfo(const std::string Name) : Name(Name) {}
  virtual ~LoopGraphInfo() = default;

  StringRef getName() const { return Name; }

protected:
  std::string Name;
};

class LoopGraph final : public DirectedGraph<LoopNode, LoopEdge>,
                        public LoopGraphInfo<LoopNode> {

public:
  LoopGraph(const LoopGraph &) = delete;
  LoopGraph(Function &F, FunctionAnalysisManager *FAM)
      : LoopGraphInfo(F.getName().str()), F(F), FAM(FAM) {
    MaxNodes = 0;
  }
  LoopGraph(Function &F) : LoopGraph(F, nullptr) {}
  ~LoopGraph() {}

  SmallVector<LoopNode *> buildLoopGraph(Function *F);
  SmallVector<LoopNode *> buildLoopGraph(Loop *CurrLoop);
  SmallVector<LoopNode *> buildLoopGraph(CallBase *CI);
  std::pair<LoopNode *, unsigned> getNextLoopChainCandidate();
  SmallVector<CandidateTy> getLoopChainCandidates();
  unsigned NodeSize() {return MaxNodes; };
  void removeNodeAndChildren(LoopNode *Node);

private:
  LoopInfo *getLoopInfoForFunction(Function *F);
  DominatorTree *getDominatorTreeForFunction(Function *F);

  Function &F;
  FunctionAnalysisManager *FAM;
  DenseMap<Function *, LoopInfo *> LIs;
  DenseMap<Function *, DominatorTree *> DTs;
  SmallVector<Function *> CallStack;
  std::unordered_map<Function *, SmallVector<LoopNode *>> FunctionCache;
  unsigned MaxNodes;
};

raw_ostream &operator<<(raw_ostream &OS, const LoopNode &N);
raw_ostream &operator<<(raw_ostream &OS, const LoopEdge &E);
raw_ostream &operator<<(raw_ostream &OS, const LoopGraph &G);

class LoopChainCandidates {
public:
  LoopChainCandidates() {}

  void addCandidate(BasicBlock *Candidate, unsigned length) {
    Candidates.push_back(std::make_pair(Candidate, length));
  }

  ArrayRef<CandidateTy> getCandidates() const {
    return Candidates;
  }

private:
  SmallVector<CandidateTy> Candidates;
};

class LoopChainDetectAnalysis
    : public AnalysisInfoMixin<LoopChainDetectAnalysis> {
public:
  typedef LoopChainCandidates Result;
  LoopChainCandidates run(Function &F, FunctionAnalysisManager &AM);

private:
  friend AnalysisInfoMixin<LoopChainDetectAnalysis>;
  static AnalysisKey Key;
};

class LoopChainDetectPrinterPass
    : public PassInfoMixin<LoopChainDetectPrinterPass> {
  raw_ostream &OS;

public:
  explicit LoopChainDetectPrinterPass(raw_ostream &OS = llvm::dbgs())
      : OS(OS) {}
  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
  bool runOnFunction(Function &F, FunctionAnalysisManager &AM);
};

} // namespace llvm

#endif // LLVM_ANALYSIS_LOOPCHAINDETECT_H