#ifndef LLVM_ANALYSIS_POINTER_CACHE_MISS_ANALYSIS_H
#define LLVM_ANALYSIS_POINTER_CACHE_MISS_ANALYSIS_H

#include "llvm/Analysis/ScalarEvolutionExpressions.h"
#include "llvm/Analysis/TBAATypeReconstruction.h"
#include "llvm/Analysis/LoopNestAnalysis.h"
#include "llvm/Analysis/ScalarEvolution.h"
#include "llvm/Analysis/PostDominators.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/IR/InstIterator.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/PassManager.h"
#include "llvm/IR/Dominators.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/IR/Value.h"
#include "llvm/Pass.h"
#define DEBUG_TYPE "pointer-cache-miss-analysis"

namespace llvm {

class PointerCacheMiss;

class PointerCacheMiss {
  public:
    struct callTreeNodes_ {
        Instruction *I;
        const SCEV *argumentSCEV;
        SmallVector<std::shared_ptr<callTreeNodes_>> children;
    };
    typedef struct callTreeNodes_ callTreeNodes;

  private:
    FunctionAnalysisManager &FAM;
    TBAATypeReconstructionAnalysis::Result &tbaaTR;

    Module &M;

    SmallVector<CallBase*> candidateMallocs;
    SmallVector<CallBase*> mallocCallSites;
    SmallVector<LoadInst*> pointerLoad;
    // SmallVector<std::shared_ptr<PointerCacheMiss::callTreeNodes>> callTreeRoots;
    SmallSetVector<Instruction*, 4> candidateInstructions;
    SmallDenseMap<Instruction *, int> candidateAllocationSizes;
    SmallVector<SmallVector<Instruction*>> candidateLoadChains;
    SmallSet<SmallVector<Instruction*>, 4> candidateLoadChainsSet;
    SmallDenseMap<Function*, SmallVector<SmallVector<Instruction*>>> candidateLoadChainsMap;
    SmallVector<SmallVector<Instruction*>> candidateLoadChainsFromGuardBlocks;

    // This function records all pointer loads in the given Loop
    void recordPointerLoads(const Loop* L);
    // This function records all the malloc call sites in the Function
    void recordMallocCallSites(Function &F);
    // This function finds all the candidate loop nests in the given function
    void findCandidateLoopNests(Function &F, const LoopInfo &LI, ScalarEvolution &SE);
    // This function finds the candidate instructions by using TBAA information and
    // matching it to allocation sites.
    void findCandidateInstructions(SmallVector<std::shared_ptr<PointerCacheMiss::callTreeNodes>> &leafNodes, std::map<std::shared_ptr<PointerCacheMiss::callTreeNodes>, APInt> &finalLeafSCEVValues);
    // This function finds the control conditions that can be cached as well.
    // Derived from candidate instructions.
    void findCandidateControlConditions();
    // This function shortlist the malloc calls of our interest.
    // The ones that can resolve to a constant value statically.
    void findCandidateMallocs();

    void performChainAnalysis();
    void performChainAnalysisImpl(Instruction *I, SmallVector<Instruction*> &loadChainImpl, int loopIVUseCountThreshold = 2);

    std::shared_ptr<PointerCacheMiss::callTreeNodes> buildMallocCallTree(CallBase *root);
    void processMallocCallTree(std::shared_ptr<PointerCacheMiss::callTreeNodes> root, SmallVector<std::shared_ptr<PointerCacheMiss::callTreeNodes>> &leafNodes, std::map<std::shared_ptr<PointerCacheMiss::callTreeNodes>, APInt> &finalLeafSCEVValues);

  public:
    PointerCacheMiss(Module &M, FunctionAnalysisManager &FAM, TBAATypeReconstructionAnalysis::Result &TR) : M(M), FAM(FAM), tbaaTR(TR) {}
    // Returns a small vector of load chains that are candidates based
    // on the algorithm.
    SmallVector<SmallVector<Instruction*>> getCandidateLoadChains() {return candidateLoadChains;}
    // Get the allocation size associated with the given instruction.
    int getAllocationSize(Instruction *I) {
      return candidateAllocationSizes.contains(I) ? candidateAllocationSizes[I]
                                                  : -1;
    }
    // Get the load chains as getCandidateLoadChains, but filtered
    // by their parent functions.
    SmallDenseMap<Function*, SmallVector<SmallVector<Instruction*>>> getCandidateLoadChainsMap() {return candidateLoadChainsMap;}

    void analyze();
    bool invalidate(Module &M, const PreservedAnalyses &PA, ModuleAnalysisManager::Invalidator &Inv);
};

class PointerCacheMissAnalysis
    : public AnalysisInfoMixin<PointerCacheMissAnalysis> {
  friend AnalysisInfoMixin<PointerCacheMissAnalysis>;
  static AnalysisKey Key;

public:
  typedef PointerCacheMiss Result;
  Result run(Module &M, ModuleAnalysisManager &AM);
};

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

public:
  explicit PointerCacheMissAnalysisPrinterPass(raw_ostream &OS) : OS(OS) {}
  explicit PointerCacheMissAnalysisPrinterPass() : OS(dbgs()) {}
  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
};
}
#endif
