// Copyright 2014 Dolphin Emulator Project
// Licensed under GPLv2+
// Refer to the license.txt file included.

#include "Core/PowerPC/JitArm64/Jit.h"
#include "Common/Assert.h"
#include "Common/BitSet.h"
#include "Common/CommonTypes.h"
#include "Common/Logging/Log.h"
#include "Core/PowerPC/JitArm64/JitArm64_RegCache.h"

using namespace Arm64Gen;

void Arm64RegCache::Init(ARM64XEmitter* emitter)
{
  m_emit = emitter;
  m_float_emit.reset(new ARM64FloatEmitter(m_emit));
  GetAllocationOrder();
}

ARM64Reg Arm64RegCache::GetReg()
{
  // If we have no registers left, dump the most stale register first
  if (!GetUnlockedRegisterCount())
    FlushMostStaleRegister();

  for (auto& it : m_host_registers)
  {
    if (!it.IsLocked())
    {
      it.Lock();
      return it.GetReg();
    }
  }
  // Holy cow, how did you run out of registers?
  // We can't return anything reasonable in this case. Return INVALID_REG and watch the failure
  // happen
  WARN_LOG(DYNA_REC, "All available registers are locked dumb dumb");
  return INVALID_REG;
}

u32 Arm64RegCache::GetUnlockedRegisterCount()
{
  u32 unlocked_registers = 0;
  for (auto& it : m_host_registers)
    if (!it.IsLocked())
      ++unlocked_registers;
  return unlocked_registers;
}

void Arm64RegCache::LockRegister(ARM64Reg host_reg)
{
  auto reg = std::find(m_host_registers.begin(), m_host_registers.end(), host_reg);
  _assert_msg_(DYNA_REC, reg != m_host_registers.end(),
               "Don't try locking a register that isn't in the cache. Reg %d", host_reg);
  reg->Lock();
}

void Arm64RegCache::UnlockRegister(ARM64Reg host_reg)
{
  auto reg = std::find(m_host_registers.begin(), m_host_registers.end(), host_reg);
  _assert_msg_(DYNA_REC, reg != m_host_registers.end(),
               "Don't try unlocking a register that isn't in the cache. Reg %d", host_reg);
  reg->Unlock();
}

void Arm64RegCache::FlushMostStaleRegister()
{
  u32 most_stale_preg = 0;
  u32 most_stale_amount = 0;
  for (u32 i = 0; i < 32; ++i)
  {
    u32 last_used = m_guest_registers[i].GetLastUsed();
    if (last_used > most_stale_amount && (m_guest_registers[i].GetType() != REG_NOTLOADED &&
                                          m_guest_registers[i].GetType() != REG_IMM))
    {
      most_stale_preg = i;
      most_stale_amount = last_used;
    }
  }
  FlushRegister(most_stale_preg, false);
}

// GPR Cache
void Arm64GPRCache::Start(PPCAnalyst::BlockRegStats& stats)
{
}

bool Arm64GPRCache::IsCalleeSaved(ARM64Reg reg)
{
  static std::vector<ARM64Reg> callee_regs = {
      X28, X27, X26, X25, X24, X23, X22, X21, X20, X19, INVALID_REG,
  };
  return std::find(callee_regs.begin(), callee_regs.end(), EncodeRegTo64(reg)) != callee_regs.end();
}

void Arm64GPRCache::FlushRegister(u32 preg, bool maintain_state)
{
  OpArg& reg = m_guest_registers[preg];
  if (reg.GetType() == REG_REG)
  {
    ARM64Reg host_reg = reg.GetReg();
    if (reg.IsDirty())
      m_emit->STR(INDEX_UNSIGNED, host_reg, PPC_REG, PPCSTATE_OFF(gpr[preg]));

    if (!maintain_state)
    {
      UnlockRegister(host_reg);
      reg.Flush();
    }
  }
  else if (reg.GetType() == REG_IMM)
  {
    if (!reg.GetImm())
    {
      m_emit->STR(INDEX_UNSIGNED, WSP, PPC_REG, PPCSTATE_OFF(gpr[preg]));
    }
    else
    {
      ARM64Reg host_reg = GetReg();

      m_emit->MOVI2R(host_reg, reg.GetImm());
      m_emit->STR(INDEX_UNSIGNED, host_reg, PPC_REG, PPCSTATE_OFF(gpr[preg]));

      UnlockRegister(host_reg);
    }

    if (!maintain_state)
      reg.Flush();
  }
}

