#ifndef LLVM_TRANSFORMS_UTILS_GCP_H
#define LLVM_TRANSFORMS_UTILS_GCP_H

#include "llvm/Analysis/LoopInfo.h"
#include "llvm/Analysis/ScalarEvolution.h"
#include "llvm/Analysis/ScalarEvolutionExpressions.h"
#include "llvm/Analysis/TypeReconstruction.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/DerivedTypes.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/Module.h"
#include "llvm/Pass.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
#include <set>
#include <queue>
#include <vector>
#include <stack>
#include <map>
using namespace llvm;
using namespace std;

#define INS_VISIT_CNT                                                          \
  5                      // while propagating constants, dont visit any ins
                         // greater than these many times
// this just indicates max values a variable can take during analysis and not
// after analysis. For post analysis value,check 'maxVarValueSetSz' in
// functionSpecialize.cpp
#define MAX_VALUES_LLVM_VAR 20

namespace llvm {
// store the vals for attributes at the callsite in the form of table.
// It has following format:
//---------------------------------------------------------------------------
// Caller func | CallSite Inst | Arg No | Set of Const values |
//---------------------------------------------------------------------------
class CallSiteArgVals {
public:
  class Caller {
  public:
    Function *F;
    Caller(Function *F) : F(F) {}
    class CallSite {
    public:
      Instruction *CS;
      CallSite(Instruction *CS) : CS(CS) {}
      class ArgVal {
      public:
        unsigned pos; // positition of the argument into the callsite
        Value *arg;
        ArgVal(unsigned &pos, Value *arg, vector<Constant *> &vec)
            : pos(pos), arg(arg) {
          this->V = std::move(vec);
        }
        vector<Constant *> V;
        unsigned getKey() const { return pos; }
        void createUniqueAttrVals();
      };
      struct cmp {
        using is_transparent = void;
        bool operator()(ArgVal *const A, ArgVal *const B) const {
          return A->getKey() < B->getKey();
        }
        bool operator()(ArgVal *const A, const unsigned &pos) const {
          return A->getKey() < pos;
        }
        bool operator()(const unsigned &pos, ArgVal *const A) const {
          return pos < A->getKey();
        }
      };
      Instruction *getKey() const { return CS; }
      set<ArgVal *, cmp> V;
      ArgVal *getArgValPtr(unsigned &pos);
      void getCombinations(vector<vector<std::pair<unsigned, Constant *>>> &,
                           unsigned, vector<unsigned> &,
                           stack<std::pair<unsigned, Constant *>> &);
      void getAttrIndices(vector<unsigned> &);
      ~CallSite() {
        for (auto v: V) {
          delete v;
        }
      }
    };
    struct cmp {
      using is_transparent = void;
      bool operator()(CallSite *const A, CallSite *const B) const {
        return A->getKey() < B->getKey();
      }
      bool operator()(CallSite *const A, Instruction *const I) const {
        return A->getKey() < I;
      }
      bool operator()(Instruction *const I, CallSite *const A) const {
        return I < A->getKey();
      }
    };
    Function *getKey() const { return F; }
    set<CallSite *, cmp> V;
    CallSite *getCallSitePtr(Instruction *I);
    void delCallSitePtr(Instruction *);
    ~Caller() {
      for (auto v: V) {
        delete v;
      }
    }
  };
  struct cmp {
    using is_transparent = void;
    bool operator()(Caller *const A, Caller *const B) const {
      return A->getKey() < B->getKey();
    }
    bool operator()(Caller *const A, Function *const F) const {
      return A->getKey() < F;
    }
    bool operator()(Function *const F, Caller *const A) const {
      return F < A->getKey();
    }
  };
  set<Caller *, cmp> V; // called func
  Caller *getCallerPtr(Function *F);
  ~CallSiteArgVals() {
    for (auto v: V) {
      delete v;
    }
  }
};

// Base class for holding a const var and a pointer var
class Var {
public:
  virtual set<Constant *> &getValSet() { return valSet; }
  virtual set<shared_ptr<Var>> &getPointeeSet() { return pointeeSet; }
  virtual Constant *getAnyVal() {
    return (valSet.empty() ? nullptr : (*(valSet.begin())));
  }
  virtual ~Var(){};

private:
  set<shared_ptr<Var>> pointeeSet;
  set<Constant *> valSet;
};

class PtrVar : public Var {
public:
  PtrVar(const shared_ptr<Var> &C) { V.insert(C); }
  set<shared_ptr<Var>> &getPointeeSet() { return V; }
  ~PtrVar() {
    // automatically decrements the counts of shared_ptr's inside it.
    V.clear();
  }

private:
  set<shared_ptr<Var>> V;
};

class ConstVar : public Var {
public:
  ConstVar(Constant *C) { V.insert(C); }
  ConstVar(vector<Constant *> &C) {
    for (Constant *c : C)
      V.insert(c);
  }
  set<Constant *> &getValSet() { return V; }
  Constant *getAnyVal() { return (V.empty() ? nullptr : (*(V.begin()))); }
  ~ConstVar() { V.clear(); }

private:
  set<Constant *> V;
};

class GlobalConstantSet {
public:
  class GlobalConstant {
  public:
    Value *llvmVal;
    vector<shared_ptr<Var>> V;
    Value *getKey() const { return llvmVal; }
    GlobalConstant(Value *llvmVal) : llvmVal(llvmVal){};
    ~GlobalConstant() { V.clear(); }
  };
  struct cmp {
    using is_transparent = void;
    bool operator()(GlobalConstant *const A, GlobalConstant *const B) const {
      return A->getKey() < B->getKey();
    }
    bool operator()(GlobalConstant *const A, Value *const V) const {
      return A->getKey() < V;
    }
    bool operator()(Value *const V, GlobalConstant *const A) const {
      return V < A->getKey();
    }
  };
  set<GlobalConstant *, cmp> V;
  GlobalConstant *getGCPtr(Value *, bool isDeleteQuery = false);
  void deleteGV(GlobalConstant *);
  ~GlobalConstantSet() {
    for (auto v: V) {
      delete v;
    }
  }
};
// corresponding to every function, find the set of values for every
// possible variable/operand
class ValsForModule {
public:
  Module *M;
  Module *getKey() { return M; }
  ValsForModule(Module *M) : M(M) {}
  class ValsForFunc {
  public:
    Function *F;

