//===-- W65816NegYIndY.cpp - Fix negative-Y indirect addressing -----------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
//
//===----------------------------------------------------------------------===//
//
// Pre-emit peephole that rewrites
//
//   LDY #imm        ; imm signed-negative (>= 0x8000 unsigned)
//   LDA (sr,S),Y    ; or STA
//
// into
//
//   LDA sr,S        ; A = ptr
//   CLC ; ADC #imm  ; A = ptr + imm  (signed add wraps within 16 bits in A)
//   TAX             ; X = adjusted ptr
//   ; for LDA path: LDA $0000,X    ; A = DBR:X
//   ; for STA path: TAY (save A) ; ... ; TYA before STA $0000,X
//
// Why: the WDC W65816 spec says (sr,S),Y computes
//
//     EA = (DBR | (mem16(sr+S) + Y)) MOD $1000000
//
// — a 24-bit add.  When Y is signed-negative (e.g. $FFFE for "-2"), the
// addition crosses bank boundaries: ptr=$5DB3 + $FFFE = $015DB1, NOT
// $005DB1.  Caught by `arr[-1]` and bubble-sort swaps with `arr[j-1]`.
//
// Using `abs,X` with operand $0000 and X = adjusted-ptr avoids the
// problem because X is < 16 bits and operand + X stays within DBR
// when the operand is small.
//
//===----------------------------------------------------------------------===//

#include "W65816.h"
#include "W65816InstrInfo.h"
#include "W65816Subtarget.h"
#include "llvm/CodeGen/MachineFunction.h"
#include "llvm/CodeGen/MachineFunctionPass.h"
#include "llvm/CodeGen/MachineInstr.h"
#include "llvm/CodeGen/MachineInstrBuilder.h"

using namespace llvm;

// DP scratch byte used to park X when the negative-Y inserter needs to
// route through TAX/TXA.  Lives in the project-wide $E0..$DF scratch
// range; $E0 is reserved for ADJCALLSTACKUP's A-preserve so we use
// $E2 here.  Coordinate with W65816AsmPrinter / W65816ISelLowering /
// W65816RegisterInfo if the layout changes.
static constexpr unsigned kDpScratchX = 0xE2;

#define DEBUG_TYPE "w65816-neg-y-indy"

namespace {

class W65816NegYIndY : public MachineFunctionPass {
public:
  static char ID;
  W65816NegYIndY() : MachineFunctionPass(ID) {}
  StringRef getPassName() const override {
    return "W65816 negative-Y indirect-Y rewriter";
  }
  bool runOnMachineFunction(MachineFunction &MF) override;
};

} // namespace

char W65816NegYIndY::ID = 0;

INITIALIZE_PASS(W65816NegYIndY, DEBUG_TYPE,
                "W65816 negative-Y indirect-Y rewriter", false, false)

FunctionPass *llvm::createW65816NegYIndY() { return new W65816NegYIndY(); }

