//===-- W65816PromoteFiToImg.cpp - Promote FrameIndex to IMG slot --------===//
//
// 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
//
//===---------------------------------------------------------------------===//
//
// Post-RA, pre-PEI pass.  Counts accesses to each i16-sized FrameIndex
// in the function and rewrites the top-K hottest ones to use IMG8..15
// DP slots ($C0/$C2/.../$CE) instead.  K = number of free IMG8..15
// slots (slots not already used by regalloc decisions).
//
// Why post-RA: at this point regalloc has decided which vregs live in
// physical registers vs spill slots.  The spills appear as the FI
// pseudo-opcodes (LDAfi/STAfi/ADCfi/SBCfi/ANDfi/ORAfi/EORfi/CMPfi),
// and the MFI tells us each FI's final size.  We see all the accesses
// and can safely rewrite — eliminateFrameIndex hasn't yet baked the
// offsets into SP-relative immediates.
//
// Why before W65816ImgCalleeSave: ImgCalleeSave scans the post-PromoteFi
// MIR for IMG8..15 usage and emits prologue PHA-bracketed saves +
// epilogue restores for each used slot.  Our promotion introduces
// fresh IMG8..15 references that ImgCalleeSave will then auto-cover.
//
// Per-access cost change:
//   STAfi → STA_DP : 5 cyc / 3 B → 4 cyc / 2 B (saves 1 cyc/1 B)
//   LDAfi → LDA_DP : 5 cyc / 3 B → 4 cyc / 2 B (saves 1 cyc/1 B)
//   ADCfi → ADC_DP : 5 cyc / 3 B → 4 cyc / 2 B (saves 1 cyc/1 B)
// Per-slot one-time overhead (added by ImgCalleeSave):
//   prologue save  : ~10 cyc / 6 B
//   epilogue restore: ~10 cyc / 6 B
// Net win if access_count * 1 > 20.  Threshold is 5 to leave margin.
//
// Restrictions:
//   - Only i16-sized FIs (2 bytes, offset 0).  Larger slots (i32 halves,
//     structs) are skipped.
//   - Skips fixed/variable-sized objects.
//   - Skips STA8fi (byte store needs SEP/REP wrap incompatible with
//     simple STA_DP — and DP stores 16 bits in M=0).
//   - Skips LDAfi_indY / STAfi_indY (indirect-Y form — different
//     addressing).
//
//===---------------------------------------------------------------------===//

#include "W65816.h"
#include "W65816InstrInfo.h"
#include "W65816Subtarget.h"
#include "llvm/ADT/BitVector.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/CodeGen/MachineFrameInfo.h"
#include "llvm/CodeGen/MachineFunction.h"
#include "llvm/CodeGen/MachineFunctionPass.h"
#include "llvm/CodeGen/MachineInstrBuilder.h"
#include "llvm/CodeGen/MachineLoopInfo.h"
#include "llvm/CodeGen/MachineRegisterInfo.h"
#include "llvm/InitializePasses.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/Format.h"

using namespace llvm;

#define DEBUG_TYPE "w65816-promote-fi-to-img"


namespace {


class W65816PromoteFiToImg : public MachineFunctionPass {
public:
  static char ID;
  W65816PromoteFiToImg() : MachineFunctionPass(ID) {}
  StringRef getPassName() const override {
    return "W65816 promote FrameIndex to IMG8..15 DP slot";
  }
  void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.addRequired<MachineLoopInfoWrapperPass>();
    AU.setPreservesCFG();
    MachineFunctionPass::getAnalysisUsage(AU);
  }
  bool runOnMachineFunction(MachineFunction &MF) override;
};


} // namespace


char W65816PromoteFiToImg::ID = 0;

INITIALIZE_PASS_BEGIN(W65816PromoteFiToImg, DEBUG_TYPE,
                      "W65816 promote FI to IMG", false, false)
INITIALIZE_PASS_DEPENDENCY(MachineLoopInfoWrapperPass)
INITIALIZE_PASS_END(W65816PromoteFiToImg, DEBUG_TYPE,
                    "W65816 promote FI to IMG", false, false)


FunctionPass *llvm::createW65816PromoteFiToImg() {
  return new W65816PromoteFiToImg();
}


