//===-- SlotIndexes.cpp - Slot Indexes Pass  ------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "llvm/CodeGen/SlotIndexes.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/CodeGen/MachineFunction.h"
#include "llvm/Config/llvm-config.h"
#include "llvm/InitializePasses.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"

using namespace llvm;

#define DEBUG_TYPE "slot-indexes"

AnalysisKey SlotIndexesAnalysis::Key;

SlotIndexesAnalysis::Result
SlotIndexesAnalysis::run(MachineFunction &MF,
                         MachineFunctionAnalysisManager &) {
  return Result(MF);
}

PreservedAnalyses
SlotIndexesPrinterPass::run(MachineFunction &MF,
                            MachineFunctionAnalysisManager &MFAM) {
  OS << "Slot indexes in machine function: " << MF.getName() << '\n';
  MFAM.getResult<SlotIndexesAnalysis>(MF).print(OS);
  return PreservedAnalyses::all();
}
char SlotIndexesWrapperPass::ID = 0;

SlotIndexesWrapperPass::SlotIndexesWrapperPass() : MachineFunctionPass(ID) {}

SlotIndexes::~SlotIndexes() {
  // The indexList's nodes are all allocated in the BumpPtrAllocator.
  indexList.clear();
}

INITIALIZE_PASS(SlotIndexesWrapperPass, DEBUG_TYPE, "Slot index numbering",
                false, false)

STATISTIC(NumLocalRenum,  "Number of local renumberings");

void SlotIndexesWrapperPass::getAnalysisUsage(AnalysisUsage &au) const {
  au.setPreservesAll();
  MachineFunctionPass::getAnalysisUsage(au);
}

void SlotIndexes::clear() {
  mi2iMap.clear();
  MBBRanges.clear();
  idx2MBBMap.clear();
  indexList.clear();
  ileAllocator.Reset();
}

void SlotIndexes::analyze(MachineFunction &fn) {

  // Compute numbering as follows:
  // Grab an iterator to the start of the index list.
  // Iterate over all MBBs, and within each MBB all MIs, keeping the MI
  // iterator in lock-step (though skipping it over indexes which have
  // null pointers in the instruction field).
  // At each iteration assert that the instruction pointed to in the index
  // is the same one pointed to by the MI iterator. This

  // FIXME: This can be simplified. The mi2iMap_, Idx2MBBMap, etc. should
  // only need to be set up once after the first numbering is computed.

  mf = &fn;

  // Check that the list contains only the sentinel.
  assert(indexList.empty() && "Index list non-empty at initial numbering?");
  assert(idx2MBBMap.empty() &&
         "Index -> MBB mapping non-empty at initial numbering?");
  assert(MBBRanges.empty() &&
         "MBB -> Index mapping non-empty at initial numbering?");
  assert(mi2iMap.empty() &&
         "MachineInstr -> Index mapping non-empty at initial numbering?");

  unsigned index = 0;
  MBBRanges.resize(mf->getMaxAnalysisBlockNumber());
  idx2MBBMap.reserve(mf->size());

  indexList.push_back(*createEntry(nullptr, index));

  // Iterate over the function.
  for (MachineBasicBlock &MBB : *mf) {
    // Insert an index for the MBB start.
    SlotIndex blockStartIndex(&indexList.back(), SlotIndex::Slot_Block);

    for (MachineInstr &MI : MBB) {
      if (MI.isDebugOrPseudoInstr())
        continue;

      // Insert a store index for the instr.
      indexList.push_back(*createEntry(&MI, index += SlotIndex::InstrDist));

      // Save this base index in the maps.
      mi2iMap.insert(std::make_pair(
          &MI, SlotIndex(&indexList.back(), SlotIndex::Slot_Block)));
    }

    // We insert one blank instructions between basic blocks.
    indexList.push_back(*createEntry(nullptr, index += SlotIndex::InstrDist));

    MBBRanges[MBB.getAnalysisNumber()].first = blockStartIndex;
    MBBRanges[MBB.getAnalysisNumber()].second =
        SlotIndex(&indexList.back(), SlotIndex::Slot_Block);
    idx2MBBMap.push_back(IdxMBBPair(blockStartIndex, &MBB));
  }

  // Sort the Idx2MBBMap
  llvm::sort(idx2MBBMap, less_first());

  LLVM_DEBUG(mf->print(dbgs(), this));
}

