//===-- W65816NarrowI32Mul.cpp - Narrow i32 multiplies -------------------===//
//
// 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
//
//===---------------------------------------------------------------------===//
//
// IR-level peephole.  Detects `mul i32 X, Y` where both X and Y have
// their top 16 bits provably zero (via LLVM's IR-level computeKnownBits)
// and rewrites to a call to `__umulhisi3` — a 16x16 -> 32 unsigned
// multiply (~30% faster than __mulsi3 for `(u32)i * i` patterns).
//
// Why an IR pass instead of a Custom SDAG lowering: LLVM's IndVarSimplify
// loop pass widens narrow induction variables (e.g. an i16 loop counter
// later zext'd to i32) into i32 PHIs.  By SDAG-build time the zext is
// gone — the MUL's operand is just `CopyFromReg %2:i32`, an opaque value.
// SDAG's computeKnownBits can't trace back across BB boundaries through
// CopyFromReg.  IR-level computeKnownBits, by contrast, walks the use-def
// graph (including PHIs) and can prove the high bits zero.
//
// Runs in addISelPrepare (right before SDAG-ISel) so it sees the
// final-shape IR.  The libcall declaration is auto-added if missing.
//
//===---------------------------------------------------------------------===//

#include "W65816.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/Analysis/ScalarEvolution.h"
#include "llvm/InitializePasses.h"
#include "llvm/Analysis/ValueTracking.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/KnownBits.h"

using namespace llvm;

#define DEBUG_TYPE "w65816-narrow-i32-mul"


namespace {


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

  StringRef getPassName() const override {
    return "W65816 narrow i32 multiplies to __umulhisi3";
  }

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

  bool runOnFunction(Function &F) override;
};


} // namespace


char W65816NarrowI32Mul::ID = 0;

INITIALIZE_PASS_BEGIN(W65816NarrowI32Mul, DEBUG_TYPE,
                      "W65816 narrow i32 multiplies", false, false)
INITIALIZE_PASS_DEPENDENCY(ScalarEvolutionWrapperPass)
INITIALIZE_PASS_END(W65816NarrowI32Mul, DEBUG_TYPE,
                    "W65816 narrow i32 multiplies", false, false)


// Get-or-declare `__umulhisi3(i16, i16) -> i32` in the module.
static FunctionCallee getUmulhisi3(Module &M) {
  LLVMContext &Ctx = M.getContext();
  Type *I16 = Type::getInt16Ty(Ctx);
  Type *I32 = Type::getInt32Ty(Ctx);
  FunctionType *FT = FunctionType::get(I32, {I16, I16}, false);
  return M.getOrInsertFunction("__umulhisi3", FT);
}


// True iff the top 16 bits of V are known zero.  Tries IR-level
// computeKnownBits first; if that doesn't prove enough, falls back
// to ScalarEvolution's unsigned-range analysis (which handles
// loop-bounded induction variables that KnownBits can't).
static bool top16Zero(Value *V, const DataLayout &DL, ScalarEvolution &SE) {
  KnownBits K = computeKnownBits(V, DL);
  if (K.countMinLeadingZeros() >= 16) {
    return true;
  }
  if (!SE.isSCEVable(V->getType())) {
    return false;
  }
  const SCEV *S = SE.getSCEV(V);
  ConstantRange R = SE.getUnsignedRange(S);
  return R.getActiveBits() <= 16;
}


