#ifndef LLVM_TRANSFORMS_UTILS_GETPOSSIBLECALLEESUTILS_H
#define LLVM_TRANSFORMS_UTILS_GETPOSSIBLECALLEESUTILS_H
#include "llvm/Analysis/CallSiteAnalysis.h"
#include "llvm/Analysis/IndirectToDirectCallAnalysis.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/Module.h"
#include "llvm/Pass.h"
#include "llvm/Passes/PassBuilder.h"
#include <map>
namespace llvm {
class CallPossiblilties {
  aocc::aocc::CallSiteAnalysis::Result *CSA = nullptr;
  aocc::IndirectToDirectCallAnalysisInfo::Result *IDCA = nullptr;
  std::map<CallBase *, std::set<Function *>> callInstToFunctionMap;
  std::map<Function *, std::set<CallBase *>> functionToCallBaseMap;
  /* Consult Indirect Call possiblity analysis and
     Callsite Analysis to get possible functions for a call
     instruction.
  */
  std::set<Function *> populatePossiblities(CallBase *D) {
    std::set<Function *> CallTargets;

    if (isIndirectCall(D) && CSA) {
      const auto &FS = CSA->getPossibleImpl(D);
      if (FS.has_value())
        CallTargets = FS.value();
    }
    if (CallTargets.empty() && IDCA && isa<CallInst>(D)) {
      auto *IDCR = IDCA->findIndirectCallPossiblities(cast<CallInst>(D));
      if (IDCR)
        CallTargets.insert(IDCR->begin(), IDCR->end());
    }
    return CallTargets;
  }

  bool isIndirectCall(CallBase *D) {
    IntrinsicInst *II = dyn_cast<IntrinsicInst>(D);
    if (II && II->getIntrinsicID() == Intrinsic::type_test)
      return true;

    return !D->getCalledFunction();
  }

public:
  CallPossiblilties(aocc::CallSiteAnalysis::Result &CSA,
                    IndirectToDirectCallAnalysisInfo::Result &IDCA)
      : CSA(&CSA), IDCA(&IDCA) {}
  CallPossiblilties() : CSA(nullptr), IDCA(nullptr) {}
  /* Create Maps from function to indirect call sites and vice versa */
  void run(Module &M, aocc::CallSiteAnalysis::Result &CallSiteA,
           IndirectToDirectCallAnalysisInfo::Result &IndirectDCA) {
    CSA = &CallSiteA;
    IDCA = &IndirectDCA;
     
    for (Function &F : M) {
      for (BasicBlock &BB : F) {
        for (Instruction &I : BB) {
          if (CallBase *callInst = dyn_cast<CallBase>(&I)) {
            if (callInstToFunctionMap.find(callInst) ==
                callInstToFunctionMap.end())
              callInstToFunctionMap[callInst] = std::set<Function *>();
            std::set<Function *> result = populatePossiblities(callInst);
            callInstToFunctionMap[callInst].insert(result.begin(),
                                                   result.end());
            for (Function *F : result) {
              if (F == nullptr)
                continue;
              if (functionToCallBaseMap.find(F) == functionToCallBaseMap.end())
                functionToCallBaseMap[F] = std::set<CallBase *>();
              functionToCallBaseMap[F].insert(callInst);
            }
          }
        }
      }
    }
    return;
  }

  std::map<CallBase *, std::set<Function *>> getCallBaseToFunctionMap() {
    return callInstToFunctionMap;
  }
  std::map<Function *, std::set<CallBase *>> getFunctionToCallBaseMap() {
    return functionToCallBaseMap;
  }
};
} // namespace llvm
#endif