// Returns the operand index of the FrameIndex for the given FI pseudo
// opcode, or -1 if this opcode isn't a promotable FI carrier.
static int getFiOperandIdx(unsigned Opc) {
  switch (Opc) {
  case W65816::LDAfi:                                   return 1;
  case W65816::STAfi:                                   return 1;
  case W65816::CMPfi:                                   return 1;
  case W65816::ADCfi:
  case W65816::SBCfi:
  case W65816::ANDfi:
  case W65816::ORAfi:
  case W65816::EORfi:                                   return 2;
  default:                                              return -1;
  }
}


// Map a promotable FI pseudo to the corresponding DP MC opcode.
static unsigned getDpOpcode(unsigned Opc) {
  switch (Opc) {
  case W65816::LDAfi: return W65816::LDA_DP;
  case W65816::STAfi: return W65816::STA_DP;
  case W65816::CMPfi: return W65816::CMP_DP;
  case W65816::ADCfi: return W65816::ADC_DP;
  case W65816::SBCfi: return W65816::SBC_DP;
  case W65816::ANDfi: return W65816::AND_DP;
  case W65816::ORAfi: return W65816::ORA_DP;
  case W65816::EORfi: return W65816::EOR_DP;
  default: return 0;
  }
}


// IMG8..IMG15 sit at DP addresses 0xC0, 0xC2, ..., 0xCE.  IMG0..IMG7
// are at 0xD0..0xDE.  Returns the DP byte for IMGn.
static uint8_t dpAddrForImg(unsigned ImgIdx) {
  assert(ImgIdx < 16 && "IMG index out of range");
  if (ImgIdx < 8) return 0xD0 + 2 * ImgIdx;
  return 0xC0 + 2 * (ImgIdx - 8);
}


