#include "LiveValues.h"

#include "llvm/IR/CFG.h"
#include "llvm/IR/InstIterator.h"
#include "llvm/IR/Instructions.h"

using namespace llvm;

static void applyMapping(InstructionSetVector& iset, llvm::DenseMap<llvm::Instruction *, llvm::Instruction *>& imap)
{
  // There will be probably be few entries in the imap, so apply them one at a time to the iset.
  for (auto& kv : imap)
  {
    if (iset.count(kv.first) != 0)
    {
      iset.remove(kv.first);
      iset.insert(kv.second);
    }
  }
}

// Compute liveness of a value at basic blocks. Roughly based on
// Algorithm 6 & 7 from the paper "Computing Liveness Sets for SSA-
// Form Programs" by Brander et al., 2011.

LiveValues::LiveValues(ArrayRef<Instruction*> computeLiveAt)
{
  m_liveSets.resize(computeLiveAt.size());

  // Build index and set of active blocks
  for (unsigned int i = 0; i < computeLiveAt.size(); i++)
  {
    Instruction* v = computeLiveAt[i];
    m_computeLiveAtIndex.insert(std::make_pair(v, i));

    m_activeBlocks.insert(v->getParent());
  }

  if (computeLiveAt.size() > 0)
  {
      m_function = computeLiveAt[0]->getParent()->getParent();
  }
}

// Go over all the instructions between begin (included) and end (excluded) and mark the given value
// live for code locations contained in the given range.
void LiveValues::markLiveRange(Instruction* value, BasicBlock::iterator begin, BasicBlock::iterator end)
{
  BasicBlock* B = begin->getParent();

  if (m_activeBlocks.count(B) == 0)
    return;  // Nothing to mark in this block

  for (BasicBlock::iterator I = begin; I != end; ++I)
  {
    if (m_computeLiveAtIndex.count(I))
    {
      // Mark this value
      unsigned int index = m_computeLiveAtIndex[I];
      m_liveSets[index].insert(value);
      m_allLiveSet.insert(value);
      // Also store for each value where it is live.
      m_liveAtIndices[value].insert(index);
    }
  }
}

void LiveValues::upAndMark(Instruction* def, Use& use, BlockSet& scanned)
{
  // Determine the starting point for the backwards search.
  // (Remember that Use represents an edge between the definition of a value and its use)
  // In the case in which the user of the use is a phi node we start the search from the terminator
  // of the preceding block.
  // This allows to avoid going through loop back-edges in cases like these:
  //                 |
  //                 | (y)
  //                 v
  //          -----------------
  //     (x)  | z = phi(x, y) |
  //    ----> | ...           |
  //    |     | x = z + 1     |
  //    |     -----------------
  //    |             |
  //    |             |
  //    |             |
  //    |             v
  //    |     -----------------
  //    |     |               |
  //    ------| INDIRECT CALL |
  //          |               |
  //          -----------------
  //                  | (Start the search for the definition of x (backwards) from here!)
  //                  v
  //
  // Notice that here x is live across the call. This case is tricky because the def comes 'after'
  // the use. The def still dominates the use because phi nodes logically use their input values on the
  // edges, i.e. on the terminator of the preceding blocks.
  //
  // This has the advantage of being able to traverse edges strictly backwards.

  Instruction* startingPoint = dyn_cast<Instruction>(use.getUser());
  if (PHINode* usePHI = dyn_cast<PHINode>(startingPoint))
  {
    BasicBlock* predecessor = usePHI->getIncomingBlock(use);
    startingPoint = predecessor->getTerminator();
  }

  BasicBlock* startingPointBB = startingPoint->getParent();
  BasicBlock* defBB = def->getParent();

  // Start a bottom-up recursive search from startingPoint to the definition of the current value.
  // Mark all the code ranges that we encounter on the way a having the current value 'live'.
  // 'scanned' contains the blocks that we have scanned to the bottom of the block and the we know
  // already having the current value 'live'.

  SmallVector<BasicBlock*, 16> worklist;
  worklist.push_back(startingPointBB);

  BlockSet visited;

  while (!worklist.empty())
  {
    BasicBlock* B = worklist.pop_back_val();

    if (scanned.count(B) != 0)
      continue;

    // We have reached the block that contains the definition of the value. We are done for this
    // branch of the search.
    if (B == defBB)
    {
      if (defBB == startingPointBB)
      {
        // If the first block that we visit is also the last mark only the range of instructions
        // between the def and the starting point.
        //    -----------------
        //    |               |
        //    | x = // def    |  <--
        //    |               |    !
        //    |               |    ! This is the range in which x is live.
        //    |               |    !
        //    | = x // use    |  <--
        //    |               |
        //    -----------------

        markLiveRange(def, ++BasicBlock::iterator(def), BasicBlock::iterator(startingPoint));
      }
      else
      {
        markLiveRange(def, ++BasicBlock::iterator(def), defBB->end());
        scanned.insert(B);
      }
    }
    else
    {
      if (B == startingPointBB)
      {
        // We are in the starting-point block.
        // This can mean two things:
        // 1. We are in the first iteration, mark the range between begin and starting point as
        // live.
        if (visited.count(B) == 0)
        {
          markLiveRange(def, B->begin(), BasicBlock::iterator(startingPoint));
        }
        // 2. We came back here because the starting point is in a loop.
        // In this case mark the whole block as live range and don't come back anymore.
        else
        {
          markLiveRange(def, B->begin(), B->end());
          scanned.insert(B);
        }

        // The if statement above allows to manage situations like this:
        //         BB0
        //        -----------------
        //        | x = ...       |
        //        -----------------
        //                |
        //                |
        //                |
        //         BB1    v
        //        -----------------<--                     <--
        //        |               |  !                       !
        //  ----->|               |  ! First range marked    !
        //  |     |               |  !                       !
        //  |     | ... = x       |<--                       ! Second and final range marked
        //  |     |               |                          !
        //  |     | INDIRECT CALL |                          !
        //  |     |               |                          !
        //  |     -----------------                        <--
        //  |              |
        //  ---------------
        // x is defined outside a loop and used inside a loop. This means that it is live inside the
        // whole loop.
        // So, we first mark the range from the use of x to the top of BB1 and, when we visit BB1
        // again (because BB1 is a predecessor of BB1) we mark the whole block as live range.
        // <rant>
        // This case could have been managed much more easily and efficiently if we had access to
        // LLVM LoopInfo analysis pass.
        // We could have done the following: x is uses in a loop and defined outside of it => mark
        // the whole loop body as live range.
        // </rant>
      }
      else
      {
        // We are in an intermediate block on the way to the definition mark it, all as live range.
        markLiveRange(def, B->begin(), B->end());
        scanned.insert(B);
      }

      visited.insert(B);

      for (pred_iterator P = pred_begin(B), PE = pred_end(B); P != PE; ++P)
      {
        worklist.push_back(*P);
      }
    }
  }
}

