//===-- W65816WidenAcc16.cpp - Promote Acc16 vregs to Wide16 ------------===//
//
// 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 promotes Acc16 vregs (constrained to physreg A only)
// to the wider Wide16 class (A + IMG0..IMG7).  Greedy regalloc gets
// 9-way pressure relief on the i16 register class; functions that
// previously failed with "ran out of registers" can now spread their
// live i16 values across A and the DP-backed imaginaries.
//
// Cross-class moves between A and IMGn are LDA/STA dp (4 cyc each way,
// 2 bytes), emitted by W65816InstrInfo::copyPhysReg.  The constraint
// that arithmetic ops require their source in A propagates back from
// the use sites — regalloc coerces Wide16 vregs to Acc16 (= {A}) at
// those sites and inserts the necessary COPYs.
//
// Calls clobber IMGn (caller-save), so any vreg in IMGn that lives
// across a call gets spilled to stack by regalloc.  This pass doesn't
// model that explicitly — it relies on the calling convention's
// regmask to mark IMGn clobbered.
//
//===----------------------------------------------------------------------===//

#include "W65816.h"
#include "W65816InstrInfo.h"
#include "W65816Subtarget.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-widen-acc16"

namespace {

class W65816WidenAcc16 : public MachineFunctionPass {
public:
  static char ID;
  W65816WidenAcc16() : MachineFunctionPass(ID) {}
  StringRef getPassName() const override {
    return "W65816 Acc16 → Wide16 promotion";
  }
  void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.setPreservesCFG();
    MachineFunctionPass::getAnalysisUsage(AU);
  }
  bool runOnMachineFunction(MachineFunction &MF) override;
};

} // namespace

char W65816WidenAcc16::ID = 0;

INITIALIZE_PASS(W65816WidenAcc16, DEBUG_TYPE,
                "W65816 Acc16 → Wide16 promotion", false, false)

FunctionPass *llvm::createW65816WidenAcc16() {
  return new W65816WidenAcc16();
}

// Returns true if the vreg has any physreg-COPY use that would conflict
// with Wide16 class assignment.  $a is a member of Wide16 (Wide16 = A +
// IMG0..15), so a COPY to $a is fine — the vreg can be Wide16 and
// regalloc will pick $a to coalesce.  $x / $y are in Idx16, NOT in
// Wide16, so a COPY to those forces the vreg to NOT be in Wide16
// (verifier would reject).
static bool flowsToIncompatiblePhysReg(Register VReg,
                                       const MachineRegisterInfo &MRI) {
  for (auto &U : MRI.use_nodbg_instructions(VReg)) {
    if (!U.isCopy()) continue;
    const MachineOperand &Dst = U.getOperand(0);
    if (!Dst.isReg() || !Dst.getReg().isPhysical()) continue;
    Register P = Dst.getReg();
    if (P == W65816::A) continue;
    if (P >= W65816::IMG0 && P <= W65816::IMG15) continue;
    return true;
  }
  return false;
}

// Returns true if VReg's def is a COPY from a physreg whose class is not
// Wide16-compatible.  copyPhysReg only handles a fixed set of source/dest
// pairs; an incompatible source physreg (e.g., DPF0, the i64-return
// high-half carrier) lowered to an IMG dest would crash with an
// "unhandled copyPhysReg" assertion at AsmPrinter time.  (Currently
// only the Phase-2 PHI widening uses this; that's disabled, so mark
// unused.)
[[maybe_unused]] static bool comesFromIncompatiblePhysReg(Register VReg,
                                         const MachineRegisterInfo &MRI) {
  for (auto &D : MRI.def_instructions(VReg)) {
    if (!D.isCopy()) continue;
    const MachineOperand &Src = D.getOperand(1);
    if (!Src.isReg() || !Src.getReg().isPhysical()) continue;
    Register P = Src.getReg();
    if (P == W65816::A) continue;
    if (P >= W65816::IMG0 && P <= W65816::IMG15) continue;
    return true;
  }
  return false;
}

