//===-- W65816ImgCalleeSave.cpp - Callee-side save/restore of IMG8..IMG15 -===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// Post-RA, pre-PEI pass that adds prologue save + epilogue restore for
// IMG8..IMG15 ($C0..$CE) in any function that uses them.  This makes
// IMG8..IMG15 behave as callee-saved AT THE ASM LEVEL without going
// through LLVM's CSR mechanism (which would shift regalloc decisions
// and break other tests — see history with
// `feedback_picol_expr_compound_or.md`).
//
// Why callee-side, not caller-side?
//
//   Callers can hold long-lived vregs in IMG8..IMG15 (regalloc treats
//   them as preserved across calls because they're not in JSLpseudo's
//   Defs).  The "obvious" fix — add them to Defs and force regalloc to
//   spill them across each call — interacts badly with stack-slot
//   coloring: the spill slot gets coalesced with another vreg whose
//   liveness appears disjoint, but the post-call reload makes the
//   lifetimes overlap and the reload reads garbage (caught by qsort,
//   strncat, etc. when IMG0..IMG7 were also in Defs).
//
//   By doing the save/restore on the CALLEE side instead, the caller
//   doesn't need to spill at all — its values in IMG8..IMG15 are
//   automatically preserved.  Only functions that USE IMG8..IMG15 pay
//   the cost (a few bytes of prologue/epilogue), and the cost is
//   amortized across the whole function (not per call).
//
// Why post-RA, not via LLVM's CSR mechanism?
//
//   Adding IMG8..IMG15 to `getCalleeSavedRegs()` makes LLVM treat them
//   as "expensive" in cost-of-use analysis.  Regalloc steers away from
//   them in functions that don't really need them, but that steering
//   changes coloring decisions in ways that broke strtol
//   (`strtol(" 0x1ABC ", &ep, 16)` returned 0).  Implementing
//   save/restore outside the CSR system keeps regalloc's decisions
//   unchanged: it sees IMG8..IMG15 as ordinary regs, uses them freely
//   under pressure, and this pass adds the asm-level bookkeeping.
//
// Why pre-PEI?
//
//   PEI is what assigns frame-index offsets and emits the actual
//   prologue/epilogue.  To add new spill slots, we need PEI to see
//   them so they get included in the frame size.  We use
//   `MFI.CreateStackObject` to register the slots, then emit STAfi /
//   LDAfi pseudos that PEI will lower to `STA d,s` / `LDA d,s`.
//
//   We also insert the save/restore as REGULAR MIR instructions BEFORE
//   PEI runs.  That means PEI sees them when it emits its frame setup,
//   and the STAfi/LDAfi disps are valid post-PEI.
//
//===----------------------------------------------------------------------===//

#include "W65816.h"
#include "W65816InstrInfo.h"
#include "W65816Subtarget.h"
#include "llvm/CodeGen/MachineFrameInfo.h"
#include "llvm/CodeGen/MachineFunction.h"
#include "llvm/CodeGen/MachineFunctionPass.h"
#include "llvm/CodeGen/MachineInstrBuilder.h"
#include "llvm/Support/Debug.h"

using namespace llvm;

#define DEBUG_TYPE "w65816-img-callee-save"

namespace {

class W65816ImgCalleeSave : public MachineFunctionPass {
public:
  static char ID;
  W65816ImgCalleeSave() : MachineFunctionPass(ID) {}
  bool runOnMachineFunction(MachineFunction &MF) override;
  StringRef getPassName() const override {
    return "W65816 callee-side save/restore for IMG8..IMG15";
  }
};

} // namespace

char W65816ImgCalleeSave::ID = 0;

INITIALIZE_PASS(W65816ImgCalleeSave, DEBUG_TYPE,
                "W65816 IMG8..IMG15 callee save/restore", false, false)

FunctionPass *llvm::createW65816ImgCalleeSave() {
  return new W65816ImgCalleeSave();
}

// IMG8..IMG15 physregs (in order so IMG_REGS[i] is the i'th high-half slot).
// Their DP addresses are $C0, $C2, ..., $CE (each slot is 16 bits = 2 bytes);
// the DP layout is also expressed via W65816Helpers::imgDPAddr.  Keep the
// parallel `IMG_DP` array for fast index→address lookup at the hot rewrite
// sites below.
static constexpr unsigned IMG_REGS[8] = {
    W65816::IMG8,  W65816::IMG9,  W65816::IMG10, W65816::IMG11,
    W65816::IMG12, W65816::IMG13, W65816::IMG14, W65816::IMG15};
