//===----------------------------------------------------------------------===//
//
/// A register allocator simplified from RegAllocFast.cpp
//
//===----------------------------------------------------------------------===//
#include "llvm/ADT/Statistic.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/CodeGen/MachineFunctionPass.h"
#include "llvm/CodeGen/MachineRegisterInfo.h"
#include "llvm/CodeGen/RegAllocRegistry.h"
#include "llvm/CodeGen/RegisterClassInfo.h"
#include "llvm/CodeGen/TargetInstrInfo.h"
#include "llvm/CodeGen/MachineFrameInfo.h"
#include "llvm/InitializePasses.h"
#include "llvm/Pass.h"
#include "llvm/CodeGen/LiveIntervals.h"
#include "llvm/CodeGen/SlotIndexes.h"
using namespace llvm;
#define DEBUG_TYPE "regalloc"
STATISTIC(NumStores, "Number of stores (spills) added");
STATISTIC(NumLoads , "Number of loads (reloads) added");
namespace {
/// This is class where you will implement your register allocator in
class RegAllocSimple : public MachineFunctionPass {
public:
  static char ID;
  RegAllocSimple() : MachineFunctionPass(ID) {}
private:
  /// Some information that might be useful for register allocation
  /// They are initialized in runOnMachineFunction
  MachineFrameInfo *MFI;
  MachineRegisterInfo *MRI;
  TargetRegisterInfo const *TRI;
  TargetInstrInfo const *TII;
  RegisterClassInfo RegClassInfo;
  // Virtual register to physical register mapping for the current block
  DenseMap<MachineOperand, MCPhysReg> LiveVirtRegs;
  // Physical register to virtual register mapping for the current block
  DenseMap<MCPhysReg, MachineOperand> LivePhysRegs;
  // Virtual register to stack slot mapping for the current block
  DenseMap<MachineOperand, int> LiveSpills;
  // Physical registers used in the current instruction
  SmallSet<MCPhysReg, 8> UsedInInstr;
public:
  StringRef getPassName() const override { return "Simple Register Allocator"; }
  void getAnalysisUsage(AnalysisUsage &AU) const override 
  {
    AU.setPreservesCFG();
    // At -O1/-O2, llc fails to schedule some required passes if this pass
    // does not preserve these anlyses; these are preserved by recomputing
    // them at the end of runOnFunction(), you can safely ignore these
    AU.addRequired<LiveIntervals>();
    AU.addPreserved<LiveIntervals>();
    AU.addRequired<SlotIndexes>();
    AU.addPreserved<SlotIndexes>();
    MachineFunctionPass::getAnalysisUsage(AU);
  }
  /// Ask the Machine IR verifier to check some simple properties
  /// Enabled with the -verify-machineinstrs flag in llc
  MachineFunctionProperties getRequiredProperties() const override 
  {
    return MachineFunctionProperties().set(
        MachineFunctionProperties::Property::NoPHIs);
  }
  MachineFunctionProperties getSetProperties() const override 
  {
    return MachineFunctionProperties().set(
        MachineFunctionProperties::Property::NoVRegs);
  }
  MachineFunctionProperties getClearedProperties() const override 
  {
    return MachineFunctionProperties().set(
      MachineFunctionProperties::Property::IsSSA);
  }
private:
  /// Allocate physical register for virtual register operand
  void allocateOperand(MachineOperand &MO, bool const isUse) 
  {
    // If we've already allocated a phys register, we don't need to do anything
    auto it = LiveVirtRegs.find(MO);
    if (it != LiveVirtRegs.end()) {
      return;
    }
    // Get a list of phys registers that can hold the operand in preferred order
    auto* RC = MRI->getRegClass(MO.getReg());
    auto AllocOrder = RegClassInfo.getOrder(RC);
    MCPhysReg Reg = [AllocOrder, this]() {
      auto const regIsUsed = [this](MCPhysReg const& Reg){ return this->LivePhysRegs.find(Reg) != this->LivePhysRegs.end(); };
      auto const regIsSpillable = [this](MCPhysReg const& Reg){ return this->UsedInInstr.contains(Reg); };
      
      auto unusedPhysRegs = AllocOrder.drop_while(regIsUsed);
      if (unusedPhysRegs.empty()) {
        auto spillablePhysRegs = AllocOrder.drop_while(regIsSpillable);
        MCPhysReg spilled = spillablePhysRegs.front();
        auto virt_spilled = this->LivePhysRegs.find(spilled);
        assert(virt_spilled != this->LivePhysRegs.end() && "Attempting to spill unbound physical register");
        this->spill(virt_spilled->second);
        return spilled;
      } else {
        return unusedPhysRegs.front();
      }
    }();
    // If we are using a virtual register and it's not already live,
    // it must be spilled, so we reload it here.
    if (isUse) {
      reload(MO, Reg);
    }
    bindPhysReg(MO, Reg);
    UsedInInstr.insert(Reg);
    LiveVirtRegs.insert({MO, Reg});
    LivePhysRegs.insert({Reg, MO});
  }
  void allocateInstruction(MachineInstr &MI) 
  {
    static SmallVector<MachineOperand, 8> Uses;
    static SmallVector<MachineOperand, 1> RegMask;
    static SmallVector<MachineOperand, 8> Defs;
    // We haven't allocated any physical registers yet
    UsedInInstr.clear();
    Uses.clear();
    RegMask.clear();
    Defs.clear();
    // Split the operands into the categories we care about
    for (auto& MO: MI.operands()) {
      if (MO.isUse())
        Uses.emplace_back(MO);
      if (MO.isRegMask())
        RegMask.emplace_back(MO);
      if (MO.isDef())
        Defs.emplace_back(MO);
    }
    auto const allocUse = [this](MachineOperand &MO){ this->allocateOperand(MO, true); };
    auto const allocDef = [this](MachineOperand &MO){ this->allocateOperand(MO, false); };
  
    // Allocate uses first
    for_each(Uses, allocUse);
    // Spill all clobbered regs before a call
    for(auto& MO: RegMask) {
      for (auto &[O, R]: LiveVirtRegs)
        if (MO.clobbersPhysReg(R)) {
          spill(O);
        }
    }
    // Allocate defs
    for_each(Defs, allocDef);
  }
  void allocateBasicBlock(MachineBasicBlock &MBB) 
  {
    // We haven't allocated any regs yet
    LiveVirtRegs.clear();
    // Allocate each instruction
    for (auto &MI: MBB) {
      allocateInstruction(MI);
      if (MI.isReturn())
        return; // We don't need to spill for returns
    }
    // Spill all live regs at the end
    for (auto &[O, _]: LiveVirtRegs) {
      spill(O);
    }
  }
  bool runOnMachineFunction(MachineFunction &MF) override 
  {
    // Get some useful information about the target
    MRI = &MF.getRegInfo();
    const TargetSubtargetInfo &STI = MF.getSubtarget();
    TRI = STI.getRegisterInfo();
    TII = STI.getInstrInfo();
    MFI = &MF.getFrameInfo();
    MRI->freezeReservedRegs(MF);
    RegClassInfo.runOnMachineFunction(MF);
    // Allocate each basic block locally
    for (MachineBasicBlock &MBB : MF) {
      allocateBasicBlock(MBB);
    }
    MRI->clearVirtRegs();
    // Recompute the analyses that we marked as preserved above, you can
    // safely ignore this code
    SlotIndexes& SI = getAnalysis<SlotIndexes>();
    SI.releaseMemory();
    SI.runOnMachineFunction(MF);
    LiveIntervals& LI = getAnalysis<LiveIntervals>();
    LI.releaseMemory();
    LI.runOnMachineFunction(MF);
    return true;
  }
  void bindPhysReg(MachineOperand &MO, MCPhysReg Reg)
  {
    auto idx = MO.getSubReg();
    if (idx) {
      Reg = TRI->getSubReg(Reg, idx);
      MO.setSubReg(0);
    }
    MO.setReg(Reg);
    if (MO.isKill())
      MO.setIsKill(false);
    else if (MO.isDead())
      MO.setIsDead(false);
    MO.setIsRenamable();
  }
  void spill(MachineOperand &MO)
  {
    assert(MO.getReg().isPhysical() && "Cannot spill something that is not a physical register");
    ++NumStores;
    // Clear the use
    LiveVirtRegs.erase(MO);
    LivePhysRegs.erase(MO.getReg());
    UsedInInstr.erase(MO.getReg());
    // Get the stack slot to spill into
    auto* RC = MRI->getRegClass(MO.getReg());
    auto slot = MFI->CreateSpillStackObject(TRI->getSpillSize(*RC), TRI->getSpillAlign(*RC));
    
    // Spill the register
    TII->storeRegToStackSlot(*MO.getMBB(), MO.getMBB()->begin(), MO.getReg(), MO.isKill(), slot, RC, TRI);   
    LiveSpills.insert({MO, slot});
  }
  void reload(MachineOperand &MO, MCPhysReg Dest)
  {
    assert(MO.getReg().isStack() && "Cannot reload something that is not a stack slot");
    ++NumLoads;
    auto* RC = MRI->getRegClass(MO.getReg());
    auto slot = LiveSpills.find(MO);
    assert(slot != LiveSpills.end() && "Did not find spill slot to reload");
    TII->loadRegFromStackSlot(*MO.getMBB(), MO.getMBB()->begin(), Dest, slot->second, RC, TRI);
    LiveSpills.erase(MO);
  }
};
} // namespace
/// Create the initializer and register the pass
char RegAllocSimple::ID = 0;
FunctionPass *llvm::createSimpleRegisterAllocator() { return new RegAllocSimple(); }
INITIALIZE_PASS(RegAllocSimple, "regallocsimple", "Simple Register Allocator", false, false)
static RegisterRegAlloc simpleRegAlloc("simple", "simple register allocator", createSimpleRegisterAllocator);