bool SlotIndexes::isBlockBoundaryIndex(SlotIndex Idx) const {
  if (getInstructionFromIndex(Idx))
    return false;

  // Adjacent blocks share a boundary entry, so a boundary is either the start
  // index of the block Idx falls in or the end index of the last block. A block
  // dropped by removeMBBFromMaps() is gone from idx2MBBMap, so its start entry
  // classifies as stale.
  assert(!idx2MBBMap.empty() && "Index -> MBB mapping is empty");
  SlotIndex Base = Idx.getBaseIndex();
  MBBIndexIterator I = std::prev(getMBBUpperBound(Base));
  return Base == I->first || Base == getMBBEndIdx(I->second);
}

SlotIndex SlotIndexes::canonicalizeIndex(SlotIndex Idx) const {
  if (!isStaleIndex(Idx))
    return Idx;

  // The block start is a boundary entry, so it bounds the walk.
  SlotIndex BlockStart = getMBBStartIdx(getMBBFromIndex(Idx));
  IndexList::iterator I = Idx.listEntry()->getIterator();
  while (&*I != BlockStart.listEntry()) {
    --I;
    if (I->getInstr())
      return SlotIndex(&*I, SlotIndex::Slot_Register);
  }
  return BlockStart;
}

void SlotIndexes::removeMachineInstrFromMaps(MachineInstr &MI,
                                             bool AllowBundled) {
  assert((AllowBundled || !MI.isBundledWithPred()) &&
         "Use removeSingleMachineInstrFromMaps() instead");
  Mi2IndexMap::iterator mi2iItr = mi2iMap.find(&MI);
  if (mi2iItr == mi2iMap.end())
    return;

  SlotIndex MIIndex = mi2iItr->second;
  IndexListEntry &MIEntry = *MIIndex.listEntry();
  assert(MIEntry.getInstr() == &MI && "Instruction indexes broken.");
  mi2iMap.erase(mi2iItr);
  // FIXME: Eventually we want to actually delete these indexes.
  MIEntry.setInstr(nullptr);
}

void SlotIndexes::removeSingleMachineInstrFromMaps(MachineInstr &MI) {
  Mi2IndexMap::iterator mi2iItr = mi2iMap.find(&MI);
  if (mi2iItr == mi2iMap.end())
    return;

  SlotIndex MIIndex = mi2iItr->second;
  IndexListEntry &MIEntry = *MIIndex.listEntry();
  assert(MIEntry.getInstr() == &MI && "Instruction indexes broken.");
  mi2iMap.erase(mi2iItr);

  // When removing the first instruction of a bundle update mapping to next
  // instruction.
  if (MI.isBundledWithSucc()) {
    // Only the first instruction of a bundle should have an index assigned.
    assert(!MI.isBundledWithPred() && "Should be first bundle instruction");

    MachineBasicBlock::instr_iterator Next = std::next(MI.getIterator());
    MachineInstr &NextMI = *Next;
    MIEntry.setInstr(&NextMI);
    mi2iMap.insert(std::make_pair(&NextMI, MIIndex));
    return;
  } else {
    // FIXME: Eventually we want to actually delete these indexes.
    MIEntry.setInstr(nullptr);
  }
}

void SlotIndexes::removeMBBFromMaps(MachineBasicBlock &MBB) {
  assert(&MBB != &MBB.getParent()->front() &&
         "Can't remove the first block of a function.");

  unsigned Num = MBB.getAnalysisNumber();
  SlotIndex StartIdx = MBBRanges[Num].first;
  SlotIndex EndIdx = MBBRanges[Num].second;

  // Give MBB's slot range to its layout predecessor so blocks stay contiguous.
  auto PrevMBB = std::prev(MBB.getIterator());
  MBBRanges[PrevMBB->getAnalysisNumber()].second = EndIdx;

  // Drop MBB's index -> MBB entry, which would dangle once MBB is erased.
  auto It = getMBBLowerBound(StartIdx);
  assert(It != MBBIndexEnd() && It->first == StartIdx && It->second == &MBB &&
         "MBB not found in index -> MBB map");
  idx2MBBMap.erase(It);

  // Clear the block-start boundary entry. MBBRanges is never renumbered, so
  // MBB's now-stale slot is simply left in place.
  StartIdx.listEntry()->setInstr(nullptr);
}

// Renumber indexes locally after curItr was inserted, but failed to get a new
// index.
void SlotIndexes::renumberIndexes(IndexList::iterator curItr) {
  // Number indexes with half the default spacing so we can catch up quickly.
  const unsigned Space = SlotIndex::InstrDist/2;
  static_assert((Space & 3) == 0, "InstrDist must be a multiple of 2*NUM");

  IndexList::iterator startItr = std::prev(curItr);
  unsigned index = startItr->getIndex();
  unsigned BeginIndex = index;
  do {
    curItr->setIndex(index += Space);
    ++curItr;
    // If the next index is bigger, we have caught up.
  } while (curItr != indexList.end() && curItr->getIndex() <= index);

  LLVM_DEBUG(dbgs() << "\n*** Renumbered SlotIndexes " << startItr->getIndex()
                    << '-' << index << " ***\n");

  // If we repack more than 20% of a function, add spacing in between the
  // instructions so that future renumberings are able to catch up
  // without also renumbering so much.
  if (index - BeginIndex >
      (getLastIndex().getIndex() - getZeroIndex().getIndex()) / 5)
    packIndexes();

  ++NumLocalRenum;
}

