#include "mir/MIR.h"

#include <stdexcept>
#include <vector>

#include "utils/Log.h"

namespace mir {
namespace {

int AlignTo(int value, int align) {
  return ((value + align - 1) / align) * align;
}

}  // namespace

void RunFrameLowering(MachineFunction& function) {
  int cursor = 0;
  for (const auto& slot : function.GetFrameSlots()) {
    cursor += slot.size;
  }
  
  // Align stack frames to 16 bytes for AArch64
  cursor = AlignTo(cursor, 16);

  cursor = 0;
  for (const auto& slot : function.GetFrameSlots()) {
    cursor += slot.size;
    function.GetFrameSlot(slot.index).offset = -cursor;
  }
  function.SetFrameSize(AlignTo(cursor, 16));

  auto& blocks = function.GetBlocks();
  if (blocks.empty()) return;

  // Insert Prologue at the start of the first block
  auto& entry_insts = blocks.front().GetInstructions();
  entry_insts.insert(entry_insts.begin(), MachineInstr(Opcode::Prologue));

  // Insert Epilogue before every Ret in all blocks
  for (auto& block : blocks) {
    auto& insts = block.GetInstructions();
    std::vector<MachineInstr> lowered;
    for (const auto& inst : insts) {
      if (inst.GetOpcode() == Opcode::Ret) {
        lowered.emplace_back(Opcode::Epilogue);
      }
      lowered.push_back(inst);
    }
    insts = std::move(lowered);
  }
}

}  // namespace mir