//===- Hello.cpp - Example code from "Writing an LLVM Pass" ---------------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// This file implements two versions of the LLVM "Hello World" pass described
// in docs/WritingAnLLVMPass.html
//
//===----------------------------------------------------------------------===//

#include "llvm/ADT/Statistic.h"
#include "llvm/IR/Function.h"
#include "llvm/Pass.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/Scalar.h"
#include "llvm/IR/InstIterator.h"
#include "llvm/IR/Instruction.h"
#include "llvm/Target/TargetLibraryInfo.h"
#include "llvm/Transforms/Utils/Local.h"
#include "llvm/Support/Debug.h"
#include "llvm/IR/Instructions.h"
#include "llvm/Analysis/ConstantFolding.h"
#include "llvm/ADT/BitVector.h"
#include "llvm/IR/ValueMap.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/CFG.h"
#include "llvm/ADT/SmallVector.h"
#include <vector>
#include <set>
using namespace llvm;

#define DEBUG_TYPE "hello"

STATISTIC(HelloCounter, "Counts number of functions greeted");

namespace {
  // Hello2 - The second implementation with getAnalysisUsage implemented.
  struct Hello2 : public FunctionPass {
    static char ID; // Pass identification, replacement for typeid
    Hello2() : FunctionPass(ID) {}

    bool runOnFunction(Function &F) override {
      ++HelloCounter;
      errs() << "Hello: ";
      errs().write_escaped(F.getName()) << '\n';
      return false;
    }

    // We don't modify the program, so we preserve all analyses.
    void getAnalysisUsage(AnalysisUsage &AU) const override {
      AU.setPreservesAll();
    }
  };
}

char Hello2::ID = 0;
static RegisterPass<Hello2>
Y("hello2", "Hello World Pass (with getAnalysisUsage implemented)");


namespace {
  // PrintBasicBlock
  struct PrintBasicBlock : public FunctionPass {
    static char ID;
    PrintBasicBlock(): FunctionPass(ID) {}

    bool runOnFunction(Function &F) override {
      errs() << "Function " << F.getName() << " basic blocks:\n";

      for (BasicBlock &BB: F) {
        printBasicBlock(BB);
      }

      return false;
    }

    void printBasicBlock(BasicBlock &BB) {
      for (Instruction &I: BB) {
        errs() << I << "\n";
      }
      errs() << "\n\n";
    }
  };
}

char PrintBasicBlock::ID = 0;
static RegisterPass<PrintBasicBlock>
Z("printbasicblock", "Prints Basic Blocks");


namespace {
  //===--------------------------------------------------------------------===//
  // DeadCodeElimination pass implementation
  //
  struct DCE : public FunctionPass {
    static char ID; // Pass identification, replacement for typeid
    DCE() : FunctionPass(ID) {
      initializeDCEPass(*PassRegistry::getPassRegistry());
    }

    bool runOnFunction(Function &F) override;

    void getAnalysisUsage(AnalysisUsage &AU) const override {
      AU.setPreservesCFG();
    }
 };
}

bool DCE::runOnFunction(Function &F) {
  if (skipOptnoneFunction(F))
    return false;

  TargetLibraryInfo *TLI = getAnalysisIfAvailable<TargetLibraryInfo>();

  // Start out with all of the instructions in the worklist...
  std::vector<Instruction*> WorkList;
  for (inst_iterator i = inst_begin(F), e = inst_end(F); i != e; ++i)
    WorkList.push_back(&*i);

  // Loop over the worklist finding instructions that are dead.  If they are
  // dead make them drop all of their uses, making other instructions
  // potentially dead, and work until the worklist is empty.
  //
  bool MadeChange = false;
  while (!WorkList.empty()) {
    Instruction *I = WorkList.back();
    WorkList.pop_back();

    errs() << "Current instruction: " << *I << "\n";
    errs() << "Instructions that use this one: \n";
    for (User *U: I->users())
      if (Instruction *Inst = dyn_cast<Instruction>(U)) {
        errs() << *Inst << "\n";
    }
    errs() << "\n\n";

    if (isInstructionTriviallyDead(I, TLI)) { // If the instruction is dead.
      // Loop over all of the values that the instruction uses, if there are
      // instructions being used, add them to the worklist, because they might
      // go dead after this one is removed.
      //

      for (User::op_iterator OI = I->op_begin(), E = I->op_end(); OI != E; ++OI) {
        if (Instruction *Used = dyn_cast<Instruction>(*OI)) {
          WorkList.push_back(Used);
        }
      }

      // Remove the instruction.
      I->eraseFromParent();

      // Remove the instruction from the worklist if it still exists in it.
      WorkList.erase(std::remove(WorkList.begin(), WorkList.end(), I),
                     WorkList.end());

      MadeChange = true;
    }
  }
  return MadeChange;
}

char DCE::ID = 0;
static RegisterPass<DCE> XX("mydce", "My DCE Pass");

FunctionPass *llvm::createDeadCodeEliminationPass() {
  return new DCE();
}


namespace {
  //===--------------------------------------------------------------------===//
  // DeadCodeElimination pass implementation
  //
  struct RemoveDeadBlocks : public FunctionPass {
    static char ID; // Pass identification, replacement for typeid
    RemoveDeadBlocks() : FunctionPass(ID) {}

    virtual bool runOnFunction(Function &F) override {
      bool first_block = true;

      std::vector<BasicBlock*> toRemove;

      for (BasicBlock &BB: F) {
        if (!first_block && pred_empty(&BB)) {
          toRemove.push_back(&BB);
        }
        first_block = false;
      }

      for (BasicBlock *BB: toRemove)
        BB->removeFromParent();

      return true;
    }

    void getAnalysisUsage(AnalysisUsage &AU) const override {
      AU.setPreservesCFG();
    }
 };
}

char RemoveDeadBlocks::ID = 0;
static RegisterPass<RemoveDeadBlocks>
RMDB("removedeadblocks", "Removes Dead Basic Blocks");


namespace {
  //===--------------------------------------------------------------------===//
  // DeadCodeElimination pass implementation
  //
  struct JumpToJump : public FunctionPass {
    static char ID; // Pass identification, replacement for typeid
    JumpToJump() : FunctionPass(ID) {}

    virtual bool runOnFunction(Function &F) override {
      for (BasicBlock &BB: F) {
        BranchInst *crt_bi;
        if (!BB.empty() && (crt_bi = dyn_cast<BranchInst>(&(*BB.begin())))) {
          if (crt_bi->getNumSuccessors() != 1)
            continue;
          BasicBlock *nextBB = crt_bi->getSuccessor(0);

          for (auto it = pred_begin(&BB); it != pred_end(&BB); ++it) {
            BasicBlock *predBB = *it;
            for (Instruction &I: *predBB) {
              Instruction *crtI = &I;
              if (BranchInst *bi = dyn_cast<BranchInst>(crtI)) {
                unsigned int n =  bi->getNumSuccessors();

                for (unsigned int i = 0; i < n; ++i) {
                  BasicBlock *succBB = bi->getSuccessor(i);
                  if (succBB == &BB) {
                    bi->setSuccessor(i, nextBB);
                  }
                }
              }
            }
          }
        }
      }

      return true;
    }

    void getAnalysisUsage(AnalysisUsage &AU) const override {
      AU.setPreservesCFG();
    }
 };
}

char JumpToJump::ID = 0;
static RegisterPass<JumpToJump>
JTJ("jumptojump", "Optimize jump to jump");