static constexpr unsigned IMG_DP[8] = {0xC0, 0xC2, 0xC4, 0xC6,
                                       0xC8, 0xCA, 0xCC, 0xCE};

static int classifyImgReg(unsigned Reg) {
  for (int i = 0; i < 8; ++i)
    if (Reg == IMG_REGS[i])
      return i;
  return -1;
}

// Classification of a DP-addressed instruction's relation to a DP slot.
enum class DpAccess {
  None,        // not a DP-imm instruction we care about
  Read,        // only reads the DP slot (e.g., LDA $C0)
  Write,       // only writes the DP slot (e.g., STA $C0)
  ReadWrite,   // both (e.g., INC $C0)
};

// Map a DP-addressed instruction's first immediate operand to an IMG
// slot index and access mode.  Returns (-1, None) if it doesn't access
// an IMG slot.
static std::pair<int, DpAccess> classifyDpImmAsImg(const MachineInstr &MI) {
  unsigned Opc = MI.getOpcode();
  DpAccess Mode;
  switch (Opc) {
  // Pure stores: write only.
  case W65816::STA_DP:
  case W65816::STZ_DP:
  case W65816::STX_DP:
  case W65816::STY_DP:
    Mode = DpAccess::Write;
    break;
  // Pure loads / compares / bit-tests: read only (writes to A/X/Y/P, not DP).
  case W65816::LDA_DP:
  case W65816::LDX_DP:
  case W65816::LDY_DP:
  case W65816::ADC_DP:
  case W65816::SBC_DP:
  case W65816::AND_DP:
  case W65816::ORA_DP:
  case W65816::EOR_DP:
  case W65816::CMP_DP:
  case W65816::CPX_DP:
  case W65816::CPY_DP:
  case W65816::BIT_DP:
    Mode = DpAccess::Read;
    break;
  // Read-modify-write.
  case W65816::INC_DP:
  case W65816::DEC_DP:
  case W65816::ASL_DP:
  case W65816::LSR_DP:
  case W65816::ROL_DP:
  case W65816::ROR_DP:
    Mode = DpAccess::ReadWrite;
    break;
  default:
    return {-1, DpAccess::None};
  }
  for (const auto &MO : MI.operands()) {
    if (!MO.isImm()) continue;
    int64_t V = MO.getImm();
    for (int i = 0; i < 8; ++i)
      if ((int64_t)IMG_DP[i] == V)
        return {i, Mode};
    return {-1, DpAccess::None};  // First imm is the dp addr; not in IMG range.
  }
  return {-1, DpAccess::None};
}

