//===-- W65816BranchExpand.cpp - Long conditional branch expansion --------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// Lengthens conditional branches that target an MBB further than +/-128
// bytes away.  The 65816 has BRL (signed-16, ±32K) for unconditional
// branches but no long *conditional* branch, so we expand
//
//     Bxx Target              -->     INV_Bxx Skip
//     fall-through Skip               BRA     Target
//                                 Skip:
//                                     fall-through
//
// The unconditional BRA is later auto-relaxed to BRL by W65816AsmBackend
// when its displacement exceeds 8 bits (in the same way that an
// assembler-time `bra label` to a label > 127 bytes away gets promoted).
//
// Algorithm:
//
//   1.  Pre-split: any MBB that has more than one conditional terminator
//       (the multi-branch SELECT_CC pattern emits two Bxx in one MBB)
//       is sliced after every conditional Bxx that isn't the LAST one.
//       After this, each MBB has at most one conditional terminator,
//       which my expansion logic can handle cleanly.
//
//   2.  Iterate to fixed-point.  In each iteration, recompute byte
//       distances (using TII::getInstSizeInBytes for accuracy) and
//       expand every conditional whose target is more than
//       EXPAND_DIST_THRESHOLD bytes away.  Each expansion adds 3 bytes
//       (the Bridge MBB's BRA), which can push another inner branch
//       over the threshold; iterate until no further expansions.
//
// Runs at addPreEmitPass, after PEI so all FrameIndex references and
// pseudo expansions have stable byte sizes.
//
//===----------------------------------------------------------------------===//

#include "W65816.h"
#include "W65816InstrInfo.h"
#include "W65816Subtarget.h"
#include "llvm/CodeGen/LivePhysRegs.h"
#include "llvm/CodeGen/MachineFunction.h"
#include "llvm/CodeGen/MachineFunctionPass.h"
#include "llvm/CodeGen/MachineInstr.h"
#include "llvm/CodeGen/MachineInstrBuilder.h"
#include "llvm/Support/raw_ostream.h"

using namespace llvm;

#define DEBUG_TYPE "w65816-branch-expand"

namespace {

class W65816BranchExpand : public MachineFunctionPass {
public:
  static char ID;
  W65816BranchExpand() : MachineFunctionPass(ID) {}
  StringRef getPassName() const override {
    return "W65816 conditional branch expansion";
  }
  bool runOnMachineFunction(MachineFunction &MF) override;
};

} // namespace

char W65816BranchExpand::ID = 0;

INITIALIZE_PASS(W65816BranchExpand, DEBUG_TYPE,
                "W65816 conditional branch expansion", false, false)

FunctionPass *llvm::createW65816BranchExpand() {
  return new W65816BranchExpand();
}

// Map a conditional branch opcode to its inverted form via the shared
// helper in W65816InstrInfo.h.  Returns 0 if not a recognised conditional Bxx.
static unsigned invertedConditional(unsigned Opc) {
  return W65816Helpers::invertCondOpcode(Opc);
}

// Byte-accurate distance estimate from a specific branch instruction
// to its target MBB.  Starts counting at the BRANCH (not at the MBB
// start) and stops at the target MBB's start.  This matters because a
// branch at the END of a large MBB has a tiny actual distance to the
// next-laid-out MBB even though the MBB itself is huge.
static unsigned estimateDistance(MachineFunction &MF,
                                 const TargetInstrInfo *TII,
                                 const MachineInstr &Br,
                                 MachineBasicBlock *To) {
  const MachineBasicBlock *From = Br.getParent();
  // Self-loop branch: target is the start of From, branch is somewhere
  // inside From.  Distance is the bytes from start of From to the
  // branch instruction (i.e., everything before Br in From).
  if (From == To) {
    unsigned Bytes = 0;
    for (const auto &MI : *From) {
      if (&MI == &Br) break;
      Bytes += TII->getInstSizeInBytes(MI);
    }
    return Bytes;
  }

  // Two cases by layout direction:
  //   forward: bytes after Br in From, plus all of MBBs strictly
  //            between, plus 0 (branch lands at To's start).
  //   backward: bytes before Br in From, plus all of MBBs strictly
  //             between, plus all of To.
  int FromIdx = -1, ToIdx = -1, Idx = 0;
  for (auto &MBB : MF) {
    if (&MBB == From) FromIdx = Idx;
    if (&MBB == To)   ToIdx   = Idx;
    Idx++;
  }
  if (FromIdx < 0 || ToIdx < 0) return 1000;  // unknown — assume far

  unsigned Bytes = 0;
  if (ToIdx > FromIdx) {
    // Forward: count from Br to end of From, then between, then 0.
    bool past = false;
    for (const auto &MI : *From) {
      if (&MI == &Br) past = true;
      if (past) Bytes += TII->getInstSizeInBytes(MI);
    }
    Idx = 0;
    for (auto &MBB : MF) {
      if (Idx > FromIdx && Idx < ToIdx)
        for (const auto &MI : MBB) Bytes += TII->getInstSizeInBytes(MI);
      Idx++;
    }
  } else {
    // Backward: count Br's preceding bytes in From, plus between, plus all of To.
    for (const auto &MI : *From) {
      if (&MI == &Br) break;
      Bytes += TII->getInstSizeInBytes(MI);
    }
    Idx = 0;
    for (auto &MBB : MF) {
      if (Idx > ToIdx && Idx < FromIdx)
        for (const auto &MI : MBB) Bytes += TII->getInstSizeInBytes(MI);
      if (Idx == ToIdx)
        for (const auto &MI : MBB) Bytes += TII->getInstSizeInBytes(MI);
      Idx++;
    }
  }
  return Bytes;
}