// Returns true if the vreg is used by any PHI.  PHI input/result must
// share the same register class (verifier requirement).  Rather than
// also widen the PHI's result and recursively all of its uses, we skip
// vregs caught up in PHIs entirely — leaves a few wins on the table
// but avoids cross-MBB analysis here.
static bool usedByPhi(Register VReg, const MachineRegisterInfo &MRI) {
  for (auto &U : MRI.use_nodbg_instructions(VReg)) {
    if (U.isPHI()) return true;
  }
  return false;
}

// Returns true if all non-debug, non-COPY uses of VReg are at operands
// whose required register class accepts Wide16 (i.e., Wide16 or a
// superclass).  COPY uses are unconstrained — fine.  PHI uses already
// filtered earlier.  If any use's operand class is strictly narrower
// than Wide16 (i.e., Acc16-only, Idx16-only, etc.), return false: the
// verifier rejects passing a Wide16 vreg to such an operand.
static bool allUsesAcceptWide(Register VReg,
                              const MachineRegisterInfo &MRI,
                              const TargetRegisterInfo &TRI,
                              const TargetInstrInfo &TII) {
  for (auto &MO : MRI.use_nodbg_operands(VReg)) {
    MachineInstr *UMI = MO.getParent();
    if (UMI->isCopy()) continue;  // COPY accepts anything
    if (UMI->isPHI()) return false;  // already filtered, but be safe
    unsigned OpIdx = UMI->getOperandNo(&MO);
    (void)TRI;
    const TargetRegisterClass *Expected =
        TII.getRegClass(UMI->getDesc(), OpIdx);
    if (!Expected) continue;  // no constraint
    if (Expected == &W65816::Wide16RegClass) continue;
    // Check superclass relationship: Wide16 must be a sub-or-equal of
    // Expected for the use to accept Wide16 vregs.  A common case:
    // Expected is a superclass that includes Wide16.  If Expected is
    // narrower (e.g., Acc16 only), reject.
    if (Expected->hasSubClassEq(&W65816::Wide16RegClass)) continue;
    return false;
  }
  return true;
}

