//===- Reg2MemHLSL.cpp - Convert registers to allocas ---------------------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//

#include "llvm/ADT/Statistic.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/CFG.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/Module.h"
#include "llvm/Pass.h"
#include "llvm/Transforms/Scalar.h"
#include "llvm/Transforms/Utils/Local.h"
#include <list>
using namespace llvm;

#define DEBUG_TYPE "reg2mem_hlsl"

STATISTIC(NumRegsDemotedHlsl, "Number of registers demoted");
STATISTIC(NumPhisDemotedHlsl, "Number of phi-nodes demoted");

namespace {
struct RegToMemHlsl : public FunctionPass {
  static char ID; // Pass identification, replacement for typeid
  RegToMemHlsl() : FunctionPass(ID) {
    initializeRegToMemHlslPass(*PassRegistry::getPassRegistry());
  }

  void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.addRequiredID(BreakCriticalEdgesID);
    AU.addPreservedID(BreakCriticalEdgesID);
  }

  bool valueEscapes(const Instruction *Inst) const {
    for (const User *U : Inst->users()) {
      const Instruction *UI = cast<Instruction>(U);
      if (isa<PHINode>(UI))
        return true;
    }
    return false;
  }

  bool runOnFunction(Function &F) override;
};

/// DemotePHIToStack - This function takes a virtual register computed by a
/// PHI node and replaces it with a slot in the stack frame allocated via
/// alloca.
/// The PHI node is deleted. It returns the pointer to the alloca inserted.
/// The difference of HLSL version is the new Alloca will be loaded for each
/// use, for case a phi inside loop be used outside the loop.
AllocaInst *DemotePHIToStack_HLSL(PHINode *P, Instruction *AllocaPoint) {
  if (P->use_empty()) {
    P->eraseFromParent();
    return nullptr;
  }

  IRBuilder<> AllocaBuilder(P);
  if (!AllocaPoint) {
    Function *F = P->getParent()->getParent();
    AllocaPoint = F->getEntryBlock().begin();
  }
  AllocaBuilder.SetInsertPoint(AllocaPoint);

  // Create a stack slot to hold the value.
  AllocaInst *Slot = AllocaBuilder.CreateAlloca(P->getType(), nullptr,
                                                P->getName() + ".reg2mem");

  // Insert a load in place of the PHI and replace all uses.
  BasicBlock::iterator InsertPt = P;

  for (; isa<PHINode>(InsertPt) || isa<LandingPadInst>(InsertPt); ++InsertPt)
    /* empty */; // Don't insert before PHI nodes or landingpad instrs.

  std::vector<Instruction *> WorkList;
  for (auto U = P->user_begin(); U != P->user_end();) {
    Instruction *I = cast<Instruction>(*(U++));
    WorkList.emplace_back(I);
  }

  for (Instruction *I : WorkList) {
    IRBuilder<> Builder(I);
    Value *Load = Builder.CreateLoad(Slot);
    I->replaceUsesOfWith(P, Load);
  }

  // Iterate over each operand inserting a store in each predecessor.
  // This should be done after load inserting because store for phi must be
  // after all other instructions of the incoming block.
  for (unsigned i = 0, e = P->getNumIncomingValues(); i < e; ++i) {
    if (InvokeInst *II = dyn_cast<InvokeInst>(P->getIncomingValue(i))) {
      assert(II->getParent() != P->getIncomingBlock(i) &&
             "Invoke edge not supported yet");
      (void)II;
    }
    Value *V = P->getIncomingValue(i);
    // Skip undef
    if (isa<UndefValue>(V))
      continue;
    new StoreInst(P->getIncomingValue(i), Slot,
                  P->getIncomingBlock(i)->getTerminator());
  }

  // Delete PHI.
  P->eraseFromParent();
  return Slot;
}

