#include "llvm/Pass.h"
#include "llvm/ADT/SmallString.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Analysis/LoopPass.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/Analysis/OrderedBasicBlock.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/InstVisitor.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/MDBuilder.h"
#include "llvm/IR/Metadata.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "llvm/Transforms/Utils/Cloning.h"
#include "llvm/Transforms/Utils/LoopUtils.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"

// Defines the value to used during debugging
#define DEBUG_TYPE "predictive"

// Defines the type of indexed expression. Possible values are {RHS, LHS}
#define IE_TYPE "loop.type"

#define IE_REPLACE "loop.replace"

// Defines the constant which replaces the induction variable
#define IE_ITER "loop.iter"

// Defines the array affected by the transformation
#define IE_ARRAY "loop.array"

// Marks the instruction which will be useless after transformation
#define IE_DEAD "loop.dead"

using namespace llvm;

namespace {
  class Pass1 : public LoopPass {
  public:
    static char ID;
    LoopInfo *LI;
    Pass1() : LoopPass(ID) {}

    void getAnalysisUsage(AnalysisUsage &AU) const override {
      AU.setPreservesAll();
      AU.addRequired<LoopInfoWrapperPass>();
      AU.addPreserved<LoopInfoWrapperPass>();
      getLoopAnalysisUsage(AU);
    }

    int getInitial(Instruction *index, Instruction *arrayptr) {
      for (auto tmpL = index,
           tmpA = arrayptr;
           tmpL != tmpA;
           tmpL = tmpL->user_back()
      ) {
        if(auto bi = dyn_cast<BinaryOperator>(tmpL)) {
          for (auto l = bi->op_begin(); l != bi->op_end(); ++l) {
            if (auto c = dyn_cast<ConstantInt>(l)) {
              return c->getSExtValue();
            }
          }
        }
      }
      return 0;
    }

    void setUsersMetadata(Instruction *start, Instruction *end, MDNode *node) {
      for (auto tmpL = start,
           tmpA = end;
           tmpL != tmpA;
           tmpL = tmpL->user_back()
      ) {
        tmpL->setMetadata(IE_DEAD, node);
      }
    }

    Constant* makeConstant(LLVMContext& context, int value) {
      return ConstantInt::get(Type::getInt32Ty(context), value);
    }

    void setAllMetadata(BasicBlock *BB) {
      OrderedBasicBlock oBB(BB);
      LLVMContext& context = BB->getContext();
      MDBuilder builder(context);
      // int count = 0;
      for (auto gepIT = BB->begin(); gepIT != BB->end(); ++gepIT) {
        if (auto *gep = dyn_cast<GetElementPtrInst>(gepIT)) {
          auto *arr = MDNode::get(context, builder.createString("array"));
          cast<Instruction>(gep->getOperand(0))->setMetadata(IE_ARRAY, arr);
          auto *RHS = MDNode::get(context, builder.createString("RHS"));
          auto *LHS = MDNode::get(context, builder.createString("LHS"));
          Instruction *dominant = nullptr;
          for (auto it = BB->begin(); it != BB->end(); ++it) {
            if (oBB.dominates(&*it, gep) && isa<LoadInst>(it)) {
              dominant = &*it;
            }
          }
          auto *dominated = gep->user_back();
          bool lhs = isa<LoadInst>(dominated);
          if (dominant != nullptr && dominated != nullptr) {
            if (lhs) {
              auto deadInst = MDNode::get(context, builder.createString("dead"));
              gepIT->setMetadata(IE_TYPE, RHS);
              gepIT->setMetadata(IE_DEAD, deadInst);
              setUsersMetadata(dominant, &*gepIT, deadInst);
              auto cnt = getInitial(dominant, &*gepIT);
              auto iterV = makeConstant(context, cnt);
              auto iterMD = MDNode::get(context, builder.createConstant(iterV));
              dominated->setMetadata(IE_ITER, iterMD);
              dominated->setMetadata(IE_DEAD, deadInst);
              dominant->setMetadata(IE_DEAD, deadInst);
            } else {
              gepIT->setMetadata(IE_TYPE, LHS);
            }
          }
        }
      }
    }