bool W65816NegYIndY::runOnMachineFunction(MachineFunction &MF) {
  const W65816InstrInfo *TII =
      MF.getSubtarget<W65816Subtarget>().getInstrInfo();
  bool Changed = false;
  for (MachineBasicBlock &MBB : MF) {
    int LastY = -1;
    MachineInstr *LastLDY = nullptr;
    for (auto It = MBB.begin(), End = MBB.end(); It != End; ) {
      MachineInstr &MI = *It++;
      if (MI.isDebugInstr()) continue;
      unsigned Opc = MI.getOpcode();
      if (Opc == W65816::LDY_Imm16 && MI.getNumOperands() >= 1 &&
          MI.getOperand(0).isImm()) {
        LastY = (int)(MI.getOperand(0).getImm() & 0xFFFF);
        LastLDY = &MI;
        continue;
      }
      bool IsLDA = Opc == W65816::LDA_StackRelIndY;
      bool IsSTA = Opc == W65816::STA_StackRelIndY;
      if ((IsLDA || IsSTA) && LastY != -1 && (LastY & 0x8000)) {
        // Negative Y.  Rewrite via TAX + LDA/STA $0000,X.
        if (MI.getNumOperands() < 1 || !MI.getOperand(0).isImm())
          continue;
        unsigned Disp = MI.getOperand(0).getImm() & 0xFF;
        DebugLoc DL = MI.getDebugLoc();
        // X-liveness check: SpillToX may have stashed a value in X
        // that's used after this rewrite.  If so, save X to DP $E1
        // (libcall scratch high half — $E0 is reserved for the A-save
        // dance in eliminateCallFramePseudoInstr) and restore after.
        // Walk forward from MI looking for an X use without a prior
        // X def; if found, X is live and we must preserve it.
        bool XLive = false;
        for (auto Scan = std::next(MachineBasicBlock::iterator(&MI));
             Scan != MBB.end(); ++Scan) {
          if (Scan->isDebugInstr()) continue;
          bool xDef = false;
          for (const MachineOperand &MO : Scan->operands()) {
            if (!MO.isReg()) continue;
            if (MO.getReg() == W65816::X) {
              if (MO.isUse()) { XLive = true; break; }
              if (MO.isDef()) xDef = true;
            }
          }
          if (XLive || xDef) break;
        }
        if (XLive) {
          // Save X to DP kDpScratchX ($E2) -- $E0 is reserved as the
          // A-preserve slot in call-frame teardown and may be live.
          BuildMI(MBB, MI, DL, TII->get(W65816::STX_DP)).addImm(kDpScratchX);
        }
        if (IsLDA) {
          // LDA disp,S ; CLC ; ADC #neg ; TAX ; LDA $0000,X
          BuildMI(MBB, MI, DL, TII->get(W65816::LDA_StackRel))
              .addImm(Disp)
              .addReg(W65816::A, RegState::ImplicitDefine);
          BuildMI(MBB, MI, DL, TII->get(W65816::CLC))
              .addReg(W65816::P, RegState::ImplicitDefine);
          BuildMI(MBB, MI, DL, TII->get(W65816::ADC_Imm16))
              .addImm(LastY)
              .addReg(W65816::A, RegState::Implicit)
              .addReg(W65816::A, RegState::ImplicitDefine)
              .addReg(W65816::P, RegState::Implicit)
              .addReg(W65816::P, RegState::ImplicitDefine);
          BuildMI(MBB, MI, DL, TII->get(W65816::TAX));
          BuildMI(MBB, MI, DL, TII->get(W65816::LDA_AbsX))
              .addImm(0)
              .addReg(W65816::A, RegState::ImplicitDefine);
        } else { // STA
          // A holds the value to store.  TAY (save A in Y) ;
          // LDA disp,S ; CLC ; ADC #neg ; TAX ; TYA ; STA $0000,X
          BuildMI(MBB, MI, DL, TII->get(W65816::TAY));
          BuildMI(MBB, MI, DL, TII->get(W65816::LDA_StackRel))
              .addImm(Disp)
              .addReg(W65816::A, RegState::ImplicitDefine);
          BuildMI(MBB, MI, DL, TII->get(W65816::CLC))
              .addReg(W65816::P, RegState::ImplicitDefine);
          BuildMI(MBB, MI, DL, TII->get(W65816::ADC_Imm16))
              .addImm(LastY)
              .addReg(W65816::A, RegState::Implicit)
              .addReg(W65816::A, RegState::ImplicitDefine)
              .addReg(W65816::P, RegState::Implicit)
              .addReg(W65816::P, RegState::ImplicitDefine);
          BuildMI(MBB, MI, DL, TII->get(W65816::TAX));
          BuildMI(MBB, MI, DL, TII->get(W65816::TYA));
          BuildMI(MBB, MI, DL, TII->get(W65816::STA_AbsX))
              .addImm(0)
              .addReg(W65816::A, RegState::Implicit);
        }
        if (XLive) {
          // Restore X from DP $E2.
          BuildMI(MBB, MI, DL, TII->get(W65816::LDX_DP)).addImm(kDpScratchX);
        }
        // Erase original LDY and the (sr,s),Y op.
        if (LastLDY) { LastLDY->eraseFromParent(); LastLDY = nullptr; }
        MI.eraseFromParent();
        LastY = -1;
        Changed = true;
        continue;
      }
      switch (Opc) {
      case W65816::TAY: case W65816::TXY:
      case W65816::INY: case W65816::DEY:
      case W65816::PLY: case W65816::LDY_DP: case W65816::LDY_Abs:
      case W65816::LDY_DPX: case W65816::LDY_AbsX:
        LastY = -1;
        LastLDY = nullptr;
        break;
      default:
        if (MI.isCall()) { LastY = -1; LastLDY = nullptr; }
        break;
      }
    }
  }
  return Changed;
}
