//===-- W65816TiedDefSpill.cpp - Pre-RA spill insertion for tied-def ----===//
//
// 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 pass: when a tied-def Acc16 instruction (ADCfi, SBCfi,
// ANDfi, ORAfi, EORfi, ADCi16imm, SBCi16imm, ANDi16imm, ORAi16imm,
// EORi16imm, ADCabs, SBCabs -- see isTiedAcc16Consumer below for the
// authoritative list) has a source vreg whose value is *also* needed
// past the consumer, fast regalloc fails to insert the necessary
// save/restore on its own.
// (Acc16 has exactly one physical register, so the consumer's
// tied-def overwrites the source; with multiple consumers/post-uses
// the source must be spilled and reloaded.)
//
// We insert that explicitly here:
//
//     %dst = TIED_OP %src, ...    (where %src is also used after)
// becomes
//     STAfi %src, freshSlot, 0
//     %dst = TIED_OP %src, ...    (now safely consumes %src)
//     %src_reload = LDAfi freshSlot, 0
//     ... post-consumer uses replaced with %src_reload
//
// Runs pre-RA so the new vregs participate in regalloc's liveness
// analysis and get assigned A.
//
//===----------------------------------------------------------------------===//

#include "W65816.h"
#include "W65816InstrInfo.h"
#include "W65816Subtarget.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/CodeGen/MachineDominators.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/MachineRegisterInfo.h"

using namespace llvm;

#define DEBUG_TYPE "w65816-tied-def-spill"

namespace {

class W65816TiedDefSpill : public MachineFunctionPass {
public:
  static char ID;
  W65816TiedDefSpill() : MachineFunctionPass(ID) {}
  StringRef getPassName() const override {
    return "W65816 tied-def spill insertion";
  }

  void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.setPreservesCFG();
    AU.addRequired<MachineDominatorTreeWrapperPass>();
    AU.addPreserved<MachineDominatorTreeWrapperPass>();
    MachineFunctionPass::getAnalysisUsage(AU);
  }

  bool runOnMachineFunction(MachineFunction &MF) override;
};

} // namespace

char W65816TiedDefSpill::ID = 0;

INITIALIZE_PASS(W65816TiedDefSpill, DEBUG_TYPE,
                "W65816 tied-def spill insertion", false, false)

FunctionPass *llvm::createW65816TiedDefSpill() {
  return new W65816TiedDefSpill();
}

// Allowlist of tied-def consumer pseudos that are known to fail
// fast regalloc when their source has multiple uses.  Restricting
// to this set avoids regressing other patterns whose existing
// regalloc behaviour is correct.
//
// All entries (see W65816Helpers::isTiedAcc16Consumer) have shape
// `(outs Acc16:$dst), (ins Acc16:$src, memfi:$addr)` or similar
// tied-source-Acc16 + side-load form, matching the failure pattern
// observed in `bump` / `eval`.  The shared predicate is reused below.