    bool runOnLoop(Loop *loop, LPPassManager &LPM) override {
      LI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
      for (auto block : loop->getBlocks()) {
        if (LI->isLoopHeader(block)) {
          continue;
        }
        if (loop->isLoopExiting(block)) {
          continue;
        }
        SmallString<10> blockName(block->getName());
        StringRef requiredPrefix("for.body", 8);
        if (!blockName.startswith(requiredPrefix)) {
          continue;
        }
        setAllMetadata(block);
      }
      return true;
    }
  };
}

char Pass1::ID = 0;
static RegisterPass<Pass1> X(
  "pass1",
  "Predictive Commoning"
);

namespace {
  struct Pass2 : public FunctionPass, public InstVisitor<Pass2> {
  public:
    static char ID;

    Pass2() : FunctionPass(ID) {}

    void getAnalysisUsage(AnalysisUsage &AU) const override {
      //AU.setPreservesCFG();
      AU.addRequired<LoopInfoWrapperPass>();
      AU.addPreserved<LoopInfoWrapperPass>();
    }

    Constant* makeConstant(LLVMContext& context, int value) {
      return ConstantInt::get(Type::getInt32Ty(context), value);
    }

    bool runOnFunction(Function &F) override {
      SmallString<10> fname(F.getName());
      if (fname.startswith("_")) {
        return false;
      }

      LoopInfo& LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
      auto entry = &F.getEntryBlock();
      AllocaInst *arr = nullptr;
      for (auto i = entry->begin(); i != entry->end(); i++) {
        if(i->getMetadata(IE_ARRAY)) {
          arr = cast<AllocaInst>(i);
        }
      }

      LLVMContext& ctx = F.getContext();
      BasicBlock* block = BasicBlock::Create(ctx, "predict", &F);
      IRBuilder<> builder(block);

      for (auto loops = LI.begin(); loops != LI.end(); loops++) {
        Loop *L = *loops;
        BasicBlock *lbody;
        SmallVector<AllocaInst*, 4> newAllocs;
        SmallVector<Instruction*, 4> oldLoads;
        for (auto it = L->block_begin(); it != L->block_end(); ++it) {
          BasicBlock *BB = *it;
          SmallString<10> bname(BB->getName());
          if (bname.startswith("for.body")) {
            lbody = BB;
            break;
          }
        }
        for (auto inst = lbody->begin(); inst != lbody->end(); inst++) {
          if (isa<GetElementPtrInst>(inst)) {
            if (auto md = inst->getMetadata(IE_TYPE)) {
              auto type = cast<MDString>(md->getOperand(0))->getString();
              if (type == "LHS") {
                continue;
              }
              auto arrMD = inst->user_back()->getMetadata(IE_ITER);
              int index = mdconst::extract<ConstantInt>(arrMD->getOperand(0))->getSExtValue();
              auto p = builder.CreateAlloca(Type::getInt32Ty(ctx), nullptr, "p");
              p->setAlignment(4);
              newAllocs.push_back(p);
              oldLoads.push_back(inst->user_back());
              auto cnt = makeConstant(ctx, 0);
              auto arrindex = makeConstant(ctx, index);
              auto gepI = builder.CreateInBoundsGEP(arr, { cnt, arrindex });
              auto lI = builder.CreateAlignedLoad(gepI, 4);
              builder.CreateAlignedStore(lI, p, 4);
            } else {
              continue;
            }
          }
        }
        IRBuilder<> loopBuilder(lbody, lbody->getFirstInsertionPt());
        int i = 0;
        for (auto it : newAllocs) {
          auto newLoad = loopBuilder.CreateAlignedLoad(it, 4);
          oldLoads[i]->replaceAllUsesWith(newLoad);
          i++;
        }
      }
      builder.CreateBr(entry->getSingleSuccessor());
      entry->getTerminator()->setSuccessor(0, block);
      return true;
    }
  };
}

char Pass2::ID = 0;
static RegisterPass<Pass2> Y(
  "pass2",
  "PredictiveCommoning pass 2"
);