//===-- W65816LowerWide32.cpp - Wide32 -> 2x i16 pre-RA lowering ---------===//
//
// 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 that decomposes Wide32 register-pair vregs into pairs
// of i16 vregs.  Without this, greedy / basic regalloc fails on i64-heavy
// functions (`RegAllocBase` crashes during `allocatePhysRegs`) because
// the i64-via-2-i32-via-Wide32 chain produces too many simultaneously
// live register-pair vregs.  After this pass, only i16 vregs remain at
// the regalloc input — Wide32 lives only inside this pass and the new
// LDAptr32S / STAptr32S / STBptr32S pseudos that take 2 i16 ptr operands
// directly.
//
// Walks the MIR and:
//   1. Finds REG_SEQUENCE producing Wide32 / Acc32 / AnyWide32; records
//      the (lo, hi) i16 source operands; queues the REG_SEQUENCE for
//      erasure.
//   2. Finds COPY whose dest is a Wide32 vreg and whose src is another
//      mapped Wide32 vreg; chains the (lo, hi) mapping forward.
//   3. Rewrites EXTRACT_SUBREG of mapped Wide32 vregs by replacing the
//      destination vreg with the appropriate half (sub_lo or sub_hi).
//   4. Rewrites LDAptr32 / STAptr32 / STBptr32 with a mapped Wide32 ptr
//      to the corresponding LDAptr32S / STAptr32S / STBptr32S pseudo
//      with two separate i16 operands.
//
// Bail / safety: any Wide32 vreg whose def we can't decompose is left
// in place — regalloc may still struggle but no miscompile.
//
//===----------------------------------------------------------------------===//

#include "W65816.h"
#include "W65816InstrInfo.h"
#include "W65816Subtarget.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallVector.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"
#include "llvm/CodeGen/TargetInstrInfo.h"
#include "llvm/CodeGen/TargetRegisterInfo.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"

using namespace llvm;

#define DEBUG_TYPE "w65816-lower-wide32"

namespace {

class W65816LowerWide32 : public MachineFunctionPass {
public:
  static char ID;
  W65816LowerWide32() : MachineFunctionPass(ID) {}
  StringRef getPassName() const override {
    return "W65816 Wide32 -> 2x i16 lowering";
  }
  void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.setPreservesCFG();
    MachineFunctionPass::getAnalysisUsage(AU);
  }
  bool runOnMachineFunction(MachineFunction &MF) override;
};

} // namespace

char W65816LowerWide32::ID = 0;

INITIALIZE_PASS(W65816LowerWide32, DEBUG_TYPE,
                "W65816 Wide32 lowering", false, false)

FunctionPass *llvm::createW65816LowerWide32() {
  return new W65816LowerWide32();
}

static bool isWide32RC(const TargetRegisterClass *RC) {
  return RC == &W65816::Wide32RegClass ||
         RC == &W65816::Acc32RegClass ||
         RC == &W65816::AnyWide32RegClass;
}