bool W65816TiedDefSpill::runOnMachineFunction(MachineFunction &MF) {
  // Only pre-RA: skip if vregs are already gone.
  if (!MF.getRegInfo().getNumVirtRegs())
    return false;
  // At -O0/optnone, the spill+reload pattern this pass introduces
  // doesn't get coalesced and ends up wasting frame space without
  // helping greedy.  Same skip rationale as WidenAcc16.
  if (MF.getFunction().hasOptNone())
    return false;

  MachineRegisterInfo &MRI = MF.getRegInfo();
  const W65816Subtarget &STI = MF.getSubtarget<W65816Subtarget>();
  const W65816InstrInfo *TII = STI.getInstrInfo();
  MachineDominatorTree &MDT =
      getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree();
  bool Changed = false;

  // Snapshot all candidate (MBB, MI, src-operand-index) tuples first;
  // we mutate the MBB during processing.
  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 each tied-source operand, check if the source vreg has
      // any use other than this MI.  If yes, queue for spill.
      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 Reg = MO.getReg();
        if (!Reg.isVirtual()) continue;
        // Count uses excluding this one.  If any other instruction
        // reads Reg, we need to preserve it across the tied-def
        // consumer.
        // Conservative: only spill when one of the OTHER uses is a
        // COPY to a *physreg* (typically a return-value setup or a
        // call-arg copy).  This is the canary pattern fast regalloc
        // mishandles — value flowing both into a tied-def consumer
        // AND into a physreg copy at the end of a BB.  Other patterns
        // (vreg-to-vreg COPY, store, etc.) tend to be handled by fast
        // correctly, and triggering on them inflates frame size
        // (vprintf-class functions overflow the 8-bit stack-rel
        // range otherwise).
        bool NeedSpill = false;
        bool BadUse = false;
        for (auto &U : MRI.use_nodbg_instructions(Reg)) {
          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;

    const TargetRegisterClass *RC = MRI.getRegClass(SrcReg);
    if (RC != &W65816::Acc16RegClass)
      continue;

    int FI = MF.getFrameInfo().CreateStackObject(2, Align(2),
                                                 /*isSpillSlot=*/true);
    DebugLoc DL = MI->getDebugLoc();

    // Insert STAfi $src, FI, 0 BEFORE MI.
    BuildMI(*MBB, MI, DL, TII->get(W65816::STAfi))
        .addReg(SrcReg)
        .addFrameIndex(FI)
        .addImm(0);

    Register NewReg = MRI.createVirtualRegister(&W65816::Acc16RegClass);
    auto InsertPos = std::next(MachineBasicBlock::iterator(MI));
    BuildMI(*MBB, InsertPos, DL, TII->get(W65816::LDAfi), NewReg)
        .addFrameIndex(FI)
        .addImm(0);

    // Only rewrite uses that come AFTER MI in program order — earlier
    // uses already saw SrcReg's original value before any tied-def
    // overwrite, so they don't need redirection.  Uses in successor
    // MBBs definitely come after; uses in MI's own MBB after the
    // LDAfi reload come after; uses before MI in its MBB are
    // pre-consumer and stay on SrcReg.
    //
    // DBG_VALUE preservation: after MI, SrcReg's original value is
    // gone (tied-def overwrote it).  Debug operands referring to SrcReg
    // AFTER MI must be redirected to NewReg so the variable continues
    // to be observable post-spill-reload.  Debug operands BEFORE MI
    // see SrcReg's pre-consumer value and stay put.  Without this, the
    // variable's location vanishes for the rest of the function after
    // any tied-def consumer.
    SmallVector<MachineOperand *, 4> ToRewrite;
    for (auto &U : MRI.use_nodbg_operands(SrcReg)) {
      if (U.getParent() == MI) continue;
      MachineBasicBlock *UseMBB = U.getParent()->getParent();
      bool After = false;
      if (UseMBB != MBB) {
        // Different block — only redirect if MI's MBB DOMINATES the
        // use's MBB.  Without dominance, there's a path from the
        // function entry to the use that bypasses MI entirely (e.g.,
        // a loop-exit edge from a pre-loop block straight into a
        // post-loop block).  Redirecting such a use to %19 (which is
        // only defined when MI runs) reads stale data — the previous
        // iter's MI value, or junk if MI never ran.  Caught by parse2/
        // printf returning N-1 because the loop's tied-def spill of n
        // was redirected to the exit block, which on the final iter
        // (loop test fails) sees iter N-1's saved value.
        if (MDT.dominates(MBB, UseMBB))
          After = true;
      } else {
        // Same block — walk forward from MI to end, see if we hit U.
        for (auto it = MachineBasicBlock::iterator(MI), e = MBB->end();
             it != e; ++it) {
          if (&*it == U.getParent()) { After = true; break; }
        }
      }
      if (After) ToRewrite.push_back(&U);
    }
    for (auto *MO : ToRewrite) {
      MO->setReg(NewReg);
      MO->setIsKill(false);
    }

    // Mirror the rewrite for debug operands that live after MI.  Walk
    // every DBG_VALUE using SrcReg; if the DBG_VALUE is in a strictly-
    // dominated MBB or after MI in MI's own MBB, point it at NewReg.
    // Snapshot first since setReg invalidates the use-list iterator.
    // Dedup because DBG_VALUE_LIST may show up multiple times in the
    // use-list (one entry per debug operand referring to SrcReg).
    SmallPtrSet<MachineInstr *, 4> DbgUserSet;
    SmallVector<MachineInstr *, 4> DbgUsers;
    for (MachineInstr &U : MRI.use_instructions(SrcReg)) {
      if (U.isDebugValue() && DbgUserSet.insert(&U).second)
        DbgUsers.push_back(&U);
    }
    for (MachineInstr *DbgMI : DbgUsers) {
      MachineBasicBlock *DbgMBB = DbgMI->getParent();
      bool After = false;
      if (DbgMBB != MBB) {
        if (MDT.dominates(MBB, DbgMBB))
          After = true;
      } else {
        for (auto It = MachineBasicBlock::iterator(MI), E = MBB->end();
             It != E; ++It) {
          if (&*It == DbgMI) { After = true; break; }
        }
      }
      if (!After) continue;
      for (MachineOperand &Op : DbgMI->debug_operands()) {
        if (Op.isReg() && Op.getReg() == SrcReg) {
          Op.setReg(NewReg);
        }
      }
    }

    Changed = true;
  }
  return Changed;
}
