//===-- W65816PreSpillCrossCall.cpp - Pre-spill cross-call Acc16 vregs ---===//
//
// 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-RA pass that pre-spills Acc16 vregs whose live range crosses a
// JSL call site.  Greedy regalloc has only one register in the Acc16
// class (A), and JSL clobbers A — so any Acc16 vreg live across a
// call MUST be spilled.  Greedy normally figures this out, but for
// functions with many such vregs (the "ok |= bit" bitmask pattern
// repeated across N if-blocks each calling a helper) greedy can run
// out of registers during spill placement, aborting compilation with
// "ran out of registers".
//
// We pre-empt the failure: walk the MBB, find cross-call Acc16
// vregs, and explicitly STAfi their value after the def + LDAfi at
// each use.  This converts the cross-call live ranges into stack-
// resident loads, dropping greedy's pressure to the point it can
// always succeed.
//
// Cost: an extra STAfi+LDAfi (~6 cyc each) per cross-call vreg.
// This is the same cost greedy would emit if it succeeded (a spill
// + reload), so we're not pessimising — just making the spill
// explicit BEFORE greedy gets confused.
//
// Heuristic: only pre-spill if the function has > 5 call sites OR a
// cross-call Acc16 vreg with > 2 uses after the call.  Below that,
// let greedy do its thing (it usually picks better placements).
//
//===----------------------------------------------------------------------===//

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

using namespace llvm;

#define DEBUG_TYPE "w65816-pre-spill-cross-call"

namespace {

class W65816PreSpillCrossCall : public MachineFunctionPass {
public:
  static char ID;
  W65816PreSpillCrossCall() : MachineFunctionPass(ID) {}

  StringRef getPassName() const override {
    return "W65816 pre-spill Acc16 vregs across calls";
  }

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

  bool runOnMachineFunction(MachineFunction &MF) override;
};

} // namespace

char W65816PreSpillCrossCall::ID = 0;

INITIALIZE_PASS(W65816PreSpillCrossCall, DEBUG_TYPE,
                "W65816 pre-spill Acc16 vregs across calls", false, false)

FunctionPass *llvm::createW65816PreSpillCrossCall() {
  return new W65816PreSpillCrossCall();
}

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

  // First pass: count call sites in the function.  Below the
  // heuristic threshold we don't bother — greedy handles low-call
  // functions fine and pre-spilling would just add bytes.
  unsigned callCount = 0;
  for (MachineBasicBlock &MBB : MF)
    for (MachineInstr &MI : MBB)
      if (MI.isCall()) callCount++;
  if (callCount < 4) return false;

  bool Changed = false;

  // Walk every Acc16 vreg in the function.  For each, find its def
  // (allowing multi-def vregs like SELECT_CC results — pick the
  // first by MachineInstr iteration), then check if any use is
  // separated from the def by a JSL call (in the same MBB).  If
  // so, pre-spill via STAfi at def + LDAfi at each post-call use.
  unsigned NumVRegs = MRI.getNumVirtRegs();
  for (unsigned i = 0; i < NumVRegs; ++i) {
    Register VReg = Register::index2VirtReg(i);
    if (MRI.def_empty(VReg)) continue;
    if (MRI.getRegClass(VReg) != &W65816::Acc16RegClass) continue;
    // Find the first def.  For PHIs we skip — pre-spilling a PHI
    // result is complex and rarely helpful for the high-pressure
    // pattern we target (which is sequential bitmask updates).
    MachineInstr *DefMI = nullptr;
    for (MachineInstr &D : MRI.def_instructions(VReg)) {
      if (D.isPHI()) { DefMI = nullptr; break; }
      if (!DefMI) DefMI = &D;
    }
    if (!DefMI) continue;
    MachineBasicBlock *MBB = DefMI->getParent();

    // Check if any use of VReg is in the same MBB AFTER a call
    // following DefMI.
    bool sawCallAfterDef = false;
    SmallVector<MachineInstr *, 4> postCallUses;
    auto Walker = std::next(DefMI->getIterator());
    while (Walker != MBB->end()) {
      MachineInstr &W = *Walker++;
      if (W.isCall()) sawCallAfterDef = true;
      if (sawCallAfterDef && W.readsRegister(VReg, /*TRI=*/nullptr))
        postCallUses.push_back(&W);
    }
    if (postCallUses.empty()) continue;

    // Pre-spill.  Fresh slot per vreg — StackSlotColoring may merge
    // slots later if their lifetimes don't overlap.
    int FI = MFI.CreateStackObject(2, Align(2), /*isSpillSlot=*/true);
    DebugLoc DL = DefMI->getDebugLoc();
    auto AfterDef = std::next(DefMI->getIterator());
    BuildMI(*MBB, AfterDef, DL, TII->get(W65816::STAfi))
        .addReg(VReg).addFrameIndex(FI).addImm(0);
    for (MachineInstr *UseMI : postCallUses) {
      Register Reload = MRI.createVirtualRegister(&W65816::Acc16RegClass);
      BuildMI(*UseMI->getParent(), UseMI->getIterator(), UseMI->getDebugLoc(),
              TII->get(W65816::LDAfi), Reload)
          .addFrameIndex(FI).addImm(0);
      // Rewrite this use's references of VReg to Reload.
      for (auto &MO : UseMI->uses()) {
        if (MO.isReg() && MO.getReg() == VReg) {
          MO.setReg(Reload);
          MO.setIsKill(false);
        }
      }
    }
    Changed = true;
  }

  return Changed;
}