/// DemoteRegToStack - This function takes a virtual register computed by an
/// Instruction and replaces it with a slot in the stack frame, allocated via
/// alloca.  This allows the CFG to be changed around without fear of
/// invalidating the SSA information for the value.  It returns the pointer to
/// the alloca inserted to create a stack slot for I.
/// The difference of HLSL version is for I is Alloca, only replace new Alloca
/// with old alloca, and HLSL don't have InvokeInst
AllocaInst *DemoteRegToStack_HLSL(Instruction &I, bool VolatileLoads,
                                  Instruction *AllocaPoint) {
  if (I.use_empty()) {
    I.eraseFromParent();
    return nullptr;
  }

  IRBuilder<> AllocaBuilder(&I);
  if (!AllocaPoint) {
    Function *F = I.getParent()->getParent();
    AllocaPoint = F->getEntryBlock().begin();
  }
  AllocaBuilder.SetInsertPoint(AllocaPoint);

  if (AllocaInst *AI = dyn_cast<AllocaInst>(&I)) {
    // Create a stack slot to hold the value.
    AllocaInst *Slot = AllocaBuilder.CreateAlloca(
        AI->getAllocatedType(), nullptr, I.getName() + ".reg2mem");
    I.replaceAllUsesWith(Slot);
    I.eraseFromParent();
    return Slot;
  }

  // Create a stack slot to hold the value.
  AllocaInst *Slot = AllocaBuilder.CreateAlloca(I.getType(), nullptr,
                                                I.getName() + ".reg2mem");
  ;

  // Change all of the users of the instruction to read from the stack slot.
  while (!I.use_empty()) {
    Instruction *U = cast<Instruction>(I.user_back());
    if (PHINode *PN = dyn_cast<PHINode>(U)) {
      // If this is a PHI node, we can't insert a load of the value before the
      // use.  Instead insert the load in the predecessor block corresponding
      // to the incoming value.
      //
      // Note that if there are multiple edges from a basic block to this PHI
      // node that we cannot have multiple loads. The problem is that the
      // resulting PHI node will have multiple values (from each load) coming
      // in
      // from the same block, which is illegal SSA form. For this reason, we
      // keep track of and reuse loads we insert.
      DenseMap<BasicBlock *, Value *> Loads;
      for (unsigned i = 0, e = PN->getNumIncomingValues(); i != e; ++i)
        if (PN->getIncomingValue(i) == &I) {
          Value *&V = Loads[PN->getIncomingBlock(i)];
          if (!V) {
            // Insert the load into the predecessor block
            V = new LoadInst(Slot, I.getName() + ".reload", VolatileLoads,
                             PN->getIncomingBlock(i)->getTerminator());
          }
          PN->setIncomingValue(i, V);
        }

    } else {
      // If this is a normal instruction, just insert a load.
      Value *V = new LoadInst(Slot, I.getName() + ".reload", VolatileLoads, U);
      U->replaceUsesOfWith(&I, V);
    }
  }

  // Insert stores of the computed value into the stack slot. We have to be
  // careful if I is an invoke instruction, because we can't insert the store
  // AFTER the terminator instruction.
  BasicBlock::iterator InsertPt;
  if (!isa<TerminatorInst>(I)) {
    InsertPt = &I;
    ++InsertPt;
    for (; isa<PHINode>(InsertPt) || isa<LandingPadInst>(InsertPt); ++InsertPt)
      /* empty */; // Don't insert before PHI nodes or landingpad instrs.
  } else {
    InvokeInst &II = cast<InvokeInst>(I);
    InsertPt = II.getNormalDest()->getFirstInsertionPt();
  }

  new StoreInst(&I, Slot, InsertPt);
  return Slot;
}
} // namespace

char RegToMemHlsl::ID = 0;
INITIALIZE_PASS_BEGIN(RegToMemHlsl, "reg2mem_hlsl",
                      "Demote values with phi-node usage to stack slots", false,
                      false)
INITIALIZE_PASS_DEPENDENCY(BreakCriticalEdges)
INITIALIZE_PASS_END(RegToMemHlsl, "reg2mem_hlsl",
                    "Demote values with phi-node usage to stack slots", false,
                    false)

bool RegToMemHlsl::runOnFunction(Function &F) {
  if (F.isDeclaration())
    return false;

  // Insert all new allocas into entry block.
  BasicBlock *BBEntry = &F.getEntryBlock();
  assert(pred_empty(BBEntry) &&
         "Entry block to function must not have predecessors!");

  // Find first non-alloca instruction and create insertion point. This is
  // safe if block is well-formed: it always have terminator, otherwise
  // we'll get and assertion.
  BasicBlock::iterator I = BBEntry->begin();
  while (isa<AllocaInst>(I))
    ++I;

  CastInst *AllocaInsertionPoint = new BitCastInst(
      Constant::getNullValue(Type::getInt32Ty(F.getContext())),
      Type::getInt32Ty(F.getContext()), "reg2mem_hlsl alloca point", I);

  // Find the escaped instructions. But don't create stack slots for
  // allocas in entry block.
  std::list<Instruction *> WorkList;
  for (Function::iterator ibb = F.begin(), ibe = F.end(); ibb != ibe; ++ibb)
    for (BasicBlock::iterator iib = ibb->begin(), iie = ibb->end(); iib != iie;
         ++iib) {
      if (!(isa<AllocaInst>(iib) && iib->getParent() == BBEntry) &&
          valueEscapes(iib)) {
        WorkList.push_front(&*iib);
      }
    }

  // Demote escaped instructions
  NumRegsDemotedHlsl += WorkList.size();
  for (std::list<Instruction *>::iterator ilb = WorkList.begin(),
                                          ile = WorkList.end();
       ilb != ile; ++ilb)
    DemoteRegToStack_HLSL(**ilb, false, AllocaInsertionPoint);

  WorkList.clear();

  // Find all phi's
  for (Function::iterator ibb = F.begin(), ibe = F.end(); ibb != ibe; ++ibb)
    for (BasicBlock::iterator iib = ibb->begin(), iie = ibb->end(); iib != iie;
         ++iib)
      if (isa<PHINode>(iib))
        WorkList.push_front(&*iib);

  // Demote phi nodes
  NumPhisDemotedHlsl += WorkList.size();
  for (std::list<Instruction *>::iterator ilb = WorkList.begin(),
                                          ile = WorkList.end();
       ilb != ile; ++ilb)
    DemotePHIToStack_HLSL(cast<PHINode>(*ilb), AllocaInsertionPoint);

  return true;
}

// createDemoteRegisterToMemoryHlsl - Provide an entry point to create this
// pass.
char &llvm::DemoteRegisterToMemoryHlslID = RegToMemHlsl::ID;
FunctionPass *llvm::createDemoteRegisterToMemoryHlslPass() {
  return new RegToMemHlsl();
}