void LiveValues::run()
{
  if (m_computeLiveAtIndex.empty())
    return;

  // for each variable v do
  for (inst_iterator I = inst_begin(m_function), E = inst_end(m_function); I != E; ++I)
  {
    Instruction* v = &*I;
    assert(v->getParent()->getParent() == m_function);

    // for each block B where v is used do
    BlockSet scanned;
    for (Value::use_iterator U = v->use_begin(), UE = v->use_end(); U != UE; ++U)
    {
      Instruction* user = cast<Instruction>(U->getUser());
      assert(user->getParent()->getParent() == m_function);
      (void)user;

      upAndMark(v, *U, scanned);
    }
  }
}

void LiveValues::remapLiveValues(llvm::DenseMap<llvm::Instruction*, llvm::Instruction*>& imap)
{
  applyMapping(m_allLiveSet, imap);
  for (auto& liveSet : m_liveSets)
    applyMapping(liveSet, imap);
}

const LiveValues::Indices* LiveValues::getIndicesWhereLive(const Value* value) const
{
  const auto& iter = m_liveAtIndices.find(value);
  if (iter == m_liveAtIndices.end())
    return nullptr;
  return &iter->second;
}

void LiveValues::setIndicesWhereLive(Value* value, const Indices* indices)
{
  for (unsigned int idx : *indices)
    setLiveAtIndex(value, idx, true);
}

bool LiveValues::liveInDisjointRegions(const Value* valueA, const Value* valueB) const
{
  const Indices* indicesA = getIndicesWhereLive(valueA);
  if (!indicesA)
    return true;

  const Indices* indicesB = getIndicesWhereLive(valueB);
  if (!indicesB)
    return true;

  for (const unsigned int index : *indicesA)
  {
    if (indicesB->count(index))
      return false;
  }

  return true;
}

void LiveValues::setLiveAtIndex(Value* value, unsigned int index, bool live)
{
  assert(index <= m_computeLiveAtIndex.size());
  if (live)
  {
    m_liveAtIndices[value].insert(index);
    Instruction* inst = cast<Instruction>(value);
    m_liveSets[index].insert(inst);
    m_allLiveSet.insert(inst);
  }
  else
  {
    m_liveAtIndices[value].remove(index);
    Instruction* inst = cast<Instruction>(value);
    m_liveSets[index].remove(inst);
    if (m_liveAtIndices[value].empty())
      m_allLiveSet.remove(inst);
  }
}

void LiveValues::setLiveAtAllIndices(llvm::Value* value, bool live)
{
  Instruction* inst = cast<Instruction>(value);
  if (live)
  {
    for (unsigned int index = 0; index < m_computeLiveAtIndex.size(); ++index)
    {
      m_liveAtIndices[value].insert(index);
      m_liveSets[index].insert(inst);
    }
    m_allLiveSet.insert(inst);
  }
  else
  {
    for (unsigned int index = 0; index < m_computeLiveAtIndex.size(); ++index)
    {
      m_liveAtIndices[value].remove(index);
      m_liveSets[index].remove(inst);
    }
    if (m_liveAtIndices[value].empty())
      m_allLiveSet.remove(inst);
  }
}

bool LiveValues::getLiveAtIndex(const Value* value, unsigned int index) const
{
  assert(index <= m_computeLiveAtIndex.size());
  const auto& it = m_liveAtIndices.find(value);
  if (it == m_liveAtIndices.end())
    return false;
  return (it->second.count(index) != 0);
}