//===-- W65816SjLjFinalize.cpp - Finish SJLJ EH 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
//
//===---------------------------------------------------------------------===//
//
// SjLjEHPrepare leaves IR with a function-context alloca, a few marker
// intrinsics (eh.sjlj.lsda / setup.dispatch / functioncontext / callsite),
// and `invoke` instructions whose unwind dest is a landing pad block that
// reads exception+selector from the function context.  The expectation
// is that the BACKEND finishes lowering by inserting an actual setjmp
// at function entry and a switch-on-call-site dispatch block.
//
// On targets like ARM that's done with custom inserter pseudos
// (Int_eh_sjlj_setjmp + EmitSjLjDispatchBlock).  We don't have any of
// that machinery, so this pass does it at IR level instead.  Concretely:
//
//   1. Find the function-context alloca and the _Unwind_SjLj_Register
//      call (SjLjEHPrepare placed both at function entry).
//   2. After the Register call, insert:
//        %r = call i16 @setjmp(ptr %jbuf)  ; jbuf is fn_ctx[5]
//        %is_unwind = icmp ne i16 %r, 0
//        br i1 %is_unwind, label %eh.dispatch, label %normal_entry
//   3. Build %eh.dispatch:
//        %cs = load i32, ptr %call_site_field
//        switch i32 %cs, label %resume_unreachable
//                       [ i32 1, label %lpad1
//                         i32 2, label %lpad2
//                         ... ]
//      The lpadN blocks already exist (originally invoke's unwind dest).
//   4. Convert each `invoke F(args) to %normal unwind to %lpad`
//      to a regular `call F(args); br %normal`.  The eh.sjlj.callsite(N)
//      intrinsic just before the invoke recorded N — we extract the
//      mapping from the explicit `store i32 N, ptr %call_site_field`
//      that SjLjEHPrepare also emitted.
//   5. Erase eh.sjlj.lsda / setup.dispatch / functioncontext / callsite
//      intrinsic calls.  lsda's result is replaced with null — our
//      personality routine doesn't consume an LSDA pointer; it knows
//      the catch types from the function context's data array which
//      _Unwind_SjLj_RaiseException populated.
//
// The runtime side (runtime/src/libcxxabiSjlj.c) provides
// _Unwind_SjLj_Register/Unregister, _Unwind_SjLj_RaiseException,
// _Unwind_SjLj_Resume, and __gxx_personality_sj0 — plus the libcxxabi
// surface (__cxa_throw, __cxa_begin_catch, etc.).
//
//===---------------------------------------------------------------------===//

#include "W65816.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/IntrinsicInst.h"
#include "llvm/IR/Intrinsics.h"
#include "llvm/IR/Module.h"
#include "llvm/Pass.h"

using namespace llvm;

#define DEBUG_TYPE "w65816-sjlj-finalize"

namespace {

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

  StringRef getPassName() const override {
    return "W65816 SJLJ EH finalize";
  }

  bool runOnFunction(Function &F) override;
};

} // namespace

char W65816SjLjFinalize::ID = 0;

INITIALIZE_PASS(W65816SjLjFinalize, DEBUG_TYPE,
                "W65816 SJLJ EH finalize", false, false)

FunctionPass *llvm::createW65816SjLjFinalize() {
  return new W65816SjLjFinalize();
}

// Match the personality recorded by clang for SJLJ EH.
static bool hasSjLjPersonality(const Function &F) {
  if (!F.hasPersonalityFn())
    return false;
  const Constant *P = F.getPersonalityFn();
  if (auto *Fn = dyn_cast<Function>(P->stripPointerCasts()))
    return Fn->getName() == "__gxx_personality_sj0";
  return false;
}

// Find the alloca SjLjEHPrepare used for the function context.  It's
// the only alloca whose first use is a GEP storing the personality fn
// (field 3) — but a more reliable marker is the eh.sjlj.functioncontext
// intrinsic call, which takes the alloca as its sole argument.
static AllocaInst *findFnCtxAlloca(Function &F) {
  for (Instruction &I : instructions(F)) {
    if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
      if (II->getIntrinsicID() == Intrinsic::eh_sjlj_functioncontext) {
        return cast<AllocaInst>(II->getArgOperand(0)->stripPointerCasts());
      }
    }
  }
  return nullptr;
}