void Arm64GPRCache::FlushRegisters(BitSet32 regs, bool maintain_state)
{
  for (int i = 0; i < 32; ++i)
  {
    if (regs[i])
    {
      if (i < 31 && regs[i + 1])
      {
        // We've got two guest registers in a row to store
        OpArg& reg1 = m_guest_registers[i];
        OpArg& reg2 = m_guest_registers[i + 1];
        if (reg1.IsDirty() && reg2.IsDirty() && reg1.GetType() == REG_REG &&
            reg2.GetType() == REG_REG)
        {
          ARM64Reg RX1 = R(i);
          ARM64Reg RX2 = R(i + 1);

          m_emit->STP(INDEX_SIGNED, RX1, RX2, PPC_REG, PPCSTATE_OFF(gpr[0]) + i * sizeof(u32));
          if (!maintain_state)
          {
            UnlockRegister(RX1);
            UnlockRegister(RX2);
            reg1.Flush();
            reg2.Flush();
          }
          ++i;
          continue;
        }
      }

      FlushRegister(i, maintain_state);
    }
  }
}

void Arm64GPRCache::Flush(FlushMode mode, PPCAnalyst::CodeOp* op)
{
  BitSet32 to_flush;
  for (int i = 0; i < 32; ++i)
  {
    bool flush = true;
    if (m_guest_registers[i].GetType() == REG_REG)
    {
      // Has to be flushed if it isn't in a callee saved register
      ARM64Reg host_reg = m_guest_registers[i].GetReg();
      flush = IsCalleeSaved(host_reg) ? flush : true;
    }

    to_flush[i] = flush;
  }
  FlushRegisters(to_flush, mode == FLUSH_MAINTAIN_STATE);
}

ARM64Reg Arm64GPRCache::R(u32 preg)
{
  OpArg& reg = m_guest_registers[preg];
  IncrementAllUsed();
  reg.ResetLastUsed();

  switch (reg.GetType())
  {
  case REG_REG:  // already in a reg
    return reg.GetReg();
    break;
  case REG_IMM:  // Is an immediate
  {
    ARM64Reg host_reg = GetReg();
    m_emit->MOVI2R(host_reg, reg.GetImm());
    reg.Load(host_reg);
    reg.SetDirty(true);
    return host_reg;
  }
  break;
  case REG_NOTLOADED:  // Register isn't loaded at /all/
  {
    // This is a bit annoying. We try to keep these preloaded as much as possible
    // This can also happen on cases where PPCAnalyst isn't feeing us proper register usage
    // statistics
    ARM64Reg host_reg = GetReg();
    reg.Load(host_reg);
    reg.SetDirty(false);
    m_emit->LDR(INDEX_UNSIGNED, host_reg, PPC_REG, PPCSTATE_OFF(gpr[preg]));
    return host_reg;
  }
  break;
  default:
    ERROR_LOG(DYNA_REC, "Invalid OpArg Type!");
    break;
  }
  // We've got an issue if we end up here
  return INVALID_REG;
}

void Arm64GPRCache::SetImmediate(u32 preg, u32 imm)
{
  OpArg& reg = m_guest_registers[preg];
  if (reg.GetType() == REG_REG)
    UnlockRegister(reg.GetReg());
  reg.LoadToImm(imm);
}

void Arm64GPRCache::BindToRegister(u32 preg, bool do_load)
{
  OpArg& reg = m_guest_registers[preg];

  reg.ResetLastUsed();

  reg.SetDirty(true);
  if (reg.GetType() == REG_NOTLOADED)
  {
    ARM64Reg host_reg = GetReg();
    reg.Load(host_reg);
    if (do_load)
      m_emit->LDR(INDEX_UNSIGNED, host_reg, PPC_REG, PPCSTATE_OFF(gpr[preg]));
  }
}

