//===-- W65816ABridgeViaX.cpp - Pre-RA bridge of Acc16 vregs via X -------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// Pre-regalloc complement to W65816TiedDefSpill.  Where TiedDefSpill
// preserves a multi-use Acc16 vreg by spilling it to a fresh stack
// slot around the tied-def consumer, this pass bridges via an Img16
// (DP-backed) vreg: park SrcReg in a fresh Img16 vreg before the
// consumer, restore to a fresh Acc16 vreg after.  Regalloc places the
// Img16 in IMG0..IMG7 (DP $D0..$DE); copyPhysReg lowers the COPYs to
// STA dp / LDA dp (4 cyc each) and no system-stack slot is allocated.
//
// (The pass name dates from an earlier prototype that bridged via X
// using TAX/TXA.  Cross-MBB X-liveness analysis was unimplemented and
// the X-bridge couldn't survive Idx16 clobbers between consumer and
// last use, so the bridge moved to Img16.  The DP-backed form has the
// same 4-cycle round-trip cost as TAX/TXA bridges with none of the
// liveness restrictions.)
//
// Win per bridged pair:
//   stack spill: STA dp,S (5 cyc) + LDA dp,S (5 cyc) + 1 frame slot
//   Img bridge : STA dp   (4 cyc) + LDA dp   (4 cyc) + no frame growth
// Net 2 cycles + (1 byte per access) saved per bridge -- and one PHA
// per avoided stack slot.
//
// Bail conditions (fall back to TiedDefSpill's stack route): any MI
// between consumer and SrcReg's last use that clobbers IMG slots,
// callees that clobber IMG0..IMG7, cross-MBB uses of SrcReg.
//
// Runs before TiedDefSpill so the latter doesn't double-process the
// same candidates.
//
//===----------------------------------------------------------------------===//

#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"
#include "llvm/CodeGen/MachineRegisterInfo.h"

using namespace llvm;

#define DEBUG_TYPE "w65816-a-bridge-via-x"

namespace {

class W65816ABridgeViaX : public MachineFunctionPass {
public:
  static char ID;
  W65816ABridgeViaX() : MachineFunctionPass(ID) {}
  StringRef getPassName() const override {
    return "W65816 Acc16 bridge via X";
  }
  void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.setPreservesCFG();
    MachineFunctionPass::getAnalysisUsage(AU);
  }
  bool runOnMachineFunction(MachineFunction &MF) override;
};

} // namespace

char W65816ABridgeViaX::ID = 0;

INITIALIZE_PASS(W65816ABridgeViaX, DEBUG_TYPE,
                "W65816 Acc16 bridge via X", false, false)

FunctionPass *llvm::createW65816ABridgeViaX() {
  return new W65816ABridgeViaX();
}

// Same predicate as TiedDefSpill via the shared helper.

// Pre-RA check for "instruction may clobber an Img16 (DP $D0..$DF)
// register."  Calls clobber them caller-save.  Any other DP load/store
// to that range would too — but we don't currently have non-libcall
// emitters into $D0..$DF, so the call check covers it.  Conservative
// extras: anything that could touch DP overall is excluded.
static bool clobbersImg(const MachineInstr &MI,
                        const MachineRegisterInfo &MRI) {
  if (MI.isCall()) return true;
  // Bail on any MI that defs an Img16 or its DP physreg — none should
  // exist before our pass runs, but cover the case for robustness.
  for (const auto &MO : MI.operands()) {
    if (!MO.isReg() || !MO.isDef()) continue;
    Register R = MO.getReg();
    if (!R.isValid()) continue;
    if (R.isPhysical()) {
      if (R == W65816::IMG0  || R == W65816::IMG1  || R == W65816::IMG2  ||
          R == W65816::IMG3  || R == W65816::IMG4  || R == W65816::IMG5  ||
          R == W65816::IMG6  || R == W65816::IMG7  ||
          R == W65816::IMG8  || R == W65816::IMG9  || R == W65816::IMG10 ||
          R == W65816::IMG11 || R == W65816::IMG12 || R == W65816::IMG13 ||
          R == W65816::IMG14 || R == W65816::IMG15)
        return true;
      continue;
    }
    const TargetRegisterClass *RC = MRI.getRegClass(R);
    if (RC == &W65816::Img16RegClass) return true;
  }
  return false;
}