bool W65816WidenAcc16::runOnMachineFunction(MachineFunction &MF) {
  if (!MF.getRegInfo().getNumVirtRegs()) return false;
  // At -O0 / optnone, register coalescing doesn't run, so the COPY we
  // insert to bridge Acc16 → Wide16 doesn't get folded; instead it
  // forces wide16 spills through DP-mapped slots that collide and
  // produce miscompiles around modify-in-place ops (lda dp; inc a;
  // sta dp; lda dp reads pre-inc value).  The promotion is purely a
  // performance optimization, so skip it for optnone functions.
  if (MF.getFunction().hasOptNone()) return false;
  MachineRegisterInfo &MRI = MF.getRegInfo();
  const W65816Subtarget &STI = MF.getSubtarget<W65816Subtarget>();
  const W65816InstrInfo *TII = STI.getInstrInfo();
  const TargetRegisterInfo *TRI = STI.getRegisterInfo();
  bool Changed = false;

  // For each Acc16 vreg, insert a COPY to a fresh Wide16 vreg right
  // after its def, then redirect all uses to the Wide16 vreg.  The
  // original Acc16 vreg keeps its tight constraint (= {A}) for the
  // def site (which is typically a pseudo whose AsmPrinter expansion
  // assumes A); the new Wide16 vreg is free for greedy to allocate
  // anywhere in {A, IMG0..IMG7}.  When both end up in A, the COPY
  // is a no-op the regalloc/coalescer collapses; when the Wide16
  // vreg lands on IMGn, the COPY becomes STA dp via copyPhysReg.
  SmallVector<Register, 16> Candidates;
  for (unsigned i = 0; i < MRI.getNumVirtRegs(); ++i) {
    Register VReg = Register::index2VirtReg(i);
    if (MRI.def_empty(VReg)) continue;
    if (MRI.getRegClass(VReg) != &W65816::Acc16RegClass) continue;
    if (flowsToIncompatiblePhysReg(VReg, MRI)) continue;
    if (usedByPhi(VReg, MRI)) continue;
    if (!MRI.hasOneDef(VReg)) continue;  // require single SSA def
    if (!allUsesAcceptWide(VReg, MRI, *TRI, *TII)) continue;
    Candidates.push_back(VReg);
  }

  for (Register VReg : Candidates) {
    MachineInstr *DefMI = &*MRI.def_instructions(VReg).begin();
    MachineBasicBlock *MBB = DefMI->getParent();
    DebugLoc DL = DefMI->getDebugLoc();
    Register WideReg = MRI.createVirtualRegister(&W65816::Wide16RegClass);
    // Insert AFTER the def, but if the def is a PHI, walk past all
    // PHIs in the block first — verifier requires all PHIs at MBB
    // entry, no non-PHI may sit between them.
    auto InsertAt = std::next(MachineBasicBlock::iterator(DefMI));
    if (DefMI->isPHI()) {
      while (InsertAt != MBB->end() && InsertAt->isPHI()) ++InsertAt;
    }
    BuildMI(*MBB, InsertAt, DL, TII->get(TargetOpcode::COPY), WideReg)
        .addReg(VReg);
    // Rewrite all non-debug uses of VReg (other than the COPY we just
    // inserted) to WideReg.
    SmallVector<MachineOperand *, 8> ToRewrite;
    for (auto &U : MRI.use_nodbg_operands(VReg)) {
      MachineInstr *UMI = U.getParent();
      if (UMI->getOpcode() == TargetOpcode::COPY &&
          UMI->getOperand(0).getReg() == WideReg) continue;
      ToRewrite.push_back(&U);
    }
    for (auto *MO : ToRewrite) {
      MO->setReg(WideReg);
      MO->setIsKill(false);
    }
    Changed = true;
  }

  // Phase 2: PHI cycle widening.  EXPERIMENTAL, currently disabled —
  // see end of pass for explanation.
  #if 0
  // PHIs whose def class is Acc16 keep
  // the value pinned to $a across iterations, forcing stack spills
  // when the PHI is live across calls or other A-clobbering ops.
  // For sumSquares-style loops with an i32 accumulator, this manifests
  // as per-iter `LDA slot ; ADC ; STA slot ; LDA slot ; STA slot` (the
  // last LDA/STA pair is the PHI-back-edge copy).  If we widen the
  // PHI's def to Wide16, regalloc can keep it in an IMG slot and the
  // back-edge PHI copy collapses to a register coalesce.
  //
  // To widen a PHI:
  //   1. Compute the SCC of Acc16 vregs connected by PHI edges (PHI
  //      def ↔ PHI incoming vreg).  This catches mutually-recursive
  //      PHIs in nested loops.
  //   2. For every member: verify all non-PHI uses accept Wide16, no
  //      flow to a physreg, single def.
  //   3. For each PHI in the SCC, walk its incoming list.  Each
  //      incoming vreg is either ALREADY in the SCC (another PHI, no
  //      bridge needed) or an external Acc16 vreg whose value flows
  //      into the SCC — bridge it by inserting `WWide = COPY W` at
  //      the end of the predecessor block and pointing the PHI's
  //      incoming at WWide.
  //   4. Change every SCC member's register class to Wide16.
  auto worklistInsertIfAcc16 = [&MRI](Register V,
                                      DenseSet<Register> &Seen,
                                      SmallVectorImpl<Register> &WL) {
    if (!V.isVirtual()) return;
    if (MRI.getRegClass(V) != &W65816::Acc16RegClass) return;
    if (!Seen.insert(V).second) return;
    WL.push_back(V);
  };

  SmallVector<MachineInstr *, 16> AcctPhis;
  for (MachineBasicBlock &MBB : MF) {
    for (MachineInstr &MI : MBB.phis()) {
      Register DefV = MI.getOperand(0).getReg();
      if (MRI.getRegClass(DefV) == &W65816::Acc16RegClass) {
        AcctPhis.push_back(&MI);
      }
    }
  }
  DenseSet<Register> ProcessedPhiVregs;
  for (MachineInstr *Seed : AcctPhis) {
    Register SeedDef = Seed->getOperand(0).getReg();
    if (ProcessedPhiVregs.count(SeedDef)) continue;
    // Build SCC by following PHI edges in both directions.
    DenseSet<Register> Comp;
    SmallVector<Register, 8> Stack;
    worklistInsertIfAcc16(SeedDef, Comp, Stack);
    while (!Stack.empty()) {
      Register V = Stack.pop_back_val();
      // Forward: V flows into other PHIs as an incoming → include those PHI defs.
      for (auto &U : MRI.use_nodbg_instructions(V)) {
        if (!U.isPHI()) continue;
        Register PhiDef = U.getOperand(0).getReg();
        worklistInsertIfAcc16(PhiDef, Comp, Stack);
      }
      // Backward: if V is itself a PHI def, include the incoming vregs.
      MachineInstr *DM = &*MRI.def_instructions(V).begin();
      if (!DM || !DM->isPHI()) continue;
      for (unsigned i = 1, e = DM->getNumOperands(); i < e; i += 2) {
        MachineOperand &MO = DM->getOperand(i);
        if (!MO.isReg() || !MO.getReg().isVirtual()) continue;
        worklistInsertIfAcc16(MO.getReg(), Comp, Stack);
      }
    }
    for (Register V : Comp) ProcessedPhiVregs.insert(V);

    // Validate every member.  PHI uses are ACCEPTED when the consumer
    // PHI is itself in the SCC (those PHIs are being widened in
    // lock-step).  Narrow-class uses (e.g., INA_PSEUDO's tied-def
    // input requires Acc16) are ALSO accepted — we'll insert a
    // Wide16→Acc16 COPY at the use site after widening.  The only
    // unrecoverable cases are: PHI uses where the consumer PHI is
    // outside the SCC (forcing cross-SCC class merging), and physreg
    // flow to $x/$y/etc. (handled separately above).
    auto usesAcceptInSCC = [&](Register V,
                               SmallVectorImpl<MachineOperand *> *NarrowSites)
        -> bool {
      for (auto &MO : MRI.use_nodbg_operands(V)) {
        MachineInstr *UMI = MO.getParent();
        if (UMI->isCopy()) continue;
        if (UMI->isPHI()) {
          Register PhiDef = UMI->getOperand(0).getReg();
          if (Comp.count(PhiDef)) continue;  // co-widened
          return false;
        }
        unsigned OpIdx = UMI->getOperandNo(&MO);
        const TargetRegisterClass *Expected =
            TII->getRegClass(UMI->getDesc(), OpIdx);
        if (!Expected) continue;
        if (Expected == &W65816::Wide16RegClass) continue;
        if (Expected->hasSubClassEq(&W65816::Wide16RegClass)) continue;
        // Expected is narrower than Wide16 (e.g., Acc16-only tied
        // input).  Mark for runtime narrowing — we'll insert a COPY
        // at apply time.
        if (NarrowSites) NarrowSites->push_back(&MO);
      }
      return true;
    };
    bool ok = true;
    SmallVector<MachineOperand *, 8> NarrowSites;
    for (Register V : Comp) {
      if (!MRI.hasOneDef(V)) { ok = false; break; }
      if (flowsToIncompatiblePhysReg(V, MRI)) { ok = false; break; }
      if (comesFromIncompatiblePhysReg(V, MRI)) { ok = false; break; }
      if (!usesAcceptInSCC(V, &NarrowSites)) { ok = false; break; }
    }
    if (!ok) continue;

    // Apply widening.  First insert bridge COPYs at predecessor edges
    // for external (non-Comp) Acc16 incomings to each PHI in Comp.
    SmallVector<std::pair<MachineInstr *, unsigned>, 16> BridgeSites;
    for (Register V : Comp) {
      MachineInstr *DM = &*MRI.def_instructions(V).begin();
      if (!DM->isPHI()) continue;
      for (unsigned i = 1, e = DM->getNumOperands(); i < e; i += 2) {
        MachineOperand &MO = DM->getOperand(i);
        if (!MO.isReg() || !MO.getReg().isVirtual()) continue;
        Register Inc = MO.getReg();
        if (Comp.count(Inc)) continue;  // in-SCC, no bridge needed
        // External incoming: ensure it's currently Acc16; if so, we'll
        // insert a COPY at the predecessor block's end.
        if (MRI.getRegClass(Inc) != &W65816::Acc16RegClass &&
            MRI.getRegClass(Inc) != &W65816::Wide16RegClass) {
          ok = false;
          break;
        }
        BridgeSites.push_back({DM, i});
      }
      if (!ok) break;
    }
    if (!ok) continue;

    // Insert bridges.
    for (auto &Site : BridgeSites) {
      MachineInstr *PhiMI = Site.first;
      unsigned OpIdx = Site.second;
      Register Inc = PhiMI->getOperand(OpIdx).getReg();
      MachineBasicBlock *PredMBB = PhiMI->getOperand(OpIdx + 1).getMBB();
      // If already Wide16 (e.g., another candidate widened it already),
      // no bridge needed — but we still need the PHI incoming to use
      // a Wide16 vreg.  Use Inc directly.
      if (MRI.getRegClass(Inc) == &W65816::Wide16RegClass) {
        continue;
      }
      // Insert COPY before the predecessor's terminator(s).
      auto InsertPos = PredMBB->getFirstTerminator();
      DebugLoc DL = (InsertPos == PredMBB->end())
                        ? PredMBB->findBranchDebugLoc()
                        : InsertPos->getDebugLoc();
      Register WideInc = MRI.createVirtualRegister(&W65816::Wide16RegClass);
      BuildMI(*PredMBB, InsertPos, DL, TII->get(TargetOpcode::COPY),
              WideInc)
          .addReg(Inc);
      PhiMI->getOperand(OpIdx).setReg(WideInc);
      PhiMI->getOperand(OpIdx).setIsKill(false);
    }

    // Force every SCC member to Img16 (IMG-only, no A).  Using Wide16
    // (A + IMG) doesn't work here: the Register Coalescer joins our
    // Wide16 vregs with adjacent Acc16 vregs (intersection = Acc16)
    // and narrows them back to A-only, defeating the widening.  Img16
    // intersects Acc16 to ∅, so the coalescer can't merge — the PHI
    // stays in IMG.  This is correct anyway for the common case (PHI
    // live across a call): A is JSL-clobbered, so it can't carry the
    // value through, and IMG8..15 is the right home.
    for (Register V : Comp) {
      MRI.setRegClass(V, &W65816::Img16RegClass);
    }
    // Insert narrowing COPYs at each narrow-class use site.  Each site
    // is `... = OP V, ...` where the operand requires Acc16 but V is
    // now Wide16.  Replace with `%Vacc = COPY V (Acc16); ... = OP %Vacc, ...`.
    for (MachineOperand *MO : NarrowSites) {
      MachineInstr *UMI = MO->getParent();
      Register OldReg = MO->getReg();
      Register NarrowReg =
          MRI.createVirtualRegister(&W65816::Acc16RegClass);
      DebugLoc DL = UMI->getDebugLoc();
      BuildMI(*UMI->getParent(), UMI, DL, TII->get(TargetOpcode::COPY),
              NarrowReg)
          .addReg(OldReg);
      MO->setReg(NarrowReg);
      MO->setIsKill(false);
    }
    Changed = true;
  }
  #endif
  // Why disabled (2026-05-13 attempt):
  // - Widening PHI cycles to Wide16 (= A + IMG0..15) is undone by the
  //   Register Coalescer: it joins our Wide16 vregs with adjacent
  //   Acc16 vregs via the bridge COPYs we insert, and the resulting
  //   joint class is `intersect(Wide16, Acc16) = Acc16`.  Net effect:
  //   no IMG, just more code through the coalescer.
  // - Switching to Img16 (= IMG0..15, no A) defeats the coalescer
  //   (intersection with Acc16 is ∅) but forces ALL widened PHIs into
  //   IMG slots even when A would be better, AND triggers cascading
  //   copyPhysReg paths that aren't all implemented (e.g., DPF0 → IMG
  //   for i64 libcall return values), aborting clang on runtime builds.
  // - A targeted fix needs either (a) a class that the coalescer
  //   refuses to join with Acc16 yet that still allows A as a member,
  //   (b) a post-coalescer pass that re-widens specific high-traffic
  //   vregs back to Img16, or (c) regalloc cost-model tuning so it
  //   prefers IMG8..15 over stack for loop-live values.
  return Changed;
}
