//===-- W65816UnLSR.cpp - Undo LSR for global-array pointer-walks --------===//
//
// 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-LSR IR pass.  LSR converts `for (i; i<n; i++) arr[i]` loops into
// pointer-walking:
//
//   %lsr.iv1 = phi ptr [@arr, ...], [%scevgep, ...]
//   %lsr.iv  = phi i16 [%n, ...], [%lsr.iv.next, ...]
//   %v = load i16, ptr %lsr.iv1
//   %scevgep = getelementptr i8, ptr %lsr.iv1, i32 stride
//   %lsr.iv.next = add i16 %lsr.iv, -1
//
// On the W65816 this is a regression: pointer-deref needs the 24-bit
// `[dp],Y` addressing (sets up $E0/$E2 each iter), whereas indexed
// access via `lda <global>, X` is 3 bytes / 5 cyc DBR-relative.  The
// LDA_AbsX SDAG combine recognizes `(load (add Wrapper, idx))` but
// LSR has destroyed that shape.
//
// This pass detects the shape and converts back: introduce a forward
// counter PHI (0, 1, 2, ...), replace loads/stores through the pointer
// PHI with indexed GEPs `getelementptr T, @global, %i`.  The old pointer
// PHI becomes dead and is DCE'd.  The existing count-down counter PHI
// stays as the loop control.
//
// Restrictions:
//  * Pointer PHI's initial value must be a `GlobalAddress` (otherwise
//    it might point at a non-DBR-bank object, and our combine doesn't
//    fire / produces wrong code).
//  * Pointer PHI must have exactly two incoming values (header + latch).
//  * Pointer PHI's only uses must be the stride GEP and loads/stores.
//  * Stride must be a positive constant.
//  * The loop must have an existing trip-count counter so we don't
//    introduce extra register pressure.
//
//===---------------------------------------------------------------------===//

#include "W65816.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/Analysis/ScalarEvolution.h"
#include "llvm/InitializePasses.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/InstIterator.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/Module.h"
#include "llvm/Pass.h"
#include "llvm/Support/Debug.h"

using namespace llvm;

#define DEBUG_TYPE "w65816-un-lsr"

namespace llvm {
void initializeW65816UnLSRPass(PassRegistry &);
}

namespace {

class W65816UnLSR : public FunctionPass {
public:
  static char ID;
  W65816UnLSR() : FunctionPass(ID) {}

  StringRef getPassName() const override {
    return "W65816 undo LSR pointer-walk for global-array access";
  }

  void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.addRequired<LoopInfoWrapperPass>();
    AU.setPreservesCFG();
  }

  bool runOnFunction(Function &F) override;

private:
  bool processLoop(Loop *L);
};

} // namespace

char W65816UnLSR::ID = 0;

INITIALIZE_PASS_BEGIN(W65816UnLSR, DEBUG_TYPE,
                      "W65816 undo LSR for global-array", false, false)
INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
INITIALIZE_PASS_END(W65816UnLSR, DEBUG_TYPE,
                    "W65816 undo LSR for global-array", false, false)

FunctionPass *llvm::createW65816UnLSR() { return new W65816UnLSR(); }

bool W65816UnLSR::runOnFunction(Function &F) {
  if (F.hasOptNone())
    return false;
  LoopInfo &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
  bool Changed = false;
  for (Loop *L : LI) {
    Changed |= processLoop(L);
    // Recurse into nested loops.
    SmallVector<Loop *, 4> Worklist(L->begin(), L->end());
    while (!Worklist.empty()) {
      Loop *Sub = Worklist.pop_back_val();
      Changed |= processLoop(Sub);
      Worklist.append(Sub->begin(), Sub->end());
    }
  }
  return Changed;
}