void Arm64GPRCache::GetAllocationOrder()
{
  // Callee saved registers first in hopes that we will keep everything stored there first
  const std::vector<ARM64Reg> allocation_order = {
      // Callee saved
      W27, W26, W25, W24, W23, W22, W21, W20, W19,

      // Caller saved
      W18, W17, W16, W15, W14, W13, W12, W11, W10, W9, W8, W7, W6, W5, W4, W3, W2, W1, W0, W30,
  };

  for (ARM64Reg reg : allocation_order)
    m_host_registers.push_back(HostReg(reg));
}

BitSet32 Arm64GPRCache::GetCallerSavedUsed()
{
  BitSet32 registers(0);
  for (auto& it : m_host_registers)
    if (it.IsLocked() && !IsCalleeSaved(it.GetReg()))
      registers[it.GetReg()] = 1;
  return registers;
}

void Arm64GPRCache::FlushByHost(ARM64Reg host_reg)
{
  for (int i = 0; i < 32; ++i)
  {
    OpArg& reg = m_guest_registers[i];
    if (reg.GetType() == REG_REG && reg.GetReg() == host_reg)
    {
      FlushRegister(i, false);
      return;
    }
  }
}

// FPR Cache
void Arm64FPRCache::Flush(FlushMode mode, PPCAnalyst::CodeOp* op)
{
  for (int i = 0; i < 32; ++i)
  {
    if (m_guest_registers[i].GetType() != REG_NOTLOADED &&
        m_guest_registers[i].GetType() != REG_IMM)
    {
      // XXX: Determine if we can keep a register in the lower 64bits
      // Which will allow it to be callee saved.
      FlushRegister(i, mode == FLUSH_MAINTAIN_STATE);
    }
  }
}

ARM64Reg Arm64FPRCache::R(u32 preg, RegType type)
{
  OpArg& reg = m_guest_registers[preg];
  IncrementAllUsed();
  reg.ResetLastUsed();
  ARM64Reg host_reg = reg.GetReg();

  switch (reg.GetType())
  {
  case REG_REG_SINGLE:
  {
    // We're asked for singles, so just return the register.
    if (type == REG_REG_SINGLE || type == REG_LOWER_PAIR_SINGLE)
      return host_reg;

    // Else convert this register back to doubles.
    m_float_emit->FCVTL(64, EncodeRegToDouble(host_reg), EncodeRegToDouble(host_reg));
    reg.Load(host_reg, REG_REG);

    // fall through
  }
  case REG_REG:  // already in a reg
  {
    return host_reg;
  }
  case REG_LOWER_PAIR_SINGLE:
  {
    // We're asked for the lower single, so just return the register.
    if (type == REG_LOWER_PAIR_SINGLE)
      return host_reg;

    // Else convert this register back to a double.
    m_float_emit->FCVT(64, 32, EncodeRegToDouble(host_reg), EncodeRegToDouble(host_reg));
    reg.Load(host_reg, REG_LOWER_PAIR);

    // fall through
  }
  case REG_LOWER_PAIR:
  {
    if (type == REG_REG)
    {
      // Load the high 64bits from the file and insert them in to the high 64bits of the host
      // register
      ARM64Reg tmp_reg = GetReg();
      m_float_emit->LDR(64, INDEX_UNSIGNED, tmp_reg, PPC_REG, PPCSTATE_OFF(ps[preg][1]));
      m_float_emit->INS(64, host_reg, 1, tmp_reg, 0);
      UnlockRegister(tmp_reg);

      // Change it over to a full 128bit register
      reg.Load(host_reg, REG_REG);
    }
    return host_reg;
  }
  case REG_DUP_SINGLE:
  {
    if (type == REG_LOWER_PAIR_SINGLE)
      return host_reg;

    if (type == REG_REG_SINGLE)
    {
      // Duplicate to the top and change over
      m_float_emit->INS(32, host_reg, 1, host_reg, 0);
      reg.Load(host_reg, REG_REG_SINGLE);
      return host_reg;
    }

    m_float_emit->FCVT(64, 32, EncodeRegToDouble(host_reg), EncodeRegToDouble(host_reg));
    reg.Load(host_reg, REG_DUP);

    // fall through
  }
  case REG_DUP:
  {
    if (type == REG_REG)
    {
      // We are requesting a full 128bit register
      // but we are only available in the lower 64bits
      // Duplicate to the top and change over
      m_float_emit->INS(64, host_reg, 1, host_reg, 0);
      reg.Load(host_reg, REG_REG);
    }
    return host_reg;
  }
  case REG_NOTLOADED:  // Register isn't loaded at /all/
  {
    host_reg = GetReg();
    u32 load_size;
    if (type == REG_REG)
    {
      load_size = 128;
      reg.Load(host_reg, REG_REG);
    }
    else
    {
      load_size = 64;
      reg.Load(host_reg, REG_LOWER_PAIR);
    }
    reg.SetDirty(false);
    m_float_emit->LDR(load_size, INDEX_UNSIGNED, host_reg, PPC_REG, PPCSTATE_OFF(ps[preg][0]));
    return host_reg;
  }
  default:
    _dbg_assert_msg_(DYNA_REC, false, "Invalid OpArg Type!");
    break;
  }
  // We've got an issue if we end up here
  return INVALID_REG;
}

