#include "mir/MIR.h"

#include <algorithm>
#include <vector>

#include "utils/Log.h"

namespace mir
{
  namespace
  {

    static int AlignTo(int value, int align)
    {
      if (value >= 0)
        return ((value + align - 1) / align) * align;
      else
        return ((value - align + 1) / align) * align;
    }

    static void ComputeFrameLayout(MachineFunction &function)
    {
      auto &slots = function.GetFrameSlots();
      const auto &callee_saved = function.GetCalleeSavedRegs();

      if (slots.empty() && callee_saved.empty())
      {
        function.SetFrameSize(0);
        return;
      }

      int callee_saved_bytes = 0;
      for (auto reg : callee_saved)
      {
        if (reg >= PhysReg::X0 && reg <= PhysReg::X30)
          callee_saved_bytes += 8;
        else
          callee_saved_bytes += 4;
      }
      callee_saved_bytes = AlignTo(callee_saved_bytes, 16);

      int offset = 0;
      for (auto &slot : slots)
      {
        if (slot.is_stack_arg || slot.is_callee_stack_arg)
          continue;
        offset -= slot.size;
        offset = AlignTo(offset, slot.size >= 8 ? 8 : 4);
        slot.offset = offset;
      }

      offset -= callee_saved_bytes;
      offset -= 16;
      offset = AlignTo(offset, 16);

      function.SetFrameSize(-offset);

      for (auto &slot : slots)
      {
        if (slot.is_callee_stack_arg)
        {
          slot.offset = 16 + slot.offset;
        }
      }
    }

    static void InsertPrologue(MachineFunction &function)
    {
      if (!function.GetEntryPtr())
        return;

      auto &entry = function.GetEntry();
      auto &insts = entry.GetInstructions();

      std::vector<MachineInstr> prologue_insts;
      prologue_insts.push_back(MachineInstr(Opcode::Prologue));

      auto it = insts.begin();
      while (it != insts.end() && it->GetOpcode() == Opcode::Prologue)
        ++it;

      insts.insert(it, prologue_insts.begin(), prologue_insts.end());
    }

    static void InsertEpilogue(MachineFunction &function)
    {
      for (auto &block : function.GetBlocks())
      {
        auto &insts = block->GetInstructions();
        for (auto it = insts.begin(); it != insts.end(); ++it)
        {
          if (it->GetOpcode() == Opcode::Ret)
          {
            std::vector<MachineInstr> epilogue_insts;
            epilogue_insts.push_back(MachineInstr(Opcode::Epilogue));

            it = insts.erase(it);
            insts.insert(it, epilogue_insts.begin(), epilogue_insts.end());
            break;
          }
        }
      }
    }

  } // namespace

  void RunFrameLowering(MachineFunction &function)
  {
    ComputeFrameLayout(function);
    InsertPrologue(function);
    InsertEpilogue(function);
  }

  void RunFrameLowering(MachineModule &module)
  {
    for (auto &function : module.GetFunctions())
    {
      if (function)
        RunFrameLowering(*function);
    }
  }

} // namespace mir