// Step 1 — pre-split: any MBB with > 1 conditional terminator gets
// sliced after each non-final conditional, so every MBB ends up with
// at most one conditional terminator.  Returns true if any MBB was
// split.
static bool splitMultiBranchMBBs(MachineFunction &MF,
                                 SmallVectorImpl<MachineBasicBlock *> &NewBlocks) {
  bool Changed = false;
  // Snapshot MBBs first (we mutate the list during iteration).
  SmallVector<MachineBasicBlock *, 16> MBBs;
  for (auto &MBB : MF) MBBs.push_back(&MBB);

  for (MachineBasicBlock *MBB : MBBs) {
    // Find the first conditional terminator that has another
    // conditional terminator after it.  Slice MBB right after it.
    bool Sliced = true;
    while (Sliced) {
      Sliced = false;
      // Walk terminators forward.
      auto firstTerm = MBB->getFirstTerminator();
      MachineBasicBlock::iterator splitAfter = MBB->end();
      MachineBasicBlock::iterator firstCond  = MBB->end();
      for (auto it = firstTerm; it != MBB->end(); ++it) {
        if (invertedConditional(it->getOpcode()) != 0) {
          if (firstCond == MBB->end()) {
            firstCond = it;
          } else {
            splitAfter = firstCond;  // split AFTER this earlier conditional
            break;
          }
        }
      }
      if (splitAfter == MBB->end()) break;

      // Create new MBB; transfer everything after splitAfter to it.
      auto *NewMBB = MF.CreateMachineBasicBlock(MBB->getBasicBlock());
      MF.insert(std::next(MBB->getIterator()), NewMBB);
      NewBlocks.push_back(NewMBB);
      // Move instructions [splitAfter+1 .. end) to NewMBB.
      auto moveStart = std::next(splitAfter);
      NewMBB->splice(NewMBB->end(), MBB, moveStart, MBB->end());
      // Transfer successors that aren't the splitAfter's target.
      MachineBasicBlock *splitTgt = nullptr;
      if (splitAfter->getNumOperands() >= 1 &&
          splitAfter->getOperand(0).isMBB())
        splitTgt = splitAfter->getOperand(0).getMBB();
      // All of MBB's existing successors that aren't splitTgt move to
      // NewMBB.  splitTgt stays as MBB's own successor (the conditional
      // branch target).  EXCEPTION: if any branch instruction we moved
      // into NewMBB *also* targets splitTgt (the multi-branch SELECT_CC
      // case where both Bxx point at the same MBB), splitTgt must also
      // be a successor of NewMBB.
      SmallVector<MachineBasicBlock *, 4> OldSuccs(MBB->successors().begin(),
                                                    MBB->successors().end());
      for (auto *S : OldSuccs) {
        if (S == splitTgt) continue;
        MBB->removeSuccessor(S);
        NewMBB->addSuccessor(S);
      }
      // Walk NewMBB's instructions; for each MBB-operand reference,
      // ensure that target is a NewMBB successor.
      for (auto &MI : *NewMBB) {
        for (unsigned i = 0; i < MI.getNumOperands(); ++i) {
          const auto &OP = MI.getOperand(i);
          if (!OP.isMBB()) continue;
          auto *RefMBB = OP.getMBB();
          if (!NewMBB->isSuccessor(RefMBB))
            NewMBB->addSuccessor(RefMBB);
        }
      }
      // MBB falls through to NewMBB now.
      MBB->addSuccessor(NewMBB);
      // The splitAfter conditional already targets splitTgt (still in
      // MBB->successors()).  Done — process the same MBB again to
      // see if another split is needed (multi-multi-branch case).
      Changed = true;
      Sliced = true;
    }
  }
  return Changed;
}