ARM64Reg Arm64FPRCache::RW(u32 preg, RegType type)
{
  OpArg& reg = m_guest_registers[preg];

  bool was_dirty = reg.IsDirty();

  IncrementAllUsed();
  reg.ResetLastUsed();

  reg.SetDirty(true);

  // If not loaded at all, just alloc a new one.
  if (reg.GetType() == REG_NOTLOADED)
  {
    reg.Load(GetReg(), type);
    return reg.GetReg();
  }

  // Only the lower value will be overwritten, so we must be extra careful to store PSR1 if dirty.
  if ((type == REG_LOWER_PAIR || type == REG_LOWER_PAIR_SINGLE) && was_dirty)
  {
    // We must *not* change host_reg as this register might still be in use. So it's fine to
    // store this register, but it's *not* fine to convert it to double. So for double convertion,
    // a temporary register needs to be used.
    ARM64Reg host_reg = reg.GetReg();
    ARM64Reg flush_reg = host_reg;

    switch (reg.GetType())
    {
    case REG_REG_SINGLE:
      flush_reg = GetReg();
      m_float_emit->FCVTL(64, EncodeRegToDouble(flush_reg), EncodeRegToDouble(host_reg));
    // fall through
    case REG_REG:
      // We are doing a full 128bit store because it takes 2 cycles on a Cortex-A57 to do a 128bit
      // store.
      // It would take longer to do an insert to a temporary and a 64bit store than to just do this.
      m_float_emit->STR(128, INDEX_UNSIGNED, flush_reg, PPC_REG, PPCSTATE_OFF(ps[preg][0]));
      break;
    case REG_DUP_SINGLE:
      flush_reg = GetReg();
      m_float_emit->FCVT(64, 32, EncodeRegToDouble(flush_reg), EncodeRegToDouble(host_reg));
    // fall through
    case REG_DUP:
      // Store PSR1 (which is equal to PSR0) in memory.
      m_float_emit->STR(64, INDEX_UNSIGNED, flush_reg, PPC_REG, PPCSTATE_OFF(ps[preg][1]));
      break;
    default:
      // All other types doesn't store anything in PSR1.
      break;
    }

    if (host_reg != flush_reg)
      Unlock(flush_reg);
  }

  reg.Load(reg.GetReg(), type);
  return reg.GetReg();
}

void Arm64FPRCache::GetAllocationOrder()
{
  const std::vector<ARM64Reg> allocation_order = {// Callee saved
                                                  Q8, Q9, Q10, Q11, Q12, Q13, Q14, Q15,

                                                  // Caller saved
                                                  Q16, Q17, Q18, Q19, Q20, Q21, Q22, Q23, Q24, Q25,
                                                  Q26, Q27, Q28, Q29, Q30, Q31, Q7, Q6, Q5, Q4, Q3,
                                                  Q2, Q1, Q0};

  for (ARM64Reg reg : allocation_order)
    m_host_registers.push_back(HostReg(reg));
}

void Arm64FPRCache::FlushByHost(ARM64Reg host_reg)
{
  for (int i = 0; i < 32; ++i)
  {
    OpArg& reg = m_guest_registers[i];
    if ((reg.GetType() != REG_NOTLOADED && reg.GetType() != REG_IMM) && reg.GetReg() == host_reg)
    {
      FlushRegister(i, false);
      return;
    }
  }
}