// DBG_VALUE preservation in this pass:
//
// This pass only ADDS instructions (PHA/LDA_DP/STAfi at function entry,
// PHA/LDAfi/STA_DP/PLA at each return-block exit).  It never erases,
// moves, or modifies user-emitted instructions, and it doesn't
// substitute one register/operand for another.  No DBG_VALUE updates
// are needed.
bool W65816ImgCalleeSave::runOnMachineFunction(MachineFunction &MF) {
  // Step 1: scan for IMG8..IMG15 WRITES.  Reads alone don't need saving
  // — if we never write IMGn, the caller's value survives untouched
  // (other functions we call also preserve IMG8..IMG15 by the same
  // convention, so no chain breaks the invariant).  Saving on read-only
  // use costs ~6 bytes per slot of needlessly-saved prologue/epilogue
  // (caught by evalAt at 1.96× Calypsi — 5 IMG slots saved when fewer
  // were actually written).
  //
  // copyPhysReg lowers `COPY $imgN = $a` to `STA_DP imm:0xCx`, so we
  // check both the physreg-DEF form AND the DP-imm-store form.
  bool UsedSlot[8] = {false};
  bool AnyWritten = false;
  for (auto &MBB : MF) {
    for (auto &MI : MBB) {
      // physreg-DEF form: $imgN appearing as a Def operand.
      for (const auto &MO : MI.operands()) {
        if (!MO.isReg() || MO.getReg() == 0 || !MO.isDef()) continue;
        int idx = classifyImgReg(MO.getReg());
        if (idx >= 0) {
          UsedSlot[idx] = true;
          AnyWritten = true;
        }
      }
      // DP-imm form: STA_DP / INC_DP / etc. write the slot at $Cx.
      auto [idx, mode] = classifyDpImmAsImg(MI);
      if (idx >= 0 &&
          (mode == DpAccess::Write || mode == DpAccess::ReadWrite)) {
        UsedSlot[idx] = true;
        AnyWritten = true;
      }
    }
  }
  if (!AnyWritten) return false;

  // Step 2: allocate one frame slot per used IMG.  Size = 2 bytes (each
  // Img16 holds a 16-bit value).  Mark as a spill slot so PEI accounts
  // for it; isSpillSlot=true means slot coloring CAN coalesce it with
  // other spill slots — but the STAfi/LDAfi we emit reference this slot
  // by FrameIndex, and the only writes to this FI are our save/restore
  // pair, so coloring can't break the round-trip.
  MachineFrameInfo &MFI = MF.getFrameInfo();
  int FrameSlots[8];
  for (int i = 0; i < 8; ++i) {
    FrameSlots[i] = -1;
    if (UsedSlot[i])
      FrameSlots[i] = MFI.CreateStackObject(2, Align(2),
                                             /*isSpillSlot=*/true);
  }

  const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
  DebugLoc DL;

  // Step 3: emit prologue save.  Insert at entry MBB.begin() so PEI's
  // emitPrologue (which inserts BEFORE existing MBB.begin) places its
  // frame setup BEFORE our saves — the right order, since our saves
  // reference frame slots whose disps require post-TCS S.
  //
  // Single PHA/PLA bracket around ALL slot saves (vs per-slot bracket).
  // For N used slots:
  //   per-slot:   N * (PHA + LDA dp + STA d,s + PLA) = 16N cyc, 6N bytes
  //   single:     PHA + N*(LDA dp + STA d,s) + PLA   = 8+8N cyc, 2+4N bytes
  // Saves 8 cyc + 2 bytes per additional slot beyond the first.
  //
  // The +2 ImmOffset on STAfi compensates for PHA's SP shift; same +2
  // applies to every slot inside the bracket since SP is constant
  // throughout.
  MachineBasicBlock &EntryMBB = MF.front();
  MachineBasicBlock::iterator EntryIt = EntryMBB.begin();
  BuildMI(EntryMBB, EntryIt, DL, TII->get(W65816::PHA));
  for (int i = 0; i < 8; ++i) {
    if (!UsedSlot[i]) continue;
    BuildMI(EntryMBB, EntryIt, DL, TII->get(W65816::LDA_DP))
        .addImm(IMG_DP[i])
        .addReg(W65816::A, RegState::ImplicitDefine);
    BuildMI(EntryMBB, EntryIt, DL, TII->get(W65816::STAfi))
        .addReg(W65816::A)
        .addFrameIndex(FrameSlots[i])
        .addImm(2)
        .addReg(W65816::A, RegState::ImplicitDefine);
  }
  BuildMI(EntryMBB, EntryIt, DL, TII->get(W65816::PLA))
      .addReg(W65816::A, RegState::ImplicitDefine);

  // Step 4: emit epilogue restore at each return MBB, just BEFORE the
  // RTL/RTS/RTI (so the IMG restore happens before the frame teardown
  // — wait, no: we want the IMG restore BEFORE we tear down the frame
  // because our STAfi/LDAfi reference frame slots).  Insert just
  // before the LAST terminator (the return).  PEI's emitEpilogue will
  // insert its frame-teardown AFTER our restores (BEFORE the return),
  // which means our `,s` disps see the post-TCS S still.
  //
  // Pattern per slot (preserving A, which may hold the return value):
  //   PHA              (preserve A; SP shifts)
  //   LDAfi A, <slot>  (A = saved IMGn value via `lda <disp+2>,s`)
  //   STA <imgN dp>    (write back to IMGn)
  //   PLA              (restore A)
  for (auto &MBB : MF) {
    if (MBB.empty()) continue;
    auto LastIt = std::prev(MBB.end());
    while (LastIt != MBB.begin() && LastIt->isDebugInstr())
      --LastIt;
    unsigned LastOpc = LastIt->getOpcode();
    if (LastOpc != W65816::RTL && LastOpc != W65816::RTS &&
        LastOpc != W65816::RTI)
      continue;

    // Single PHA/PLA bracket for all restores (same optimization as save).
    BuildMI(MBB, LastIt, DL, TII->get(W65816::PHA));
    for (int i = 7; i >= 0; --i) {
      if (!UsedSlot[i]) continue;
      BuildMI(MBB, LastIt, DL, TII->get(W65816::LDAfi))
          .addReg(W65816::A, RegState::Define)
          .addFrameIndex(FrameSlots[i])
          .addImm(2);
      BuildMI(MBB, LastIt, DL, TII->get(W65816::STA_DP))
          .addImm(IMG_DP[i])
          .addReg(W65816::A, RegState::Implicit);
    }
    BuildMI(MBB, LastIt, DL, TII->get(W65816::PLA))
        .addReg(W65816::A, RegState::ImplicitDefine);
  }

  return true;
}