// Repair indexes after adding and removing instructions.
void SlotIndexes::repairIndexesInRange(MachineBasicBlock *MBB,
                                       MachineBasicBlock::iterator Begin,
                                       MachineBasicBlock::iterator End) {
  bool includeStart = (Begin == MBB->begin());
  SlotIndex startIdx;
  if (includeStart)
    startIdx = getMBBStartIdx(MBB);
  else
    startIdx = getInstructionIndex(*--Begin);

  SlotIndex endIdx;
  if (End == MBB->end())
    endIdx = getMBBEndIdx(MBB);
  else
    endIdx = getInstructionIndex(*End);

  // FIXME: Conceptually, this code is implementing an iterator on MBB that
  // optionally includes an additional position prior to MBB->begin(), indicated
  // by the includeStart flag. This is done so that we can iterate MIs in a MBB
  // in parallel with SlotIndexes, but there should be a better way to do this.
  IndexList::iterator ListB = startIdx.listEntry()->getIterator();
  IndexList::iterator ListI = endIdx.listEntry()->getIterator();
  MachineBasicBlock::iterator MBBI = End;
  bool pastStart = false;
  bool OldIndexesRemoved = false;
  while (ListI != ListB || MBBI != Begin || (includeStart && !pastStart)) {
    assert(ListI->getIndex() >= startIdx.getIndex() &&
           (includeStart || !pastStart) &&
           "Decremented past the beginning of region to repair.");

    MachineInstr *SlotMI = ListI->getInstr();
    MachineInstr *MI = (MBBI != MBB->end() && !pastStart) ? &*MBBI : nullptr;
    bool MBBIAtBegin = MBBI == Begin && (!includeStart || pastStart);
    bool MIIndexNotFound = MI && !mi2iMap.contains(MI);
    bool SlotMIRemoved = false;

    if (SlotMI == MI && !MBBIAtBegin) {
      --ListI;
      if (MBBI != Begin)
        --MBBI;
      else
        pastStart = true;
    } else if (MIIndexNotFound || OldIndexesRemoved) {
      if (MBBI != Begin)
        --MBBI;
      else
        pastStart = true;
    } else {
      // We ran through all the indexes on the interval
      //   -> The only thing left is to go through all the
      //   remaining MBB instructions and update their indexes
      if (ListI == ListB)
        OldIndexesRemoved = true;
      else
        --ListI;
      if (SlotMI) {
        removeMachineInstrFromMaps(*SlotMI);
        SlotMIRemoved = true;
      }
    }

    MachineInstr *InstrToInsert = SlotMIRemoved ? SlotMI : MI;

    // Insert instruction back into the maps after passing it/removing the index
    if ((MIIndexNotFound || SlotMIRemoved) && InstrToInsert->getParent() &&
        !InstrToInsert->isDebugOrPseudoInstr())
      insertMachineInstrInMaps(*InstrToInsert);
  }
}

void SlotIndexes::packIndexes() {
  for (auto [Index, Entry] : enumerate(indexList))
    Entry.setIndex(Index * SlotIndex::InstrDist);
}

void SlotIndexes::print(raw_ostream &OS) const {
  for (const IndexListEntry &ILE : indexList) {
    OS << ILE.getIndex() << ' ';

    if (ILE.getInstr())
      OS << *ILE.getInstr();
    else
      OS << '\n';
  }

  for (const MachineBasicBlock &MBB : *mf)
    OS << printMBBReference(MBB) << "\t[" << getMBBStartIdx(&MBB) << ';'
       << getMBBEndIdx(&MBB) << ")\n";
}

#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
LLVM_DUMP_METHOD void SlotIndexes::dump() const { print(dbgs()); }
#endif

// Print a SlotIndex to a raw_ostream.
void SlotIndex::print(raw_ostream &os) const {
  if (isValid())
    os << listEntry()->getIndex() << "Berd"[getSlot()];
  else
    os << "invalid";
}

#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
// Dump a SlotIndex to stderr.
LLVM_DUMP_METHOD void SlotIndex::dump() const {
  print(dbgs());
  dbgs() << "\n";
}
#endif