bool W65816PromoteFiToImg::runOnMachineFunction(MachineFunction &MF) {
  // DISABLED again 2026-05-13 (3rd-attempt write-up).  Two new findings:
  //   1. With kMaxPromote=2 and IMG0..7 (caller-save, skip ImgCalleeSave),
  //      sumSquares regressed 56 → 72 inst because the FIs picked by
  //      access-count (fi#2, fi#3) are intermediate spill temps, not
  //      the i32-accumulator's halves (which are different FIs).  The
  //      loop body ends up using BOTH IMG and stack slots for related
  //      values.
  //   2. To pick the RIGHT FIs (those corresponding to PHI-cycled
  //      values like the i32 accumulator), we need either:
  //      (a) IR-level analysis BEFORE FI assignment, or
  //      (b) Post-RA dataflow analysis to identify "long-lived" FIs
  //          (active across the loop back-edge with no def/use boundary).
  //   This is the next blocker.  Disabled until either (a) or (b) is
  //   implemented.
  return false;
  if (skipFunction(MF.getFunction())) return false;
  const W65816Subtarget &STI = MF.getSubtarget<W65816Subtarget>();
  const W65816InstrInfo *TII = STI.getInstrInfo();
  MachineFrameInfo &MFI = MF.getFrameInfo();

  // 1. Walk all instructions, count FI accesses for promotable opcodes.
  // Weight by loop depth: an access inside a depth-N loop counts as
  // 10^N to model the dynamic execution count (an inner-loop slot
  // gets executed many times per outer call).
  MachineLoopInfo &MLI =
      getAnalysis<MachineLoopInfoWrapperPass>().getLI();
  DenseMap<int, unsigned> AccessCount;
  DenseMap<int, SmallVector<MachineInstr *, 8>> AccessSites;
  for (MachineBasicBlock &MBB : MF) {
    unsigned LoopDepth = MLI.getLoopDepth(&MBB);
    unsigned Weight = 1;
    for (unsigned i = 0; i < LoopDepth && i < 3; ++i) Weight *= 10;
    for (MachineInstr &MI : MBB) {
      int FiIdx = getFiOperandIdx(MI.getOpcode());
      if (FiIdx < 0) continue;
      const MachineOperand &MO = MI.getOperand(FiIdx);
      if (!MO.isFI()) continue;
      int FI = MO.getIndex();
      if (MFI.isVariableSizedObjectIndex(FI)) continue;
      if (MFI.getObjectSize(FI) != 2) continue;
      if (FI < 0) continue;
      const MachineOperand &OffMO = MI.getOperand(FiIdx + 1);
      if (!OffMO.isImm() || OffMO.getImm() != 0) continue;
      AccessCount[FI] += Weight;
      AccessSites[FI].push_back(&MI);
    }
  }
  if (AccessCount.empty()) return false;

  // 2. Determine which IMG0..7 slots are already in use (caller-save).
  // Use caller-save IMG0..7 instead of callee-save IMG8..15: this lets
  // us skip ImgCalleeSave entirely (no prologue/epilogue overhead).
  // The trade-off: any call inside the function clobbers IMG0..7.  Mark
  // any function with calls as "callees might clobber" → skip promotion.
  // This restricts wins to leaf functions (no internal calls).
  BitVector UsedImg(8, false);
  for (MachineBasicBlock &MBB : MF) {
    for (MachineInstr &MI : MBB) {
      // Skip CALL instructions — their `implicit-def dead $img0..7`
      // operand list marks every IMG slot used, but that's just the
      // caller-save annotation, not actual value-bearing usage.
      if (MI.isCall()) continue;
      for (const MachineOperand &MO : MI.operands()) {
        if (!MO.isReg() || !MO.getReg().isPhysical()) continue;
        Register R = MO.getReg();
        unsigned ImgIdx = 16;
        if      (R == W65816::IMG0)  ImgIdx = 0;
        else if (R == W65816::IMG1)  ImgIdx = 1;
        else if (R == W65816::IMG2)  ImgIdx = 2;
        else if (R == W65816::IMG3)  ImgIdx = 3;
        else if (R == W65816::IMG4)  ImgIdx = 4;
        else if (R == W65816::IMG5)  ImgIdx = 5;
        else if (R == W65816::IMG6)  ImgIdx = 6;
        else if (R == W65816::IMG7)  ImgIdx = 7;
        if (ImgIdx < 8) UsedImg.set(ImgIdx);
      }
    }
  }

  // 3. Sort FIs by access count (descending).
  SmallVector<int, 16> Ordered;
  for (auto &P : AccessCount) Ordered.push_back(P.first);
  std::sort(Ordered.begin(), Ordered.end(),
            [&](int A, int B) { return AccessCount[A] > AccessCount[B]; });

  // 4. Assign IMG slots greedily.  Each IMG8..15 slot used triggers
  //    a save/restore pair in W65816ImgCalleeSave (~20 cyc + ~12 B
  //    per slot per CALL into this function).  For recursive or
  //    deep-call-stack functions, that overhead dominates the per-
  //    access savings — measured: promoting 4 slots in fib(10)
  //    regressed it 38% (12617 → 17391 cyc).  Gate on a very high
  //    threshold + bail entirely if the function has any calls (the
  //    save/restore cost compounds with recursion / call frequency
  //    in ways the static access count can't capture).
  bool HasCalls = false;
  bool IsRecursive = false;
  StringRef SelfName = MF.getName();
  for (MachineBasicBlock &MBB : MF) {
    for (MachineInstr &MI : MBB) {
      if (MI.isCall()) {
        HasCalls = true;
        // Check for self-call (recursive).
        for (const MachineOperand &MO : MI.operands()) {
          if (MO.isGlobal() && MO.getGlobal()->getName() == SelfName)
            IsRecursive = true;
          else if (MO.isSymbol() && SelfName == MO.getSymbolName())
            IsRecursive = true;
        }
      }
    }
  }
  // Recursive functions: skip — recursion makes per-call overhead
  // compound (each level of recursion pays the save/restore).
  if (IsRecursive) return false;
  // Caller-save IMG0..7 strategy: any internal call clobbers them, so
  // the only safe promoted slots are those whose lifetime doesn't
  // cross a call.  For now, only promote in leaf functions (no internal
  // calls at all).  This catches simple loops like sumSquares (which
  // calls __umulhisi3 — but that's in libgcc.s and doesn't actually
  // touch IMG0..7; treat libgcc multiplies as IMG-safe).
  //
  // Whitelist of libgcc functions known to not touch IMG0..7.
  auto isImgSafeLibcall = [](const MachineInstr &MI) -> bool {
    if (!MI.isCall()) return false;
    for (const MachineOperand &MO : MI.operands()) {
      StringRef Name;
      if (MO.isGlobal()) Name = MO.getGlobal()->getName();
      else if (MO.isSymbol()) Name = MO.getSymbolName();
      else continue;
      // libgcc.s multiply/divide/shift helpers — verified to only use
      // $E0..$E9 internally, no IMG0..7 touch.
      if (Name == "__umulhisi3" || Name == "__mulhi3" ||
          Name == "__mulsi3" || Name == "__udivhi3" ||
          Name == "__umodhi3" || Name == "__divhi3" ||
          Name == "__modhi3" || Name == "__udivsi3" ||
          Name == "__umodsi3" || Name == "__divsi3" ||
          Name == "__modsi3" || Name == "__ashlhi3" ||
          Name == "__lshrhi3" || Name == "__ashrhi3" ||
          Name == "__ashlsi3" || Name == "__lshrsi3" ||
          Name == "__ashrsi3")
        return true;
      return false;
    }
    return false;
  };
  bool AllCallsImgSafe = true;
  for (MachineBasicBlock &MBB : MF) {
    for (MachineInstr &MI : MBB) {
      if (MI.isCall() && !isImgSafeLibcall(MI)) {
        AllCallsImgSafe = false;
        break;
      }
    }
    if (!AllCallsImgSafe) break;
  }
  if (HasCalls && !AllCallsImgSafe) return false;
  // Threshold: per-access save is 1 cyc, no save/restore overhead.  We
  // just need the access count to be > 0 to win.  Use a small threshold
  // for safety (avoid promoting marginal slots).
  const unsigned kAccessThreshold = 5u;
  const unsigned kMaxPromote = 2u;
  DenseMap<int, unsigned> FiToImgIdx;
  unsigned NextFreeImg = 0;
  for (int FI : Ordered) {
    if (AccessCount[FI] < kAccessThreshold) break;
    if (FiToImgIdx.size() >= kMaxPromote) break;
    while (NextFreeImg < 8 && UsedImg.test(NextFreeImg)) ++NextFreeImg;
    if (NextFreeImg >= 8) break;
    FiToImgIdx[FI] = NextFreeImg;  // Map to IMG0..7 (caller-save)
    ++NextFreeImg;
  }
  if (FiToImgIdx.empty()) return false;

  // 5. Rewrite each access.  Insert the new DP MC inst before the
  //    pseudo, then erase the pseudo.  Preserve flags and tied-def
  //    semantics via implicit operands.
  bool Changed = false;
  for (auto &P : FiToImgIdx) {
    int FI = P.first;
    unsigned ImgIdx = P.second;
    uint8_t DpAddr = dpAddrForImg(ImgIdx);
    LLVM_DEBUG(dbgs() << "Promote fi#" << FI << " -> IMG"
                      << ImgIdx << " ($" << format("%02x", DpAddr)
                      << "), " << AccessCount[FI] << " accesses\n");
    for (MachineInstr *MI : AccessSites[FI]) {
      unsigned Opc = MI->getOpcode();
      unsigned NewOpc = getDpOpcode(Opc);
      if (!NewOpc) continue;
      MachineBasicBlock *MBB = MI->getParent();
      DebugLoc DL = MI->getDebugLoc();
      MachineInstrBuilder NewMI =
          BuildMI(*MBB, MI, DL, TII->get(NewOpc)).addImm(DpAddr);
      // Carry implicit-def $a (LDA/ADC/SBC/AND/ORA/EOR all write $a)
      // and implicit-use $a (STA/CMP/ADC/SBC/AND/ORA/EOR all read $a).
      // ADCfi/SBCfi additionally use $p; their DP equivalents read $p
      // implicitly via the tablegen Defs/Uses.  But since we built the
      // new MI from TII->get(NewOpc), the implicit operands from the
      // descriptor are auto-added.  We only need to copy non-FI explicit
      // operands... which for our pseudos are register operands.  The
      // physical register defs/uses they carry must be preserved.
      for (const MachineOperand &MO : MI->operands()) {
        if (MO.isReg() && MO.getReg().isPhysical() && MO.isImplicit()) {
          // Skip — already added by descriptor.
          continue;
        }
        if (MO.isReg() && MO.getReg().isPhysical() && !MO.isImplicit()) {
          // Explicit physreg operand (e.g., the $a in STAfi $a, fi, 0).
          // Convert to implicit so the DP MC inst's descriptor matches.
          RegState Flags = MO.isDef() ? RegState::ImplicitDefine
                                       : RegState::Implicit;
          if (MO.isKill()) Flags = Flags | RegState::Kill;
          NewMI.addReg(MO.getReg(), Flags);
        }
        // FI/offset operands are skipped — replaced by the DP imm above.
        // VReg defs/uses should be gone post-RA; if any survived, skip.
      }
      MI->eraseFromParent();
      Changed = true;
    }
    // Mark the FI as dead so PEI can skip allocating stack for it.
    // MFI doesn't expose RemoveStackObject publicly, but setting size
    // to 0 also works in most code paths.  Actually leave it alive —
    // a 2-byte unused slot is cheap, and removing exposes us to
    // PEI bugs.
  }
  return Changed;
}