bool W65816SjLjFinalize::runOnFunction(Function &F) {
  if (!hasSjLjPersonality(F))
    return false;

  AllocaInst *FnCtx = findFnCtxAlloca(F);
  if (!FnCtx)
    return false;

  Module &M = *F.getParent();
  LLVMContext &Ctx = F.getContext();
  Type *I32Ty = Type::getInt32Ty(Ctx);
  Type *I16Ty = Type::getInt16Ty(Ctx);
  Type *PtrTy = PointerType::getUnqual(Ctx);
  Type *FnCtxTy = FnCtx->getAllocatedType();

  // Walk invokes; build the call-site → landing-pad map.  The
  // call-site index for each invoke is the i32 stored to fn_ctx[1]
  // immediately before the invoke (SjLjEHPrepare placed it).
  SmallVector<InvokeInst *, 4> Invokes;
  DenseMap<int, BasicBlock *> CSToLPad;
  for (BasicBlock &BB : F) {
    if (auto *II = dyn_cast<InvokeInst>(BB.getTerminator())) {
      Invokes.push_back(II);
      // Walk backward from the invoke for the most recent
      // `store i32 <const>, ptr <call_site_gep>`.
      int CS = -1;
      for (auto It = std::next(BB.rbegin()); It != BB.rend(); ++It) {
        if (auto *SI = dyn_cast<StoreInst>(&*It)) {
          if (auto *C = dyn_cast<ConstantInt>(SI->getValueOperand())) {
            if (auto *GEP = dyn_cast<GetElementPtrInst>(SI->getPointerOperand())) {
              if (GEP->getPointerOperand() == FnCtx) {
                CS = (int)C->getSExtValue();
                break;
              }
            }
          }
        }
      }
      if (CS > 0)
        CSToLPad[CS] = II->getUnwindDest();
    }
  }

  if (Invokes.empty())
    return false;

  // Find the call to _Unwind_SjLj_Register — our setjmp insertion point
  // is right after it (so the function context is fully populated).
  CallInst *RegisterCall = nullptr;
  for (Instruction &I : F.getEntryBlock()) {
    if (auto *CI = dyn_cast<CallInst>(&I)) {
      if (CI->getCalledFunction() &&
          CI->getCalledFunction()->getName() == "_Unwind_SjLj_Register") {
        RegisterCall = CI;
        break;
      }
    }
  }
  if (!RegisterCall)
    return false;

  // Materialize: %r = call setjmp(ptr %jbuf)
  // jbuf is fn_ctx field 5.
  BasicBlock *EntryBB = RegisterCall->getParent();
  IRBuilder<> Builder(RegisterCall->getNextNode());
  Value *Jbuf = Builder.CreateStructGEP(FnCtxTy, FnCtx, 5, "jbuf");
  // Our setjmp signature: int setjmp(void *jb).  We treat its return
  // as i16 here to match the W65816 ABI; the runtime returns 0 on the
  // initial call and a nonzero value when longjmp comes back.
  FunctionCallee SetjmpFn =
      M.getOrInsertFunction("setjmp", FunctionType::get(I16Ty, {PtrTy}, false));
  // Mark setjmp as returns_twice so LLVM doesn't optimize across it.
  if (auto *SF = dyn_cast<Function>(SetjmpFn.getCallee())) {
    SF->addFnAttr(Attribute::ReturnsTwice);
  }
  CallInst *SetjmpCall = Builder.CreateCall(SetjmpFn, {Jbuf}, "sj.r");
  SetjmpCall->setCanReturnTwice();
  Value *IsUnwind = Builder.CreateICmpNE(
      SetjmpCall, ConstantInt::get(I16Ty, 0), "sj.is_unwind");

  // Split entry block: instructions after our br go into a new block;
  // we'll branch there on the first-time setjmp return.
  BasicBlock *EHContinueBB =
      EntryBB->splitBasicBlock(Builder.GetInsertPoint(), "sj.first_entry");
  // splitBasicBlock created an unconditional br at the split point;
  // replace it with our conditional branch to dispatch vs continue.
  EntryBB->getTerminator()->eraseFromParent();
  Builder.SetInsertPoint(EntryBB);

  // Build dispatch block: switch on call_site.
  BasicBlock *DispatchBB = BasicBlock::Create(Ctx, "sj.dispatch", &F);
  IRBuilder<> DBuilder(DispatchBB);
  Value *CallSiteGEP =
      DBuilder.CreateStructGEP(FnCtxTy, FnCtx, 1, "cs.gep");
  LoadInst *CallSite =
      DBuilder.CreateLoad(I32Ty, CallSiteGEP, /*isVolatile=*/true, "cs");
  // Default case: if we ever land here with an unknown call_site, just
  // unreachable — the runtime should never longjmp with an out-of-range
  // index.  But a defensive fallback is to spin (terminate).
  BasicBlock *DispatchUnreachable =
      BasicBlock::Create(Ctx, "sj.dispatch.unreachable", &F);
  new UnreachableInst(Ctx, DispatchUnreachable);
  SwitchInst *SI =
      DBuilder.CreateSwitch(CallSite, DispatchUnreachable, CSToLPad.size());
  for (auto &KV : CSToLPad) {
    SI->addCase(cast<ConstantInt>(ConstantInt::get(I32Ty, KV.first)),
                KV.second);
  }

  // Final entry-block terminator: if (is_unwind) goto dispatch else continue.
  Builder.CreateCondBr(IsUnwind, DispatchBB, EHContinueBB);

  // The landing-pad blocks each start with a `landingpad { ptr, i32 }`
  // instruction.  We need its catch-clause typeinfo arguments to build
  // the per-function catch table further down, so capture them BEFORE
  // erasing — the catchtab loop below uses the saved data instead of
  // re-reading from the IR.
  //
  // Capture: per call_site, list of (typeinfo Constant*) for each
  // catch clause (skipping null = catch-all).  De-dup landingpads
  // (multiple call_sites can share a landing pad).
  struct LPadInfo {
    SmallVector<Constant *, 2> CatchTypes;
  };
  DenseMap<BasicBlock *, LPadInfo> LPadCatches;
  SmallVector<LandingPadInst *, 4> LPads;
  for (auto &KV : CSToLPad) {
    BasicBlock *LPadBB = KV.second;
    if (LPadCatches.count(LPadBB))
      continue;
    LandingPadInst *LP = LPadBB->getLandingPadInst();
    if (!LP)
      continue;
    LPadInfo Info;
    for (unsigned i = 0; i < LP->getNumClauses(); i++) {
      if (LP->isCatch(i)) {
        Constant *TIClause = cast<Constant>(LP->getClause(i));
        if (TIClause->isNullValue())
          continue;
        Info.CatchTypes.push_back(TIClause);
      }
    }
    LPadCatches[LPadBB] = std::move(Info);
    LPads.push_back(LP);
  }
  // After we convert invokes to plain calls, landingpad blocks are no
  // longer reached via an unwind edge — the verifier requires landing-
  // pads to be reached only from invoke unwind dests.  Erase them now
  // (they're no-ops post-finalize; the real exception ptr / selector
  // come from explicit fn_ctx.data loads SjLjEHPrepare emitted right
  // after).  Replace landingpad uses with poison since downstream
  // code reads via GEPs, not via the inst's result.
  for (LandingPadInst *LP : LPads) {
    LP->replaceAllUsesWith(PoisonValue::get(LP->getType()));
    LP->eraseFromParent();
  }
  // The function's "personality" attribute references @__gxx_personality_sj0,
  // which would normally require landingpads.  Drop it since we have none.
  F.setPersonalityFn(nullptr);

  // Convert each invoke to a regular call + br to its normal dest.
  for (InvokeInst *II : Invokes) {
    SmallVector<Value *, 8> Args(II->args());
    SmallVector<OperandBundleDef, 1> Bundles;
    II->getOperandBundlesAsDefs(Bundles);
    CallInst *CI = CallInst::Create(II->getFunctionType(),
                                    II->getCalledOperand(), Args, Bundles, "",
                                    II->getIterator());
    CI->setCallingConv(II->getCallingConv());
    CI->setAttributes(II->getAttributes());
    CI->setDebugLoc(II->getDebugLoc());
    if (!II->getType()->isVoidTy())
      II->replaceAllUsesWith(CI);
    BasicBlock *NormalDest = II->getNormalDest();
    BasicBlock *UnwindDest = II->getUnwindDest();
    BranchInst::Create(NormalDest, II->getIterator());
    // Drop the unwind-dest PHI predecessor entries for this invoke's BB.
    UnwindDest->removePredecessor(II->getParent());
    II->eraseFromParent();
  }

  // Build per-function catch table.  Format (flat array of i16 pairs):
  //   [cs1, ti_addr1, cs1, ti_addr2, ..., cs2, ti_addr1, ..., 0, 0]
  // Each (call_site, typeinfo_address) row encodes "if a throw is in
  // flight while call_site is active, try to catch with this typeinfo".
  // Terminated by a (0, 0) sentinel.  Catch table address is stored in
  // fn_ctx[4] (the lsda field) by replacing eh.sjlj.lsda's result.
  //
  // To make the landing pad's selector compare work without a real
  // selector value: we set selector = (i32)(uintptr_t)&typeinfo at
  // longjmp time, and rewrite the landing pad's
  // eh.typeid.for(@TI) calls to also yield (i32)(uintptr_t)&TI.  The
  // icmp eq then succeeds for the matched catch.
  SmallVector<Constant *, 16> TableRows;
  // Walk each invoke's call_site → unwind_dest mapping; emit a row per
  // (call_site, catch typeinfo) pair using the LPadCatches data we
  // captured BEFORE erasing the landingpad insts above.
  for (auto &KV : CSToLPad) {
    int CS = KV.first;
    BasicBlock *LPadBB = KV.second;
    auto It = LPadCatches.find(LPadBB);
    if (It == LPadCatches.end())
      continue;
    for (Constant *TIClause : It->second.CatchTypes) {
      TableRows.push_back(ConstantInt::get(I16Ty, CS));
      TableRows.push_back(ConstantExpr::getPtrToInt(TIClause, I16Ty));
    }
  }
  // Append (0, 0) sentinel.
  TableRows.push_back(ConstantInt::get(I16Ty, 0));
  TableRows.push_back(ConstantInt::get(I16Ty, 0));

  ArrayType *TableArrTy = ArrayType::get(I16Ty, TableRows.size());
  Constant *TableInit = ConstantArray::get(TableArrTy, TableRows);
  std::string TableName = "_W65SJLJ_CATCHTAB_" + F.getName().str();
  GlobalVariable *Table = new GlobalVariable(
      M, TableArrTy, /*isConstant=*/true, GlobalValue::InternalLinkage,
      TableInit, TableName);

  // Replace eh.sjlj.lsda → catch-table address; rewrite eh.typeid.for
  // to (i32) ptrtoint of its typeinfo arg; erase setup_dispatch /
  // functioncontext / callsite intrinsic markers.
  SmallVector<Instruction *, 8> ToErase;
  for (Instruction &I : instructions(F)) {
    if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
      switch (II->getIntrinsicID()) {
      case Intrinsic::eh_sjlj_lsda: {
        Constant *TableAddr = ConstantExpr::getBitCast(Table, PtrTy);
        II->replaceAllUsesWith(TableAddr);
        ToErase.push_back(II);
        break;
      }
      case Intrinsic::eh_typeid_for: {
        // Replace with: zext (ptrtoint TI to i16) to i32.
        IRBuilder<> TBuilder(II);
        Value *TI = II->getArgOperand(0);
        Value *AsI16 = TBuilder.CreatePtrToInt(TI, I16Ty);
        Value *AsI32 = TBuilder.CreateZExt(AsI16, I32Ty);
        II->replaceAllUsesWith(AsI32);
        ToErase.push_back(II);
        break;
      }
      case Intrinsic::eh_sjlj_setup_dispatch:
      case Intrinsic::eh_sjlj_functioncontext:
      case Intrinsic::eh_sjlj_callsite:
        ToErase.push_back(II);
        break;
      default:
        break;
      }
    }
  }
  for (Instruction *I : ToErase)
    I->eraseFromParent();

  return true;
}