// Drop conditional branches whose target matches the unconditional
// branch immediately following them (both edges go to the same MBB,
// so the conditional is dead).  This pattern survives upstream cleanup
// when the branches were emitted by the W65816 SELECT_CC inserter or
// by codegenprepare on an `br i1 %c, label %X, label %X` IR shape.
// Returns true if any MI was deleted.
// Invert a conditional branch when its taken-target equals the
// immediately-following MBB in layout order.  Pattern:
//
//   MBB ends with:
//     Bxx CondTarget
//     BRA UncondTarget
//   next-laid-out MBB == CondTarget
//
// Rewrite to (no BRA, fall through to CondTarget which is the next MBB):
//
//     Bxx_inverted UncondTarget
//     <fall through to CondTarget>
//
// Common pattern: tight loops `beq exit ; bra loop` where the exit
// block is placed right after the loop body.  After this, the backedge
// becomes a single `bne loop` (3 cyc taken) instead of `beq exit + bra
// loop` (6 cyc when looping).  sumOfSquares hits this 50× per call.
static bool invertCondToFallThrough(MachineFunction &MF,
                                    const TargetInstrInfo *TII) {
  bool Changed = false;
  auto MBBIt = MF.begin();
  while (MBBIt != MF.end()) {
    MachineBasicBlock &MBB = *MBBIt;
    auto NextMBB = std::next(MBBIt);
    ++MBBIt;
    if (NextMBB == MF.end()) {
      continue;
    }
    auto T = MBB.getFirstTerminator();
    if (T == MBB.end()) {
      continue;
    }
    unsigned CondOpc = T->getOpcode();
    unsigned InvOpc = invertedConditional(CondOpc);
    if (InvOpc == 0) {
      continue;
    }
    if (T->getNumOperands() < 1 || !T->getOperand(0).isMBB()) {
      continue;
    }
    MachineBasicBlock *CondTarget = T->getOperand(0).getMBB();
    auto N = std::next(T);
    if (N == MBB.end()) {
      continue;
    }
    unsigned UncondOpc = N->getOpcode();
    if (UncondOpc != W65816::BRA && UncondOpc != W65816::BRL) {
      continue;
    }
    if (N->getNumOperands() < 1 || !N->getOperand(0).isMBB()) {
      continue;
    }
    MachineBasicBlock *UncondTarget = N->getOperand(0).getMBB();
    if (CondTarget != &*NextMBB) {
      continue;
    }
    if (CondTarget == UncondTarget) {
      continue;
    }
    T->setDesc(TII->get(InvOpc));
    T->getOperand(0).setMBB(UncondTarget);
    N->eraseFromParent();
    Changed = true;
  }
  return Changed;
}


static bool dropDeadConditionalsToBRATarget(MachineFunction &MF) {
  bool Changed = false;
  for (auto &MBB : MF) {
    auto T = MBB.getFirstTerminator();
    while (T != MBB.end()) {
      auto Next = std::next(T);
      if (Next == MBB.end()) break;
      unsigned CondOpc = T->getOpcode();
      if (invertedConditional(CondOpc) == 0) { ++T; continue; }
      unsigned UncondOpc = Next->getOpcode();
      if (UncondOpc != W65816::BRA && UncondOpc != W65816::BRL) {
        ++T; continue;
      }
      if (T->getNumOperands() < 1 || !T->getOperand(0).isMBB()) { ++T; continue; }
      if (Next->getNumOperands() < 1 || !Next->getOperand(0).isMBB()) { ++T; continue; }
      if (T->getOperand(0).getMBB() != Next->getOperand(0).getMBB()) { ++T; continue; }
      // Conditional and unconditional target the same MBB.  Drop the
      // conditional; the unconditional already covers both edges.
      auto Erase = T++;
      Erase->eraseFromParent();
      Changed = true;
    }
  }
  return Changed;
}

