//===----------------------------------------------------------------------===//
//
/// 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;
    }
    assert(MO.getReg().isVirtual() && "Can only allocate a virtual register");
    
    // 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){ 
      assert(MO.isUse());
      LLVM_DEBUG(dbgs() << "Regalloc: Allocating use " << MO << "\n");
      this->allocateOperand(MO, true); 
    };
    auto const allocDef = [this](MachineOperand &MO){ 
      assert(MO.isDef());
      LLVM_DEBUG(dbgs() << "Regalloc: Allocating def " << MO << "\n");
      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)) {
          LLVM_DEBUG(dbgs() << "Regalloc: Spilling clobbered reg " << R << "\n");
          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)
  {
    assert(MO.getReg().isVirtual() && "Cannot bind a physical register to a non-virtual register");

    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);