//===- LowerLLVMExtPass.cpp - Lower LLVM Ext operations  ------------------ //
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This file implements a pass to lower custom ops generated by the LLVM Ext
// dialect.
//
//===----------------------------------------------------------------------===//

#include "Dialect/LLVMExt/LLVMExt.h"
#include "Passes/Passes.h"

#include "mlir/Dialect/LLVMIR/LLVMDialect.h"

#include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h"

namespace mlir {
namespace enzyme {
using namespace mlir::enzyme;
#define GEN_PASS_DEF_LOWERLLVMEXTPASS
#include "Passes/Passes.h.inc"
} // namespace enzyme
} // namespace mlir

namespace {
using namespace mlir;
using namespace enzyme;

LogicalResult tryLoweringToAlloca(llvm_ext::AllocOp alloc,
                                  uint64_t staticThreshold) {
  llvm::APInt size;
  if (!matchPattern(alloc.getSize(), m_ConstantInt(&size)))
    return failure();

  if (size.getZExtValue() > staticThreshold)
    return failure();

  Operation *free = nullptr;
  for (auto user : alloc.getResult().getUsers()) {
    auto oFree = dyn_cast<llvm_ext::FreeOp>(user);
    if (!oFree)
      continue;

    if (free)
      return failure(); // multiple frees

    if (oFree->getBlock() != alloc->getBlock()) // free not in same block
      return failure();

    free = user;

    if (!alloc->isBeforeInBlock(free))
      return failure();
  }

  OpBuilder builder(alloc);
  auto alloca = LLVM::AllocaOp::create(builder, alloc.getLoc(),
                                       alloc.getResult().getType(),
                                       builder.getI8Type(), alloc.getSize());
  alloc.getResult().replaceAllUsesWith(alloca.getResult());

  LLVM::LifetimeStartOp::create(builder, alloc.getLoc(), alloca.getResult());

  builder.setInsertionPoint(free);
  LLVM::LifetimeEndOp::create(builder, alloc.getLoc(), alloca.getResult());

  free->erase();
  alloc->erase();

  return success();
}

void lowerAlloc(llvm_ext::AllocOp alloc, uint64_t staticThreshold) {
  if (tryLoweringToAlloca(alloc, staticThreshold).succeeded())
    return;

  SymbolTable symtable(SymbolTable::getNearestSymbolTable(alloc));

  auto mallocFn = symtable.lookup<LLVM::LLVMFuncOp>("malloc");
  if (!mallocFn) {
    Block *b = &symtable.getOp()->getRegion(0).front();
    OpBuilder builder(b, b->begin());

    auto fnType = mlir::LLVM::LLVMFunctionType::get(
        LLVM::LLVMPointerType::get(alloc.getContext()), builder.getI64Type(),
        /*isVarArg=*/false);

    mallocFn =
        LLVM::LLVMFuncOp::create(builder, alloc.getLoc(), "malloc", fnType);
  }

  OpBuilder builder(alloc);
  auto mallocCall = LLVM::CallOp::create(builder, alloc.getLoc(), mallocFn,
                                         alloc->getOperands());
  alloc.getResult().replaceAllUsesWith(mallocCall.getResult());
  alloc.erase();
}

void lowerFree(llvm_ext::FreeOp free) {
  SymbolTable symtable(SymbolTable::getNearestSymbolTable(free));

  auto freeFn = symtable.lookup<LLVM::LLVMFuncOp>("free");
  if (!freeFn) {
    Block *b = &symtable.getOp()->getRegion(0).front();
    OpBuilder builder(b, b->begin());

    auto fnType = mlir::LLVM::LLVMFunctionType::get(
        LLVM::LLVMVoidType::get(free.getContext()),
        LLVM::LLVMPointerType::get(free.getContext()),
        /*isVarArg=*/false);

    freeFn = LLVM::LLVMFuncOp::create(builder, free.getLoc(), "free", fnType);
  }

  OpBuilder builder(free);
  LLVM::CallOp::create(builder, free.getLoc(), freeFn, free->getOperands());

  free.erase();
}

struct LowerLLVMExtPass
    : public enzyme::impl::LowerLLVMExtPassBase<LowerLLVMExtPass> {
  using LowerLLVMExtPassBase::LowerLLVMExtPassBase;

  void runOnOperation() override {
    Operation *op = getOperation();

    op->walk([](llvm_ext::PtrSizeHintOp psh) { psh.erase(); });

    SmallVector<llvm_ext::AllocOp> allocs;
    op->walk([&](llvm_ext::AllocOp alloc) { allocs.push_back(alloc); });

    for (auto alloc : allocs)
      lowerAlloc(alloc, lowerToAllocaThreshold);

    SmallVector<llvm_ext::FreeOp> frees;
    op->walk([&](llvm_ext::FreeOp free) { frees.push_back(free); });

    for (auto free : frees)
      lowerFree(free);
  }
};

} // end anonymous namespace