bool W65816UnLSR::processLoop(Loop *L) {
  BasicBlock *Header = L->getHeader();
  BasicBlock *Latch = L->getLoopLatch();
  if (!Latch)
    return false;
  // Single-block loops are fine (header == latch).

  // Find a pointer PHI whose initial value is a GlobalAddress and whose
  // latch-incoming is a constant-stride GEP off itself.
  SmallVector<PHINode *, 4> Candidates;
  for (PHINode &PN : Header->phis()) {
    if (!PN.getType()->isPointerTy()) continue;
    if (PN.getNumIncomingValues() != 2) continue;
    Candidates.push_back(&PN);
  }

  bool Changed = false;
  for (PHINode *PtrPHI : Candidates) {
    // Identify which incoming is from the preheader vs the latch.
    Value *InitVal = nullptr;
    Value *StepVal = nullptr;
    for (unsigned i = 0; i < PtrPHI->getNumIncomingValues(); ++i) {
      BasicBlock *Pred = PtrPHI->getIncomingBlock(i);
      Value *Inc = PtrPHI->getIncomingValue(i);
      if (L->contains(Pred))
        StepVal = Inc;
      else
        InitVal = Inc;
    }
    if (!InitVal || !StepVal) continue;

    // InitVal must be a GlobalAddress (GlobalVariable).  Could relax to
    // any constant base, but globals are the only case our combine
    // handles correctly.
    auto *InitGV = dyn_cast<GlobalVariable>(InitVal);
    if (!InitGV) continue;

    // StepVal must be a `getelementptr i8, %PtrPHI, ConstantInt`.
    auto *StepGEP = dyn_cast<GetElementPtrInst>(StepVal);
    if (!StepGEP) continue;
    if (StepGEP->getPointerOperand() != PtrPHI) continue;
    if (StepGEP->getNumIndices() != 1) continue;
    auto *StrideCI = dyn_cast<ConstantInt>(StepGEP->getOperand(1));
    if (!StrideCI) continue;
    int64_t StrideBytes = StrideCI->getSExtValue();
    if (StrideBytes <= 0) continue;
    // Determine element type for the indexed GEP.  Walk all loads/
    // stores through PtrPHI; require they all access the same primitive
    // type whose size equals StrideBytes (typical: i16 array → stride 2).
    Type *ElemTy = nullptr;
    SmallVector<Instruction *, 4> Accesses;
    bool ok = true;
    for (User *U : PtrPHI->users()) {
      if (U == StepGEP) continue;
      if (auto *LD = dyn_cast<LoadInst>(U)) {
        if (LD->getPointerOperand() != PtrPHI) { ok = false; break; }
        Type *T = LD->getType();
        if (!T->isIntegerTy()) { ok = false; break; }
        if (!ElemTy) ElemTy = T;
        else if (ElemTy != T) { ok = false; break; }
        Accesses.push_back(LD);
        continue;
      }
      if (auto *ST = dyn_cast<StoreInst>(U)) {
        if (ST->getPointerOperand() != PtrPHI) { ok = false; break; }
        Type *T = ST->getValueOperand()->getType();
        if (!T->isIntegerTy()) { ok = false; break; }
        if (!ElemTy) ElemTy = T;
        else if (ElemTy != T) { ok = false; break; }
        Accesses.push_back(ST);
        continue;
      }
      // Any other use → bail.
      ok = false;
      break;
    }
    if (!ok || !ElemTy || Accesses.empty()) continue;
    Module &M = *Header->getModule();
    const DataLayout &DL = M.getDataLayout();
    uint64_t ElemBytes = DL.getTypeAllocSize(ElemTy);
    if ((int64_t)ElemBytes != StrideBytes) continue;

    // We have a clean global-array pointer-walk.  Build a forward
    // counter PHI starting at 0, incrementing by 1.  Replace the
    // pointer-PHI's loads/stores with `GEP ElemTy, @global, %i`.
    LLVMContext &Ctx = Header->getContext();
    IntegerType *I16 = Type::getInt16Ty(Ctx);
    IRBuilder<> B(&Header->front());
    PHINode *Counter = B.CreatePHI(I16, 2, "unlsr.i");
    BasicBlock *Preheader = L->getLoopPreheader();
    if (!Preheader) continue;
    Counter->addIncoming(ConstantInt::get(I16, 0), Preheader);
    // Build the counter increment in the latch, just before the
    // terminator.
    IRBuilder<> BL(Latch->getTerminator());
    Value *CounterNext = BL.CreateAdd(Counter, ConstantInt::get(I16, 1),
                                       "unlsr.i.next");
    Counter->addIncoming(CounterNext, Latch);

    // Replace each load/store with a GEP-then-load/store.
    for (Instruction *Acc : Accesses) {
      IRBuilder<> BA(Acc);
      Value *Idx = Counter;
      // GEP type is i16 elem, idx is i16.
      Value *Addr = BA.CreateGEP(ElemTy, InitGV, Idx, "unlsr.addr");
      if (auto *LD = dyn_cast<LoadInst>(Acc)) {
        LD->setOperand(LoadInst::getPointerOperandIndex(), Addr);
      } else if (auto *ST = dyn_cast<StoreInst>(Acc)) {
        ST->setOperand(StoreInst::getPointerOperandIndex(), Addr);
      }
    }
    // Erase the now-unused step GEP and pointer PHI.
    if (StepGEP->use_empty()) StepGEP->eraseFromParent();
    if (PtrPHI->use_empty()) PtrPHI->eraseFromParent();

    // Try to eliminate the count-down LSR counter (`%lsr.iv = phi i16
    // [%n, preheader], [%lsr.iv.next, latch]; %lsr.iv.next = add %lsr.iv,
    // -1; %exitcond = icmp eq %lsr.iv.next, 0`).  Replace the exit
    // comparison with `icmp eq %unlsr.i.next, %n`, then erase the
    // count-down chain.  Saves one i16 PHI + dec + cmp per loop.
    for (PHINode &LSR : Header->phis()) {
      if (&LSR == Counter) continue;
      if (!LSR.getType()->isIntegerTy(16)) continue;
      if (LSR.getNumIncomingValues() != 2) continue;
      // Identify preheader vs latch incomings.
      Value *LsrInit = nullptr;
      Value *LsrStep = nullptr;
      for (unsigned i = 0; i < LSR.getNumIncomingValues(); ++i) {
        BasicBlock *Pred = LSR.getIncomingBlock(i);
        if (L->contains(Pred)) LsrStep = LSR.getIncomingValue(i);
        else                    LsrInit = LSR.getIncomingValue(i);
      }
      if (!LsrInit || !LsrStep) continue;
      // %lsr.iv.next must be `add %lsr.iv, -1`.
      auto *LsrStepBO = dyn_cast<BinaryOperator>(LsrStep);
      if (!LsrStepBO || LsrStepBO->getOpcode() != Instruction::Add) continue;
      Value *OtherOp = nullptr;
      if (LsrStepBO->getOperand(0) == &LSR) OtherOp = LsrStepBO->getOperand(1);
      else if (LsrStepBO->getOperand(1) == &LSR) OtherOp = LsrStepBO->getOperand(0);
      if (!OtherOp) continue;
      auto *DecCI = dyn_cast<ConstantInt>(OtherOp);
      if (!DecCI || !DecCI->isMinusOne()) continue;
      // %lsr.iv.next typically has 2 uses: the icmp exit comparison and
      // the PHI back-edge (LSR.users includes LsrStepBO as the latch
      // incoming).  Allow that; find the icmp among the uses.
      ICmpInst *ExitCmp = nullptr;
      bool extraUse = false;
      for (User *U : LsrStepBO->users()) {
        if (U == &LSR) continue;
        if (auto *IC = dyn_cast<ICmpInst>(U)) {
          if (ExitCmp) { extraUse = true; break; }
          ExitCmp = IC;
          continue;
        }
        extraUse = true;
        break;
      }
      if (extraUse || !ExitCmp) continue;
      if (ExitCmp->getPredicate() != ICmpInst::ICMP_EQ &&
          ExitCmp->getPredicate() != ICmpInst::ICMP_NE) continue;
      Value *CmpOther = nullptr;
      if (ExitCmp->getOperand(0) == LsrStepBO)
        CmpOther = ExitCmp->getOperand(1);
      else if (ExitCmp->getOperand(1) == LsrStepBO)
        CmpOther = ExitCmp->getOperand(0);
      if (!CmpOther) continue;
      auto *CmpZero = dyn_cast<ConstantInt>(CmpOther);
      if (!CmpZero || !CmpZero->isZero()) continue;
      bool lsrOnlyIncrUse = true;
      for (User *U : LSR.users()) {
        if (U == LsrStepBO) continue;
        lsrOnlyIncrUse = false; break;
      }
      if (!lsrOnlyIncrUse) continue;
      // CounterNext was inserted just before the latch terminator, but
      // ExitCmp may live higher in the block (originally placed there
      // by LSR).  Move CounterNext to just before ExitCmp so it
      // dominates the cmp's use.
      if (auto *CounterNextI = dyn_cast<Instruction>(CounterNext))
        CounterNextI->moveBefore(ExitCmp->getIterator());
      // Replace exit comparison: icmp eq %unlsr.i.next, %lsrInit.
      ExitCmp->setOperand(0, CounterNext);
      ExitCmp->setOperand(1, LsrInit);
      // RAUW each value with undef before erasing — they have a mutual
      // reference (PHI → BO via latch-incoming, BO → PHI via add operand).
      Type *I16T = LSR.getType();
      LSR.replaceAllUsesWith(UndefValue::get(I16T));
      LsrStepBO->replaceAllUsesWith(UndefValue::get(I16T));
      LsrStepBO->eraseFromParent();
      LSR.eraseFromParent();
      break;
    }

    Changed = true;
  }
  return Changed;
}