bool W65816ABridgeViaX::runOnMachineFunction(MachineFunction &MF) {
  if (!MF.getRegInfo().getNumVirtRegs()) return false;
  if (MF.getFunction().hasOptNone()) return false;
  MachineRegisterInfo &MRI = MF.getRegInfo();
  const W65816Subtarget &STI = MF.getSubtarget<W65816Subtarget>();
  const W65816InstrInfo *TII = STI.getInstrInfo();
  bool Changed = false;

  // Snapshot candidates before mutating MIR.
  struct Candidate {
    MachineBasicBlock *MBB;
    MachineInstr *MI;
    unsigned OpIdx;
  };
  SmallVector<Candidate, 8> Candidates;

  for (auto &MBB : MF) {
    for (auto &MI : MBB) {
      if (!W65816Helpers::hasTiedAcc16Src(MI)) continue;
      for (unsigned i = 0; i < MI.getNumOperands(); ++i) {
        const MachineOperand &MO = MI.getOperand(i);
        if (!MO.isReg() || !MO.isUse()) continue;
        if (!MI.isRegTiedToDefOperand(i)) continue;
        Register R = MO.getReg();
        if (!R.isVirtual()) continue;
        if (MRI.getRegClass(R) != &W65816::Acc16RegClass) continue;

        // Mirror TiedDefSpill's "needs spill" criterion exactly:
        // SrcReg has a post-consumer COPY to a physreg.
        bool needSpill = false;
        bool badUse = false;
        for (auto &U : MRI.use_nodbg_instructions(R)) {
          if (&U == &MI) continue;
          if (U.isPHI()) { badUse = true; break; }
          if (U.isCopy()) {
            const MachineOperand &Dst = U.getOperand(0);
            if (Dst.isReg() && Dst.getReg().isPhysical()) {
              needSpill = true;
              continue;
            }
          }
        }
        if (needSpill && !badUse) {
          Candidates.push_back({&MBB, &MI, i});
        }
      }
    }
  }

  for (auto C : Candidates) {
    MachineInstr *MI = C.MI;
    MachineBasicBlock *MBB = C.MBB;
    unsigned OpIdx = C.OpIdx;
    Register SrcReg = MI->getOperand(OpIdx).getReg();
    if (!SrcReg.isVirtual()) continue;
    if (MRI.getRegClass(SrcReg) != &W65816::Acc16RegClass) continue;

    // Determine the post-consumer-use range in MI's MBB.  All uses
    // outside MBB disqualify (cross-MBB X/Y liveness too complex
    // for first cut — fall through to TiedDefSpill).
    bool sameMBBOnly = true;
    auto LastUseIt = MBB->end();
    for (auto &U : MRI.use_nodbg_instructions(SrcReg)) {
      if (&U == MI) continue;
      if (U.getParent() != MBB) { sameMBBOnly = false; break; }
      // Track latest use (in MBB order).
      auto It = MachineBasicBlock::iterator(&U);
      bool afterMI = false;
      for (auto Walk = MachineBasicBlock::iterator(MI), End = MBB->end();
           Walk != End; ++Walk) {
        if (Walk == It) { afterMI = true; break; }
      }
      if (!afterMI) continue;  // pre-consumer use stays on SrcReg
      // Pick the latest such It as LastUseIt.
      bool isLater = (LastUseIt == MBB->end());
      if (!isLater) {
        for (auto Walk = std::next(It); Walk != MBB->end(); ++Walk) {
          if (Walk == LastUseIt) { isLater = true; break; }
        }
      }
      if (isLater) LastUseIt = It;
    }
    if (!sameMBBOnly || LastUseIt == MBB->end()) continue;

    // Scan from just after MI to LastUseIt: bail if anything could
    // clobber an IMGn (calls and other DP-touchers).
    bool imgClobbered = false;
    for (auto It = std::next(MachineBasicBlock::iterator(MI));
         It != LastUseIt; ++It) {
      if (It->isDebugInstr()) continue;
      if (clobbersImg(*It, MRI)) { imgClobbered = true; break; }
    }
    if (imgClobbered) continue;

    // Bridge.  Park SrcReg in an Img16 (DP-backed) vreg around the
    // consumer; restore via COPY back to a fresh Acc16 vreg afterward.
    // Regalloc allocates the Img16 vreg to one of IMG0..IMG7 (DP slots
    // $D0..$DE).  copyPhysReg lowers the COPYs to STA dp / LDA dp
    // (4 cyc each); spills don't touch the system stack at all.
    DebugLoc DL = MI->getDebugLoc();
    Register ImgReg = MRI.createVirtualRegister(&W65816::Img16RegClass);
    BuildMI(*MBB, MI, DL, TII->get(TargetOpcode::COPY), ImgReg)
        .addReg(SrcReg);
    Register NewReg = MRI.createVirtualRegister(&W65816::Acc16RegClass);
    auto AfterMI = std::next(MachineBasicBlock::iterator(MI));
    BuildMI(*MBB, AfterMI, DL, TII->get(TargetOpcode::COPY), NewReg)
        .addReg(ImgReg);

    // Rewrite uses of SrcReg that come AFTER MI in the same MBB.
    SmallVector<MachineOperand *, 4> ToRewrite;
    for (auto &U : MRI.use_nodbg_operands(SrcReg)) {
      if (U.getParent() == MI) continue;
      MachineBasicBlock *UseMBB = U.getParent()->getParent();
      if (UseMBB != MBB) continue;
      bool After = false;
      for (auto Walk = MachineBasicBlock::iterator(MI),
                End = MBB->end(); Walk != End; ++Walk) {
        if (&*Walk == U.getParent()) { After = true; break; }
      }
      if (After) ToRewrite.push_back(&U);
    }
    for (auto *MO : ToRewrite) {
      MO->setReg(NewReg);
      MO->setIsKill(false);
    }
    Changed = true;
  }

  return Changed;
}