bool W65816LowerWide32::runOnMachineFunction(MachineFunction &MF) {
  MachineRegisterInfo &MRI = MF.getRegInfo();
  const W65816Subtarget &STI = MF.getSubtarget<W65816Subtarget>();
  const W65816InstrInfo &TII = *STI.getInstrInfo();

  // Map: Wide32 vreg -> (loVreg, hiVreg) of i16 type.
  DenseMap<Register, std::pair<Register, Register>> wideMap;

  // Pass 1: collect all Wide32 vregs.
  SmallVector<Register, 16> wide32Vregs;
  for (unsigned i = 0, e = MRI.getNumVirtRegs(); i != e; ++i) {
    Register R = Register::index2VirtReg(i);
    if (MRI.reg_nodbg_empty(R))
      continue;
    if (!isWide32RC(MRI.getRegClass(R)))
      continue;
    wide32Vregs.push_back(R);
  }

  if (wide32Vregs.empty())
    return false;

  // Pass 2: process REG_SEQUENCE / chained-COPY / multi-subreg-def
  // shapes; build the mapping.  Iterate to fixed point because COPY
  // chains depend on prior mappings.
  SmallVector<MachineInstr *, 16> toErase;
  bool changed = true;
  while (changed) {
    changed = false;
    for (Register W : wide32Vregs) {
      if (wideMap.count(W))
        continue;
      MachineInstr *DefMI = MRI.getUniqueVRegDef(W);
      if (DefMI && DefMI->getOpcode() == TargetOpcode::REG_SEQUENCE) {
        Register Lo, Hi;
        bool Bail = false;
        for (unsigned op = 1; op + 1 < DefMI->getNumOperands(); op += 2) {
          if (!DefMI->getOperand(op).isReg() ||
              !DefMI->getOperand(op + 1).isImm())
            continue;
          unsigned idx = DefMI->getOperand(op + 1).getImm();
          Register Src = DefMI->getOperand(op).getReg();
          unsigned SrcSub = DefMI->getOperand(op).getSubReg();
          // If the source has a sub-register specifier (e.g.
          // `%W.sub_lo:wide32` is a slice of a wide32 vreg), the
          // effective "half" is the corresponding half of that source.
          // Resolve via wideMap when the parent is already mapped;
          // otherwise defer until a later iteration picks it up.
          if (SrcSub != 0) {
            if (!Src.isVirtual() || !wideMap.count(Src)) {
              Bail = true;
              break;
            }
            auto [SrcLo, SrcHi] = wideMap[Src];
            Src = (SrcSub == llvm::sub_lo) ? SrcLo : SrcHi;
          }
          if (idx == llvm::sub_lo)
            Lo = Src;
          else if (idx == llvm::sub_hi)
            Hi = Src;
        }
        if (Bail) continue;
        if (Lo && Hi) {
          wideMap[W] = {Lo, Hi};
          toErase.push_back(DefMI);
          changed = true;
          continue;
        }
      }
      if (DefMI && DefMI->isCopy()) {
        Register Src = DefMI->getOperand(1).getReg();
        if (Src.isVirtual() && wideMap.count(Src)) {
          // Copy the source mapping into a local before writing wideMap[W]:
          // DenseMap may rehash when inserting W, invalidating any reference
          // returned by operator[](Src).  Naive `wideMap[W] = wideMap[Src]`
          // dangles after rehash and writes zeros into the new entry.
          auto SrcPair = wideMap[Src];
          wideMap[W] = SrcPair;
          toErase.push_back(DefMI);
          changed = true;
          continue;
        }
      }
      // Multi-subreg-def shape: separate sub-reg COPYs build %W:
      //   undef %W.sub_lo:wide32 = COPY %A:acc16
      //   %W.sub_hi:wide32 = COPY %B:acc16
      // Equivalent to a REG_SEQUENCE %A, sub_lo, %B, sub_hi.  softDouble
      // at -O2 generates this heavily; without handling it the Wide32
      // vreg survives to regalloc, which then asks for a spill/reload
      // from a non-pair physreg and trips load/storeRegToStackSlot's
      // llvm_unreachable.
      Register LoSrc, HiSrc;
      MachineInstr *LoDefMI = nullptr;
      MachineInstr *HiDefMI = nullptr;
      bool ok = true;
      bool Bail = false;
      for (MachineInstr &MI : MRI.def_instructions(W)) {
        if (!MI.isCopy()) { ok = false; break; }
        const MachineOperand &Dst = MI.getOperand(0);
        const MachineOperand &Src = MI.getOperand(1);
        if (!Dst.isReg() || Dst.getReg() != W) { ok = false; break; }
        unsigned SubIdx = Dst.getSubReg();
        Register S = Src.getReg();
        unsigned SrcSub = Src.getSubReg();
        // If the source has a sub-register specifier, resolve through
        // wideMap[parent].  Symmetric with the REG_SEQUENCE handler
        // above — without this, `%W.sub_lo = COPY %V.sub_lo:wide32`
        // records the wide32 parent %V instead of %V's i16 sub_lo.
        if (SrcSub != 0) {
          if (!S.isVirtual() || !wideMap.count(S)) { Bail = true; break; }
          auto [SL, SH] = wideMap[S];
          S = (SrcSub == llvm::sub_lo) ? SL : SH;
        }
        if (SubIdx == llvm::sub_lo) {
          if (LoDefMI) { ok = false; break; }
          LoDefMI = &MI;
          LoSrc = S;
        } else if (SubIdx == llvm::sub_hi) {
          if (HiDefMI) { ok = false; break; }
          HiDefMI = &MI;
          HiSrc = S;
        } else {
          ok = false;
          break;
        }
      }
      if (Bail) continue;
      if (ok && LoSrc && HiSrc) {
        wideMap[W] = {LoSrc, HiSrc};
        if (LoDefMI) toErase.push_back(LoDefMI);
        if (HiDefMI) toErase.push_back(HiDefMI);
        changed = true;
      }
    }
  }

  // Pass 2b: handle PHIs whose result is a Wide32 vreg by splitting
  // into 2 PHIs (one per half).  Iterate to fixed point: a PHI becomes
  // resolvable only after all its sources have been mapped.
  changed = true;
  while (changed) {
    changed = false;
    for (Register W : wide32Vregs) {
      if (wideMap.count(W))
        continue;
      MachineInstr *DefMI = MRI.getUniqueVRegDef(W);
      if (!DefMI || !DefMI->isPHI())
        continue;
      bool AllMapped = true;
      for (unsigned op = 1; op + 1 < DefMI->getNumOperands(); op += 2) {
        Register Src = DefMI->getOperand(op).getReg();
        if (!Src.isVirtual() || !wideMap.count(Src)) {
          AllMapped = false;
          break;
        }
      }
      if (!AllMapped)
        continue;
      Register NewLo = MRI.createVirtualRegister(&W65816::Acc16RegClass);
      Register NewHi = MRI.createVirtualRegister(&W65816::Acc16RegClass);
      MachineBasicBlock *MBB = DefMI->getParent();
      DebugLoc DL = DefMI->getDebugLoc();
      auto PHILo = BuildMI(*MBB, DefMI, DL, TII.get(TargetOpcode::PHI), NewLo);
      auto PHIHi = BuildMI(*MBB, DefMI, DL, TII.get(TargetOpcode::PHI), NewHi);
      for (unsigned op = 1; op + 1 < DefMI->getNumOperands(); op += 2) {
        Register Src = DefMI->getOperand(op).getReg();
        MachineBasicBlock *PredMBB = DefMI->getOperand(op + 1).getMBB();
        auto [SrcLo, SrcHi] = wideMap[Src];
        PHILo.addReg(SrcLo).addMBB(PredMBB);
        PHIHi.addReg(SrcHi).addMBB(PredMBB);
      }
      wideMap[W] = {NewLo, NewHi};
      toErase.push_back(DefMI);
      changed = true;
    }
  }

  // Pass 3: rewrite uses.
  SmallVector<MachineInstr *, 32> useToErase;
  // Hand a half to the users of the vreg projected out of the pair, Dst.
  // A user whose operand class does not take the half's own class gets a
  // COPY of it into that class right before it, as TwoAddress does for a
  // tied operand (so tied uses are left to it).  Replacing unconditionally
  // gave an Img16 half (an argument's high word arrives in X and lives in a
  // DP slot) to CMP, whose register operand is implicitly A: the compare
  // read whatever A held.  The COPY goes at the user rather than at the
  // projection so that A, a single register, is held only across the use.
  //
  // wideMap was filled in Pass 2 and is not updated by the replacements.  A
  // pair's half can itself be projected out of another pair (`%h = COPY
  // %p.sub_hi`, then `%q = REG_SEQUENCE %h, ...`, as an i32 shift by 16
  // builds); once %h is replaced and its COPY erased, handing out %h for
  // %q's half would leave uses with no def.  So every half read from
  // wideMap goes through `current`, which follows the replacements.
  const TargetRegisterInfo *TRI = STI.getRegisterInfo();
  DenseMap<Register, Register> forwarded;
  auto current = [&](Register R) {
    for (auto It = forwarded.find(R); It != forwarded.end();
         It = forwarded.find(R))
      R = It->second;
    return R;
  };
  auto forwardHalf = [&](MachineInstr *MI, Register Dst, Register Half) {
    Half = current(Half);
    const TargetRegisterClass *HalfRC = MRI.getRegClass(Half);
    for (MachineOperand &MO :
         make_early_inc_range(MRI.use_nodbg_operands(Dst))) {
      MachineInstr *User = MO.getParent();
      unsigned OpNo = User->getOperandNo(&MO);
      const TargetRegisterClass *Need =
          User->getRegClassConstraint(OpNo, &TII, TRI);
      if (!Need || Need->hasSubClassEq(HalfRC) ||
          User->isRegTiedToDefOperand(OpNo))
        continue;
      Register Fit = MRI.createVirtualRegister(Need);
      BuildMI(*User->getParent(), User, User->getDebugLoc(),
              TII.get(TargetOpcode::COPY), Fit)
          .addReg(Half);
      MO.setReg(Fit);
    }
    MRI.replaceRegWith(Dst, Half);
    forwarded[Dst] = Half;
    useToErase.push_back(MI);
  };
  for (MachineBasicBlock &MBB : MF) {
    for (MachineBasicBlock::iterator It = MBB.begin(); It != MBB.end();) {
      MachineInstr *MI = &*It++;

      // EXTRACT_SUBREG of a mapped Wide32 vreg: replace the dest vreg
      // with the appropriate half (sub_lo or sub_hi).
      if (MI->getOpcode() == TargetOpcode::EXTRACT_SUBREG) {
        Register Src = MI->getOperand(1).getReg();
        if (Src.isVirtual() && wideMap.count(Src)) {
          unsigned SubIdx = MI->getOperand(2).getImm();
          auto [Lo, Hi] = wideMap[Src];
          Register Half = (SubIdx == llvm::sub_lo) ? Lo : Hi;
          forwardHalf(MI, MI->getOperand(0).getReg(), Half);
          continue;
        }
      }

      // COPY %V.sub_lo / %V.sub_hi (partial-reg COPY where source has a
      // sub-reg specifier and the source vreg is a mapped Wide32).
      // LLVM emits this shape instead of EXTRACT_SUBREG when projecting
      // a half out of a Wide32 vreg.  Only the shape with a full-reg
      // destination is handled here — partial-reg destinations would
      // imply the dst is itself a Wide32 sub-reg def, which the def-side
      // multi-subreg-def handling covers separately.
      if (MI->isCopy()) {
        const MachineOperand &SrcOp = MI->getOperand(1);
        const MachineOperand &DstOp = MI->getOperand(0);
        if (SrcOp.isReg() && SrcOp.getReg().isVirtual() &&
            wideMap.count(SrcOp.getReg()) && SrcOp.getSubReg() != 0 &&
            DstOp.isReg() && DstOp.getSubReg() == 0) {
          unsigned SubIdx = SrcOp.getSubReg();
          auto [Lo, Hi] = wideMap[SrcOp.getReg()];
          Register Half = (SubIdx == llvm::sub_lo) ? Lo : Hi;
          forwardHalf(MI, DstOp.getReg(), Half);
          continue;
        }
      }

      // LDAptr32 / STAptr32 / STBptr32 with a mapped Wide32 ptr:
      // rewrite to LDAptr32S / STAptr32S / STBptr32S.
      unsigned Opc = MI->getOpcode();
      bool isPtrOp = (Opc == W65816::LDAptr32 || Opc == W65816::STAptr32 ||
                      Opc == W65816::STBptr32);
      if (isPtrOp) {
        Register Ptr = MI->getOperand(1).getReg();
        if (Ptr.isVirtual() && wideMap.count(Ptr)) {
          Register Lo = current(wideMap[Ptr].first);
          Register Hi = current(wideMap[Ptr].second);
          unsigned NewOpc = (Opc == W65816::LDAptr32) ? W65816::LDAptr32S
                          : (Opc == W65816::STAptr32) ? W65816::STAptr32S
                                                      : W65816::STBptr32S;
          DebugLoc DL = MI->getDebugLoc();
          MachineBasicBlock *ParentMBB = MI->getParent();
          if (Opc == W65816::LDAptr32) {
            Register Dst = MI->getOperand(0).getReg();
            BuildMI(*ParentMBB, MI->getIterator(), DL, TII.get(NewOpc), Dst)
                .addReg(Lo)
                .addReg(Hi);
          } else {
            Register Val = MI->getOperand(0).getReg();
            BuildMI(*ParentMBB, MI->getIterator(), DL, TII.get(NewOpc))
                .addReg(Val)
                .addReg(Lo)
                .addReg(Hi);
          }
          useToErase.push_back(MI);
          continue;
        }
      }
    }
  }

  // Erase use-side instructions (EXTRACT_SUBREG, LDAptr32-family) first
  // so the Wide32 vreg becomes dead.
  //
  // DBG_VALUE preservation: MRI.replaceRegWith above already redirected
  // any debug operands to the new half-vreg (replaceRegWith walks
  // reg_operands which includes debug ops).  No further work needed for
  // the EXTRACT_SUBREG / partial-COPY paths since the user-visible value
  // is preserved in the half vreg.  For LDAptr32 rewrites we don't move
  // a user-named value so DBG_VALUE attribution is unaffected.
  for (auto *MI : useToErase)
    MI->eraseFromParent();

  // Now check each REG_SEQUENCE / chained-COPY def: only erase if the
  // Wide32 vreg has no remaining uses.  Any leftover use means the pass
  // didn't cover that opcode — leaving the def in place keeps the MIR
  // well-formed (at the cost of pair-allocation pressure for that
  // specific case).
  //
  // Iterate to fixed point: a chained-COPY pattern like
  //   %114:wide32 = REG_SEQUENCE ...
  //   %74:wide32 = COPY %114
  //   ... uses of %74 ...
  // queues both the REG_SEQUENCE and the COPY for erasure.  Pass 3
  // rewrites %74's uses, leaving %74 dead.  In a single-pass erase,
  // %114 still has its COPY use at the time we check, so the REG_-
  // SEQUENCE is skipped — but then we erase the COPY, leaving %114
  // dead too.  Loop until no more erasures.
  bool eraseAny = !useToErase.empty();
  bool progress = true;
  while (progress) {
    progress = false;
    for (auto *&MI : toErase) {
      if (!MI)
        continue;
      if (MI->getNumOperands() == 0) {
        MI = nullptr;
        continue;
      }
      Register Dst = MI->getOperand(0).getReg();
      if (!Dst.isVirtual() || MRI.use_nodbg_empty(Dst)) {
        // DBG_VALUE preservation: use_nodbg_empty ignores debug uses, so
        // erasing here can strand DBG_VALUEs referring to Dst.  The
        // wide32 vreg has no surviving half mapping that's directly
        // usable (a wide32-typed variable is implemented as two i16
        // halves but DBG_VALUE expects a single location), so mark
        // dependent DBG_VALUEs as undef rather than emit a wrong
        // location.  A future pass could split the variable across
        // DW_OP_piece(lo)+DW_OP_piece(hi); not in scope here.
        if (Dst.isVirtual()) {
          // Dedup: DBG_VALUE_LIST may appear multiple times in the
          // use-list when several of its operands reference Dst.
          SmallPtrSet<MachineInstr *, 4> DbgUserSet;
          SmallVector<MachineInstr *, 4> DbgUsers;
          for (MachineInstr &U : MRI.use_instructions(Dst)) {
            if (U.isDebugValue() && DbgUserSet.insert(&U).second)
              DbgUsers.push_back(&U);
          }
          for (MachineInstr *DbgMI : DbgUsers) {
            for (MachineOperand &Op : DbgMI->debug_operands()) {
              if (Op.isReg() && Op.getReg() == Dst) {
                Op.setReg(0);
                Op.setSubReg(0);
              }
            }
          }
        }
        MI->eraseFromParent();
        MI = nullptr;
        eraseAny = true;
        progress = true;
      }
    }
  }

  return eraseAny;
}