//===-- W65816SpillToX.cpp - Replace stack spills with TAX/TXA -----------===//
//
// 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 peephole: replace stack-spill/reload pairs with TAX/TXA (or
// TAY/TYA) when the index register is dead during the spill window.
//
// Fast regalloc spills A to stack via STAfi/LDAfi, costing ~12 cycles
// per round-trip (sta is 5 cycles + lda is 5 cycles + the displacement
// dispatch).  But the W65816 has TAX (2 cycles) + TXA (2 cycles), a
// 3x speedup if X is free during the spill window.
//
// We scan each basic block for the pattern:
//
//     STAfi $a, slot, 0
//     ...   (instructions that don't touch X or A's slot, don't kill A)
//     LDAfi $a, slot, 0
//
// If no instruction in the gap reads or writes X (or P-flags-dependent
// X side effects, etc.), we rewrite the pair as:
//
//     TAX
//     ...
//     TXA
//
// This saves 4 bytes (stack-rel addressing is 2 bytes per op vs TAX/TXA
// at 1 byte each) AND saves the memory traffic.  Net: ~8 cycles per
// converted pair.
//
// Conservative liveness: we treat X as "in use" if ANY instruction in
// the gap references W65816::X (def or use).  False positives mean
// we keep the slow stack form; false negatives are correctness bugs.
//
//===----------------------------------------------------------------------===//

#include "W65816.h"
#include "W65816InstrInfo.h"
#include "W65816Subtarget.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/CodeGen/MachineFrameInfo.h"
#include "llvm/CodeGen/MachineFunction.h"
#include "llvm/CodeGen/MachineFunctionPass.h"
#include "llvm/CodeGen/MachineInstr.h"
#include "llvm/CodeGen/MachineInstrBuilder.h"
#include "llvm/CodeGen/TargetRegisterInfo.h"

using namespace llvm;

#define DEBUG_TYPE "w65816-spill-to-x"

namespace {

class W65816SpillToX : public MachineFunctionPass {
public:
  static char ID;
  W65816SpillToX() : MachineFunctionPass(ID) {}
  StringRef getPassName() const override {
    return "W65816 spill-to-X peephole";
  }

  void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.setPreservesCFG();
    MachineFunctionPass::getAnalysisUsage(AU);
  }

  bool runOnMachineFunction(MachineFunction &MF) override;
};

} // namespace

char W65816SpillToX::ID = 0;

INITIALIZE_PASS(W65816SpillToX, DEBUG_TYPE, "W65816 spill-to-X peephole",
                false, false)

FunctionPass *llvm::createW65816SpillToX() {
  return new W65816SpillToX();
}

// Classifies how an MI interacts with X.
enum XEffect { XNone = 0, XReads = 1, XDefs = 2, XBoth = 3 };

// Most W65816 transfer/index opcodes (TAX, INX, LDX, STX, CPX, etc.)
// are tablegen'd as `InstImplied` with no Defs/Uses metadata, so the
// MCInstrDesc carries no implicit X operand and a generic operand
// scan misses them.  We hard-code the X-effect per opcode instead.
// Calls clobber X under our caller-saved-X ABI.
static XEffect xEffect(const MachineInstr &MI, const TargetRegisterInfo *TRI) {
  switch (MI.getOpcode()) {
  case W65816::TAX:                        // X := A
  case W65816::TYX:                        // X := Y
  case W65816::TSX:                        // X := SP
  case W65816::PLX:                        // X := pop
    return XDefs;
  case W65816::TXA:                        // A := X
  case W65816::TXY:                        // Y := X
  case W65816::TXS:                        // SP := X
  case W65816::PHX:                        // push X
    return XReads;
  case W65816::INX:                        // X := X+1
  case W65816::DEX:                        // X := X-1
    return XBoth;
  default:
    break;
  }
  if (MI.isCall()) return XBoth;           // caller-clobbered X
  // Generic operand scan for opcodes that carry X explicitly (LDX/STX/CPX
  // pseudos) or any properly-modelled implicit defs/uses.
  int eff = XNone;
  for (const auto &MO : MI.operands()) {
    if (!MO.isReg()) continue;
    Register R = MO.getReg();
    if (!R.isPhysical()) continue;
    bool isX = R == W65816::X || (TRI && TRI->regsOverlap(R, W65816::X));
    if (!isX) continue;
    if (MO.isDef()) eff |= XDefs; else eff |= XReads;
  }
  return (XEffect)eff;
}