bool W65816BranchExpand::runOnMachineFunction(MachineFunction &MF) {
  const auto &STI = MF.getSubtarget<W65816Subtarget>();
  const auto *TII = STI.getInstrInfo();
  bool AnyChanged = false;
  // Blocks created below (split tails + `BRL Target` bridges).  These are
  // wired into the CFG with successors but no live-in lists; we recompute
  // live-ins for exactly these blocks at the end so downstream passes that
  // query successor live-ins (W65816SepRepCleanup) stay sound.  Existing
  // blocks' live-ins are unaffected by inserting predecessors, so there's
  // no need to recompute the whole function.
  SmallVector<MachineBasicBlock *, 8> NewBlocks;

  // Step -1: invert Bxx + BRA when cond target is the next MBB.  Saves
  // 3 cyc per backedge in tight loops (sumOfSquares 50 iters: -150 cyc).
  AnyChanged |= invertCondToFallThrough(MF, TII);

  // Step 0: drop dead conditionals (Bxx X immediately followed by BRA X
  // — both edges to the same MBB).  Cheap and removes false-positive
  // candidates from the distance-based expansion below.
  AnyChanged |= dropDeadConditionalsToBRATarget(MF);

  // Step 1: split multi-conditional-terminator MBBs.
  AnyChanged |= splitMultiBranchMBBs(MF, NewBlocks);

  // Step 2: iterate to fixed-point.  Each expansion adds 3 bytes
  // (bridge BRA), which may push another previously-OK branch over
  // the threshold.  Cap at MAX_ITERS to avoid pathological cases.
  const unsigned EXPAND_DIST_THRESHOLD = 90;  // tighter margin under +/-128
  const unsigned MAX_ITERS = 10;
  for (unsigned iter = 0; iter < MAX_ITERS; ++iter) {
    bool Changed = false;

    // Promote long BRA to BRL.  The assembler's BRA→BRL relaxation
    // sometimes fails to fire when the target symbol resolves early
    // in MC layout — the linker then sees a PCREL8 reloc that's out
    // of range.  Force the BRL ourselves when the estimate exceeds
    // the safe threshold; saves one byte if BRA would have fit, but
    // beats a hard link error.
    for (auto &MBB : MF) {
      for (auto &MI : MBB.terminators()) {
        if (MI.getOpcode() != W65816::BRA) continue;
        if (MI.getNumOperands() < 1 || !MI.getOperand(0).isMBB()) continue;
        MachineBasicBlock *Target = MI.getOperand(0).getMBB();
        unsigned Dist = estimateDistance(MF, TII, MI, Target);
        if (Dist > EXPAND_DIST_THRESHOLD) {
          MI.setDesc(TII->get(W65816::BRL));
          Changed = true;
        }
      }
    }

    // Collect candidates.  After step 1, each MBB has at most one
    // conditional terminator, so we walk terminators().
    SmallVector<std::pair<MachineBasicBlock *, MachineInstr *>, 8> Candidates;
    for (auto &MBB : MF) {
      for (auto &MI : MBB.terminators()) {
        unsigned Opc = MI.getOpcode();
        if (invertedConditional(Opc) == 0) continue;
        if (MI.getNumOperands() < 1 || !MI.getOperand(0).isMBB()) continue;
        MachineBasicBlock *Target = MI.getOperand(0).getMBB();
        unsigned Dist = estimateDistance(MF, TII, MI, Target);
        if (Dist > EXPAND_DIST_THRESHOLD)
          Candidates.emplace_back(&MBB, &MI);
      }
    }

    for (auto [MBB, BrMI] : Candidates) {
      unsigned Opc = BrMI->getOpcode();
      unsigned InvOpc = invertedConditional(Opc);
      MachineBasicBlock *Target = BrMI->getOperand(0).getMBB();
      DebugLoc DL = BrMI->getDebugLoc();

      // Layout transformation:
      //   MBB:    ... ; Bxx Target ; (fall-through Skip)
      // Becomes:
      //   MBB:    ... ; INV_Bxx Skip
      //   Bridge:       BRA Target
      //   Skip:   (= original MBB's fall-through successor)
      //
      // After splitMultiBranchMBBs, MBB has ONE conditional terminator
      // (BrMI) and at most one unconditional terminator after it (which
      // we leave alone — it's the fall-through-or-explicit branch).
      // MBB's successors are {Target, Skip} where Skip is whichever
      // is not Target.
      MachineBasicBlock *Skip = nullptr;
      for (auto *S : MBB->successors()) {
        if (S != Target) { Skip = S; break; }
      }
      if (!Skip) continue;  // function-end conditional — rare; skip

      // Create Bridge MBB.
      MachineBasicBlock *Bridge =
          MF.CreateMachineBasicBlock(MBB->getBasicBlock());
      MF.insert(std::next(MBB->getIterator()), Bridge);
      NewBlocks.push_back(Bridge);

      // Replace successor edges: MBB used to have {Target, Skip}; now
      // it has {Bridge, Skip}.  Bridge has {Target}.
      MBB->removeSuccessor(Target);
      MBB->addSuccessor(Bridge);
      Bridge->addSuccessor(Target);

      // Erase original Bxx, emit inverted Bxx targeting Skip.
      BrMI->eraseFromParent();
      // Insert at MBB's terminator position so any unconditional
      // fall-through marker after stays after.
      auto insertPt = MBB->getFirstTerminator();
      BuildMI(*MBB, insertPt, DL, TII->get(InvOpc)).addMBB(Skip);
      // After the rewrite, MBB falls through to Bridge (which now sits
      // immediately after MBB in layout).  Any unconditional BRA/BRL
      // already at the end of MBB used to direct the fall-through to
      // Skip — but with Bridge interposed, that BRA would skip past
      // Bridge entirely and Bridge becomes unreachable.  Remove it.
      // (Skip is still reachable via INV_Bxx; Target is reachable via
      // fall-through-to-Bridge then BRL.)  Caught by vprintf crashing
      // because dropDeadConditionalsToBRATarget then dropped the
      // INV_Bxx as redundant with the leftover BRA Skip.
      while (insertPt != MBB->end()) {
        unsigned NextOpc = insertPt->getOpcode();
        if (NextOpc == W65816::BRA || NextOpc == W65816::BRL) {
          if (insertPt->getNumOperands() >= 1 &&
              insertPt->getOperand(0).isMBB() &&
              insertPt->getOperand(0).getMBB() == Skip) {
            insertPt = insertPt->eraseFromParent();
            continue;
          }
        }
        ++insertPt;
      }

      // Bridge: BRL Target.  Always emit the long form rather than
      // relying on the assembler to relax BRA→BRL — the relaxation
      // path is fragile in mixed-fragment scenarios (MC layout
      // doesn't always re-evaluate after layout shifts) and we'd
      // rather pay 1 extra byte per long branch than risk a silent
      // PCREL8 fixup that can't be resolved at link time.
      BuildMI(Bridge, DL, TII->get(W65816::BRL)).addMBB(Target);

      Changed = true;
    }
    AnyChanged = AnyChanged || Changed;
    if (!Changed) break;
  }

  // Step 3: re-run the dead-conditional sweep.  Expansion introduces
  // `INV_Bxx Skip ; BRA Target` pairs; when the original codegen
  // already had `BRA Skip` after the (now-erased) Bxx, those collapse
  // into `INV_Bxx X ; BRA X` — the conditional is dead.
  AnyChanged |= dropDeadConditionalsToBRATarget(MF);

  // Step 4: drop trailing `BRA next_MBB` / `BRL next_MBB` when the
  // target is the immediately-following layout MBB.  Block-placement
  // sometimes leaves these as explicit branches even though
  // fall-through suffices.  Saves 3 bytes / 3 cycles each.
  for (auto MBBIt = MF.begin(); MBBIt != MF.end(); ++MBBIt) {
    auto NextMBB = std::next(MBBIt);
    if (NextMBB == MF.end()) continue;
    auto Last = MBBIt->getLastNonDebugInstr();
    if (Last == MBBIt->end()) continue;
    unsigned Op = Last->getOpcode();
    if (Op != W65816::BRA && Op != W65816::BRL) continue;
    if (Last->getNumOperands() < 1 || !Last->getOperand(0).isMBB()) continue;
    if (Last->getOperand(0).getMBB() != &*NextMBB) continue;
    Last->eraseFromParent();
    AnyChanged = true;
  }

  // Bridge/split blocks created above are wired into the CFG with successors
  // but no live-in lists, so they are left with EMPTY live-ins even though
  // values flow through them to their targets.  Downstream peepholes that
  // query successor live-in sets (W65816SepRepCleanup's PHI-copy hoist and
  // lagged-ptr sink guards both call Succ->isLiveIn(A)) would then wrongly
  // conclude A is dead at a block exit and hoist a clobbering def across a
  // still-live value.  This is the root cause of the malloc over-heap
  // miscompile: the over-heap NULL return (A=0) lived out only through a
  // Bridge block, so the hoist replaced it with the bounds-compare result.
  // Recompute live-ins for exactly the new blocks (their successors are
  // existing blocks whose live-ins are already correct; fullyRecomputeLiveIns
  // converges across chained new blocks).  Existing blocks are untouched, so
  // this stays O(new blocks) instead of O(function) per invocation.
  if (!NewBlocks.empty())
    fullyRecomputeLiveIns(NewBlocks);
  return AnyChanged;
}