bool Arm64FPRCache::IsCalleeSaved(ARM64Reg reg)
{
  static std::vector<ARM64Reg> callee_regs = {
      Q8, Q9, Q10, Q11, Q12, Q13, Q14, Q15, INVALID_REG,
  };
  return std::find(callee_regs.begin(), callee_regs.end(), reg) != callee_regs.end();
}

void Arm64FPRCache::FlushRegister(u32 preg, bool maintain_state)
{
  OpArg& reg = m_guest_registers[preg];
  ARM64Reg host_reg = reg.GetReg();
  RegType type = reg.GetType();
  bool dirty = reg.IsDirty();

  // If we're in single mode, just convert it back to a double.
  if (type == REG_REG_SINGLE)
  {
    if (dirty)
      m_float_emit->FCVTL(64, EncodeRegToDouble(host_reg), EncodeRegToDouble(host_reg));
    type = REG_REG;
  }
  if (type == REG_DUP_SINGLE || type == REG_LOWER_PAIR_SINGLE)
  {
    if (dirty)
      m_float_emit->FCVT(64, 32, EncodeRegToDouble(host_reg), EncodeRegToDouble(host_reg));

    if (type == REG_DUP_SINGLE)
      type = REG_DUP;
    else
      type = REG_LOWER_PAIR;
  }

  if (type == REG_REG || type == REG_LOWER_PAIR)
  {
    u32 store_size;
    if (type == REG_REG)
      store_size = 128;
    else
      store_size = 64;

    if (dirty)
      m_float_emit->STR(store_size, INDEX_UNSIGNED, host_reg, PPC_REG, PPCSTATE_OFF(ps[preg][0]));

    if (!maintain_state)
    {
      UnlockRegister(host_reg);
      reg.Flush();
    }
  }
  else if (type == REG_DUP)
  {
    if (dirty)
    {
      // If the paired registers were at the start of ppcState we could do an STP here.
      // Too bad moving them would break savestate compatibility between x86_64 and AArch64
      // m_float_emit->STP(64, INDEX_SIGNED, host_reg, host_reg, PPC_REG,
      // PPCSTATE_OFF(ps[preg][0]));
      m_float_emit->STR(64, INDEX_UNSIGNED, host_reg, PPC_REG, PPCSTATE_OFF(ps[preg][0]));
      m_float_emit->STR(64, INDEX_UNSIGNED, host_reg, PPC_REG, PPCSTATE_OFF(ps[preg][1]));
    }

    if (!maintain_state)
    {
      UnlockRegister(host_reg);
      reg.Flush();
    }
  }
}

void Arm64FPRCache::FlushRegisters(BitSet32 regs, bool maintain_state)
{
  for (int j : regs)
    FlushRegister(j, maintain_state);
}

BitSet32 Arm64FPRCache::GetCallerSavedUsed()
{
  BitSet32 registers(0);
  for (auto& it : m_host_registers)
    if (it.IsLocked())
      registers[it.GetReg() - Q0] = 1;
  return registers;
}

bool Arm64FPRCache::IsSingle(u32 preg, bool lower_only)
{
  RegType type = m_guest_registers[preg].GetType();
  return type == REG_REG_SINGLE || type == REG_DUP_SINGLE ||
         (lower_only && type == REG_LOWER_PAIR_SINGLE);
}

void Arm64FPRCache::FixSinglePrecision(u32 preg)
{
  OpArg& reg = m_guest_registers[preg];
  ARM64Reg host_reg = reg.GetReg();
  switch (reg.GetType())
  {
  case REG_DUP:  // only PS0 needs to be converted
    m_float_emit->FCVT(32, 64, EncodeRegToDouble(host_reg), EncodeRegToDouble(host_reg));
    reg.Load(host_reg, REG_DUP_SINGLE);
    break;
  case REG_REG:  // PS0 and PS1 needs to be converted
    m_float_emit->FCVTN(32, EncodeRegToDouble(host_reg), EncodeRegToDouble(host_reg));
    reg.Load(host_reg, REG_REG_SINGLE);
    break;
  default:
    break;
  }
}