// Convenience wrapper: returns true if MI references X in any way.
static bool touchesX(const MachineInstr &MI, const TargetRegisterInfo *TRI) {
  return xEffect(MI, TRI) != XNone;
}

// Returns true if MI is `STAfi $a, slot, 0`.
static int matchSTAfi(const MachineInstr &MI) {
  if (MI.getOpcode() != W65816::STAfi) return -1;
  if (MI.getNumOperands() < 3) return -1;
  if (!MI.getOperand(0).isReg() || MI.getOperand(0).getReg() != W65816::A)
    return -1;
  if (!MI.getOperand(1).isFI()) return -1;
  if (!MI.getOperand(2).isImm() || MI.getOperand(2).getImm() != 0) return -1;
  return MI.getOperand(1).getIndex();
}

// Returns FI if MI is `LDAfi slot, 0` defining $a, else -1.
static int matchLDAfi(const MachineInstr &MI) {
  if (MI.getOpcode() != W65816::LDAfi) return -1;
  if (MI.getNumOperands() < 3) return -1;
  if (!MI.getOperand(0).isReg() || MI.getOperand(0).getReg() != W65816::A)
    return -1;
  if (!MI.getOperand(1).isFI()) return -1;
  if (!MI.getOperand(2).isImm() || MI.getOperand(2).getImm() != 0) return -1;
  return MI.getOperand(1).getIndex();
}

// Returns true if MI reads or writes the slot at FrameIndex FI.
static bool referencesSlot(const MachineInstr &MI, int FI) {
  for (const auto &MO : MI.operands()) {
    if (MO.isFI() && MO.getIndex() == FI) return true;
  }
  return false;
}