bool W65816NarrowI32Mul::runOnFunction(Function &F) {
  Module *M = F.getParent();
  const DataLayout &DL = M->getDataLayout();
  Type *I16 = Type::getInt16Ty(F.getContext());
  ScalarEvolution &SE = getAnalysis<ScalarEvolutionWrapperPass>().getSE();

  SmallVector<BinaryOperator *, 8> Worklist;
  for (Instruction &I : instructions(F)) {
    auto *BO = dyn_cast<BinaryOperator>(&I);
    if (!BO || BO->getOpcode() != Instruction::Mul) {
      continue;
    }
    if (!BO->getType()->isIntegerTy(32)) {
      continue;
    }
    if (!top16Zero(BO->getOperand(0), DL, SE)) {
      continue;
    }
    if (!top16Zero(BO->getOperand(1), DL, SE)) {
      continue;
    }
    Worklist.push_back(BO);
  }

  if (Worklist.empty()) {
    return false;
  }

  // When the i32 operand is `zext i16 X to i32`, use X directly instead
  // of emitting `trunc i32 (zext i16 X) to i16` — that trunc-of-zext is
  // semantically the identity but keeps the zext (= a fresh i32 SSA
  // value) live, which materializes a Wide32 vreg pair at ISel and
  // forces a 4-byte spill slot (the canonical sumSquares `conv` pattern
  // burned slots 0xd / 0x5 this way).  Skipping the trunc lets the
  // post-replaceAll DCE drop the zext entirely, freeing the slot.
  auto narrowOperand = [&](Value *V, IRBuilder<> &B) -> Value * {
    if (auto *ZE = dyn_cast<ZExtInst>(V)) {
      if (ZE->getSrcTy() == I16) return ZE->getOperand(0);
    }
    if (auto *AE = dyn_cast<SExtInst>(V)) {
      // Sext from i16 also has the right low 16 bits.
      if (AE->getSrcTy() == I16) return AE->getOperand(0);
    }
    return B.CreateTrunc(V, I16);
  };

  FunctionCallee Callee = getUmulhisi3(*M);
  SmallPtrSet<Instruction *, 8> MaybeDeadSet;
  for (BinaryOperator *BO : Worklist) {
    IRBuilder<> B(BO);
    Value *AOp = BO->getOperand(0);
    Value *BOp = BO->getOperand(1);
    Value *A = narrowOperand(AOp, B);
    Value *Bv = narrowOperand(BOp, B);
    Value *Call = B.CreateCall(Callee, {A, Bv});
    BO->replaceAllUsesWith(Call);
    BO->eraseFromParent();
    // If the original operands were zext/sext nodes, they may now be
    // dead.  Add to a SET — `mul %x, %x` would push the same node
    // twice, and the cleanup loop would then double-free.
    if (auto *I = dyn_cast<Instruction>(AOp)) MaybeDeadSet.insert(I);
    if (auto *I = dyn_cast<Instruction>(BOp)) MaybeDeadSet.insert(I);
  }
  // Cleanup: any extension that's now use-less can be deleted.
  for (Instruction *I : MaybeDeadSet) {
    if (I->use_empty() && (isa<ZExtInst>(I) || isa<SExtInst>(I) ||
                            isa<TruncInst>(I))) {
      I->eraseFromParent();
    }
  }

  // Phase 2: narrow LSR-introduced i32 PHIs whose only uses (after
  // the mul-rewrite above) are trunc-to-i16 + a single self-feeding
  // `add %P, const` increment.  Without this, even though the mul
  // operates on i16, the i32 PHI still requires 4 bytes of frame +
  // an i32 increment chain (post-PEI).  LSR widened these from i16
  // to i32 to support a sub-expression that we've now narrowed —
  // the i32 representation has become dead weight.
  //
  // Guard with SCEV: `getUnsignedRange(%P).getActiveBits() <= 16`
  // proves the PHI never escapes u16, so the i16 add gives the same
  // low-16 bits as the original i32 add at every observable point
  // (the back-edge value can wrap on the exit iteration but is
  // never observed — exit takes the trip-end branch first).
  SmallVector<PHINode *, 4> PhiWorklist;
  for (BasicBlock &BB : F) {
    for (PHINode &PN : BB.phis()) {
      if (PN.getType()->isIntegerTy(32)) PhiWorklist.push_back(&PN);
    }
  }
  for (PHINode *PN : PhiWorklist) {
    // Classify every use.
    SmallVector<TruncInst *, 4> Truncs;
    BinaryOperator *Incr = nullptr;
    bool ok = true;
    for (User *U : PN->users()) {
      if (auto *TI = dyn_cast<TruncInst>(U)) {
        if (!TI->getDestTy()->isIntegerTy(16)) { ok = false; break; }
        Truncs.push_back(TI);
        continue;
      }
      auto *BO = dyn_cast<BinaryOperator>(U);
      if (!BO || BO->getOpcode() != Instruction::Add) { ok = false; break; }
      if (!isa<ConstantInt>(BO->getOperand(1))) { ok = false; break; }
      // BO must feed back to this PHI via at least one incoming edge.
      bool feedsBack = false;
      for (Value *Inc : PN->incoming_values()) {
        if (Inc == BO) { feedsBack = true; break; }
      }
      if (!feedsBack) { ok = false; break; }
      if (Incr) { ok = false; break; }
      Incr = BO;
    }
    if (!ok || !Incr || Truncs.empty()) continue;

    // Increment const must fit i16.
    auto *IncrCI = cast<ConstantInt>(Incr->getOperand(1));
    if (IncrCI->getValue().getActiveBits() > 16) continue;
    // Non-back-edge incomings must be i16-representable constants.
    for (Value *Inc : PN->incoming_values()) {
      if (Inc == Incr) continue;
      auto *CIv = dyn_cast<ConstantInt>(Inc);
      if (!CIv) { ok = false; break; }
      if (CIv->getValue().getActiveBits() > 16) { ok = false; break; }
    }
    if (!ok) continue;
    // SCEV bound check.
    if (!SE.isSCEVable(PN->getType())) continue;
    ConstantRange R = SE.getUnsignedRange(SE.getSCEV(PN));
    if (R.getActiveBits() > 16) continue;

    // Narrow.  Build %narrow_phi in same BB, then %narrow_incr right
    // before Incr; patch incoming values to match.
    IRBuilder<> B(PN);
    PHINode *NewPN = B.CreatePHI(I16, PN->getNumIncomingValues(),
                                  PN->getName() + ".narrow");
    // Add placeholders for the back-edge incomings; we'll patch them
    // after building NewIncr.
    for (unsigned i = 0; i < PN->getNumIncomingValues(); ++i) {
      Value *Inc = PN->getIncomingValue(i);
      BasicBlock *Pred = PN->getIncomingBlock(i);
      if (Inc == Incr) {
        NewPN->addIncoming(UndefValue::get(I16), Pred);
      } else {
        auto *CIv = cast<ConstantInt>(Inc);
        NewPN->addIncoming(
            ConstantInt::get(I16, CIv->getZExtValue() & 0xFFFF),
            Pred);
      }
    }
    IRBuilder<> B2(Incr);
    Value *NewIncr = B2.CreateAdd(
        NewPN,
        ConstantInt::get(I16, IncrCI->getZExtValue() & 0xFFFF),
        Incr->getName() + ".narrow");
    if (auto *NewIncrBO = dyn_cast<BinaryOperator>(NewIncr)) {
      NewIncrBO->setHasNoUnsignedWrap(Incr->hasNoUnsignedWrap());
      NewIncrBO->setHasNoSignedWrap(Incr->hasNoSignedWrap());
    }
    for (unsigned i = 0; i < NewPN->getNumIncomingValues(); ++i) {
      if (isa<UndefValue>(NewPN->getIncomingValue(i))) {
        NewPN->setIncomingValue(i, NewIncr);
      }
    }
    // Replace trunc uses with the new narrow PHI, then break the
    // PHI/Incr use-cycle before erasing.
    for (TruncInst *TI : Truncs) {
      TI->replaceAllUsesWith(NewPN);
      TI->eraseFromParent();
    }
    // Incr is `add %PN, const`; PN's back-edge incoming references Incr.
    // Replace Incr's uses with undef so PN's back-edge becomes a dead
    // reference, then erase Incr, then PN.
    Incr->replaceAllUsesWith(UndefValue::get(Incr->getType()));
    Incr->eraseFromParent();
    PN->eraseFromParent();
  }
  return true;
}


FunctionPass *llvm::createW65816NarrowI32Mul() {
  return new W65816NarrowI32Mul();
}