    class ValsForIns {
    public:
      Instruction *I;
      ValsForIns(Instruction *I) : I(I) {}
      vector<shared_ptr<Var>> V;
      Instruction *getKey() const { return I; }
    };
    struct cmpI {
      using is_transparent = void;
      bool operator()(ValsForIns *const A, ValsForIns *const B) const {
        return A->getKey() < B->getKey();
      }
      bool operator()(ValsForIns *const A, Instruction *const I) const {
        return A->getKey() < I;
      }
      bool operator()(Instruction *const I, ValsForIns *const A) const {
        return I < A->getKey();
      }
    };
    set<ValsForIns *, cmpI> V;
    ValsForFunc(Function *F) : F(F) {}
    Function *getKey() const { return F; }
    ValsForIns *getIPtr(Instruction *);
    ~ValsForFunc() {
      for (auto v: V) {
        delete v;
      }
    }
  };
  struct cmpF {
    using is_transparent = void;
    bool operator()(ValsForFunc *const A, ValsForFunc *const B) const {
      return A->getKey() < B->getKey();
    }
    bool operator()(ValsForFunc *const A, Function *const F) const {
      return A->getKey() < F;
    }
    bool operator()(Function *const F, ValsForFunc *const A) const {
      return F < A->getKey();
    }
  };
  set<ValsForFunc *, cmpF> V;
  ~ValsForModule() {
    for (auto v: V) {
      delete v;
    }
  }
  ValsForFunc *getFPtr(Function *);
};
struct GeneralizedConstPropagation {
public:
  CallSiteArgVals *CSAV;
  ValsForModule *VFM;
  GlobalConstantSet *GCS;
  TypeReconstructionAnalysis::Result& TR;
  GeneralizedConstPropagation(TypeReconstructionAnalysis::Result& TR) :TR(TR) { }

  // data structure to check a parent for global variable only in following
  // situations:
  // int a[] = {1,2,3,4,5}
  // int* arr[] = {a+1,a+2}
  // Parent for arr is 'a' because it refers 'a'
  map<Value *, set<Value *>> parent;

  bool isAllowedAggregateContainedTy(Type *) const;
  void getGlobalList(Module &);
  void flattenArray(Constant *arrayInitializer, vector<Constant *> &V);
  void flattenStruct(Constant *structConst, vector<Constant *> &V);
  void recordSructArrInitializer(GlobalVariable *);
  inline bool isAllowedArrEleTy(Type *);
  void recordArrInitializer(GlobalVariable *);
  void experiment(Module &M, ScalarEvolution *SE);
  // for globals with const vals
  void makeEntryIntoGlobalValMap(GlobalVariable *, vector<Constant *> &);
  bool insRefersGV(Instruction *);
  void initialize();
  void propagateConstants(Function *, LoopInfo *LI, ScalarEvolution *);
  bool isAllowedAggregateTy(Type *) const;
  Type *getPrimitiveContainedTy(Type *);
  inline bool isAllowedStructEleTy(Type *);
  void makeGlobalModAnalysis(Module &);
  bool operatorPossiblyModifiesOpnd(Value *);
  bool insPossiblyModifiesOpnd(Instruction *, Instruction *);
  bool insPossiblyModifiesOpnd(Value *, Instruction *);
  void collectGVUses(Module *,
                     map<Function *, map<Value *, queue<Instruction *>>> &);
  void collectGEPS(Function *, map<Value *, queue<Instruction *>> &);
  bool inspectInstForAllOpnds(Instruction *);
  void getCombinations(set<Constant *> &, set<Constant *> &,
                       vector<std::pair<Constant *, Constant *>> &);
  bool isPtrToPtr(Value *);
  void getArrBounds(Type *, vector<unsigned> &);
  void getCombinations(map<unsigned, vector<Var *>> &, int,
                       vector<vector<Var *>> &, stack<Var *> &);
  bool processIns(Instruction *, ScalarEvolution *);
  bool processIntArithmeticIns(Instruction *, unsigned, ScalarEvolution *);
  bool processShlIns(Instruction *, ScalarEvolution *);
  bool processGEPIns(Instruction *, ScalarEvolution *);
  bool processSExtIns(Instruction *, ScalarEvolution *);
  bool processStoreIns(Instruction *);
  bool processLoadIns(Instruction *);
  bool processPHIIns(Instruction *);
  bool isLoopInstruction(Instruction *, set<Instruction *> &);
  void analyseFunction(Function *, ScalarEvolution *, LoopInfo *);
  void analyseModule(Module &);
  void analyseFuncForLoops(Function *, LoopInfo *,
                           map<Function *, set<unsigned>> &);
  CallSiteArgVals *queryAttrValsAtCallSites(Module &);
  ~GeneralizedConstPropagation() {
    delete GCS;
    delete CSAV;
    delete VFM;
  }
};
} // namespace llvm

#endif // ends LLVM_TRANSFORMS_UTILS_GCP