bool W65816SpillToX::runOnMachineFunction(MachineFunction &MF) {
  const W65816Subtarget &STI = MF.getSubtarget<W65816Subtarget>();
  const W65816InstrInfo *TII = STI.getInstrInfo();
  const TargetRegisterInfo *TRI = STI.getRegisterInfo();
  bool Changed = false;
  // Slots whose last reference we erased — candidates for reclamation.
  SmallSet<int, 8> SlotsTouched;

  for (auto &MBB : MF) {
    // Pass 1: collect (STAfi, slot) entries.
    SmallVector<std::pair<MachineInstr *, int>, 8> Stas;
    for (auto &MI : MBB) {
      int FI = matchSTAfi(MI);
      if (FI != -1) Stas.push_back({&MI, FI});
    }

    // For each STAfi, scan forward for the matching LDAfi with no
    // intervening X touch or slot reference.  Process in REVERSE
    // order so any nested pair is converted first; the outer pair's
    // gap scan then sees the inner TAX/TXA (which touches X) and
    // bails — preventing a mid-bridge X clobber.
    for (auto It = Stas.rbegin(); It != Stas.rend(); ++It) {
      auto [StaMI, FI] = *It;
      bool xTouched = false;
      bool gapEmpty = true;
      MachineInstr *LdaMI = nullptr;
      for (auto Scan = std::next(MachineBasicBlock::iterator(StaMI));
           Scan != MBB.end(); ++Scan) {
        MachineInstr &MI2 = *Scan;
        if (MI2.isDebugInstr()) continue;

        // Look for the matching LDAfi.  TAX preserves A so we don't
        // need to check A liveness — only whether X was free.
        if (matchLDAfi(MI2) == FI) { LdaMI = &MI2; break; }

        // Bail if X is touched (use or def, including implicit on
        // calls) or if the slot is referenced by something else
        // (which would invalidate the saved value).
        if (touchesX(MI2, TRI)) { xTouched = true; break; }
        if (referencesSlot(MI2, FI)) break;
        gapEmpty = false;
      }

      // Defer empty-gap pairs to StackSlotCleanup, which deletes both
      // (A still holds the stored value across an empty gap).  That
      // beats our TAX+TXA conversion (0 instr vs 2 instr).
      if (!LdaMI || xTouched || gapEmpty) continue;

      // X-live-after-LDA check: TXA (the LDAfi replacement) clobbers X.
      // If anything downstream of the LDA reads X — including the next
      // JSL's implicit $x — then we'd silently corrupt X.  Caught by
      // i32 first-arg functions where $x is live-in (= arg0_hi) and
      // a libcall later in the block expects $x intact.  Scan from just
      // past LDA to end-of-block; if any instr uses X, bail.
      bool xUsedAfter = false;
      for (auto Scan = std::next(MachineBasicBlock::iterator(LdaMI));
           Scan != MBB.end(); ++Scan) {
        const MachineInstr &MI3 = *Scan;
        if (MI3.isDebugInstr()) continue;
        XEffect eff = xEffect(MI3, TRI);
        if (eff & XReads) { xUsedAfter = true; break; }
        if (eff & XDefs) break;  // X redefined; no longer live
      }
      // Also bail if X is live-in to MBB and nothing has defined X
      // between MBB start and STA — the live-in value is needed past
      // the LDA point.
      if (!xUsedAfter && MBB.isLiveIn(W65816::X)) {
        bool xRedefBeforeSta = false;
        for (auto Scan = MBB.begin();
             Scan != MachineBasicBlock::iterator(StaMI); ++Scan) {
          const MachineInstr &MI3 = *Scan;
          if (MI3.isDebugInstr()) continue;
          if (xEffect(MI3, TRI) & XDefs) { xRedefBeforeSta = true; break; }
        }
        if (!xRedefBeforeSta) xUsedAfter = true;
      }
      if (xUsedAfter) continue;

      // Cross-block use check: if the slot is referenced anywhere
      // OUTSIDE the [STA, LDA] window (including other blocks), the
      // STA we'd erase is feeding those other reads — eliding it
      // would silently corrupt them.  Caught by sumTable() returning
      // a stale phi value because the loop's STA-to-merge-slot was
      // eliminated; the merge block's LDA then read the bb.0-init 0
      // instead of the loop's accumulated sum.
      bool externalUse = false;
      for (auto &OtherMBB : MF) {
        for (auto &OtherMI : OtherMBB) {
          if (&OtherMI == StaMI || &OtherMI == LdaMI) continue;
          // Walk inside-window range and skip those refs.
          if (&OtherMBB == &MBB) {
            // We already verified the gap doesn't reference FI; only
            // STA/LDA themselves are allowed users in this block.
          }
          if (referencesSlot(OtherMI, FI)) {
            externalUse = true;
            break;
          }
        }
        if (externalUse) break;
      }
      if (externalUse) continue;

      // Replace STAfi with TAX, LDAfi with TXA.
      DebugLoc StaDL = StaMI->getDebugLoc();
      DebugLoc LdaDL = LdaMI->getDebugLoc();
      MachineBasicBlock *MBB2 = StaMI->getParent();
      auto StaIt = MachineBasicBlock::iterator(StaMI);
      auto LdaIt = MachineBasicBlock::iterator(LdaMI);
      BuildMI(*MBB2, StaIt, StaDL, TII->get(W65816::TAX));
      BuildMI(*MBB2, LdaIt, LdaDL, TII->get(W65816::TXA))
          .addReg(W65816::A, RegState::ImplicitDefine);
      StaMI->eraseFromParent();
      LdaMI->eraseFromParent();
      SlotsTouched.insert(FI);
      Changed = true;
    }

    // Post-pass: collapse `TAX ; TXA` (or `TXA ; TAX`) pairs whose
    // observable effect is dead.  These appear when an inner STA/LDA
    // pair (originally between an outer pair we converted) was deleted
    // by StackSlotCleanup or coalesced by stack-slot-coloring, leaving
    // our TAX/TXA bookends adjacent.
    //
    // Distinct effect per ordering:
    //   TAX;TXA  : net effect is `X := A` (A unchanged, X clobbered).
    //              Removable iff X dead afterwards.
    //   TXA;TAX  : net effect is `A := X` (X unchanged, A clobbered).
    //              Removable iff A dead afterwards.
    //
    // The earlier code mis-handled TXA;TAX as if it clobbered X; in
    // fact X comes through the pair unchanged.
    auto It = MBB.begin();
    while (It != MBB.end()) {
      auto Next = std::next(It);
      if (Next == MBB.end()) break;
      bool isTaxThenTxa = It->getOpcode() == W65816::TAX &&
                          Next->getOpcode() == W65816::TXA;
      bool isTxaThenTax = It->getOpcode() == W65816::TXA &&
                          Next->getOpcode() == W65816::TAX;
      if (!isTaxThenTxa && !isTxaThenTax) { ++It; continue; }

      // Choose which physreg's liveness matters based on which value
      // the pair clobbers.
      Register Clobbered = isTaxThenTxa ? W65816::X : W65816::A;

      bool observed = false;
      bool killedByDef = false;
      for (auto Tail = std::next(Next); Tail != MBB.end(); ++Tail) {
        if (Tail->isDebugInstr()) continue;
        if (Tail->readsRegister(Clobbered, TRI)) { observed = true; break; }
        // Calls clobber both A and X (caller-saved).
        if (Tail->isCall()) { killedByDef = true; break; }
        // Opcode-based defs (TAX/TXA tablegen has no Defs metadata).
        if (Clobbered == W65816::X) {
          XEffect E = xEffect(*Tail, TRI);
          if (E & XReads) { observed = true; break; }
          if (E & XDefs)  { killedByDef = true; break; }
        } else {
          // For A: any LDA*/PLA/TXA/TYA/INA/DEA/arith op redefines A.
          unsigned Op = Tail->getOpcode();
          if (Op == W65816::TXA || Op == W65816::TYA ||
              Op == W65816::INA || Op == W65816::DEA ||
              Op == W65816::PLA) { killedByDef = true; break; }
          if (Tail->modifiesRegister(W65816::A, TRI)) {
            killedByDef = true; break;
          }
        }
      }
      if (observed) { ++It; continue; }
      if (!killedByDef) {
        bool liveOut = false;
        for (MachineBasicBlock *Succ : MBB.successors()) {
          if (Succ->isLiveIn(Clobbered)) { liveOut = true; break; }
        }
        if (liveOut) { ++It; continue; }
      }

      auto Erase1 = It++;
      auto Erase2 = It++;
      Erase1->eraseFromParent();
      Erase2->eraseFromParent();
      Changed = true;
    }
  }

  // Reclaim frame slots whose last reference we just erased.  Without
  // this, PEI still allocates space for them and emits the prologue
  // PHA, even though the slot is unused — wastes 1 PHA (4 cyc) and
  // 1 PLY per call.  RemoveStackObject marks the slot dead by setting
  // its size to ~0ULL; PEI ignores those when computing frame size.
  if (!SlotsTouched.empty()) {
    MachineFrameInfo &MFI = MF.getFrameInfo();
    for (int FI : SlotsTouched) {
      bool stillUsed = false;
      for (auto &MBB : MF) {
        for (auto &MI : MBB) {
          if (referencesSlot(MI, FI)) { stillUsed = true; break; }
        }
        if (stillUsed) break;
      }
      if (!stillUsed) MFI.RemoveStackObject(FI);
    }
  }
  return Changed;
}