//===- MemoryUtils.cpp ----------------------------------------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//

#include "flang/Optimizer/Transforms/MemoryUtils.h"
#include "flang/Optimizer/Builder/FIRBuilder.h"
#include "flang/Optimizer/Builder/Todo.h"
#include "mlir/Dialect/OpenACC/OpenACC.h"
#include "mlir/IR/Dominance.h"
#include "llvm/ADT/STLExtras.h"

namespace {
/// Helper class to detect if an alloca is inside an mlir::Block that can be
/// reached again before its deallocation points via block successors. This
/// analysis is only valid if the deallocation points are inside (or nested
/// inside) the same region as alloca because it does not consider region CFG
/// (for instance, the block inside a fir.do_loop is obviously inside a loop,
/// but is not a loop formed by blocks). The dominance of the alloca on its
/// deallocation points implies this pre-condition (although it is more
/// restrictive).
class BlockCycleDetector {
public:
  bool allocaIsInCycle(fir::AllocaOp alloca,
                       llvm::ArrayRef<mlir::Operation *> deallocationPoints);

private:
  // Cache for blocks owning alloca that have been analyzed. In many Fortran
  // programs, allocas are usually made in the same blocks with no block cycles.
  // So getting a fast "no" is beneficial.
  llvm::DenseMap<mlir::Block *, /*isInCycle*/ bool> analyzed;
};
} // namespace

namespace {
class AllocaReplaceImpl {
public:
  AllocaReplaceImpl(fir::AllocaRewriterCallBack allocaRewriter,
                    fir::DeallocCallBack deallocGenerator)
      : allocaRewriter{allocaRewriter}, deallocGenerator{deallocGenerator} {}
  bool replace(mlir::RewriterBase &, fir::AllocaOp);

private:
  mlir::Region *findDeallocationPointsAndOwner(
      fir::AllocaOp alloca,
      llvm::SmallVectorImpl<mlir::Operation *> &deallocationPoints);
  bool
  allocDominatesDealloc(fir::AllocaOp alloca,
                        llvm::ArrayRef<mlir::Operation *> deallocationPoints) {
    return llvm::all_of(deallocationPoints, [&](mlir::Operation *deallocPoint) {
      return this->dominanceInfo.properlyDominates(alloca.getOperation(),
                                                   deallocPoint);
    });
  }
  void
  genIndirectDeallocation(mlir::RewriterBase &, fir::AllocaOp,
                          llvm::ArrayRef<mlir::Operation *> deallocationPoints,
                          mlir::Value replacement, mlir::Region &owningRegion);

private:
  fir::AllocaRewriterCallBack allocaRewriter;
  fir::DeallocCallBack deallocGenerator;
  mlir::DominanceInfo dominanceInfo;
  BlockCycleDetector blockCycleDetector;
};
} // namespace

static bool
allocaIsInCycleImpl(mlir::Block *allocaBlock,
                    llvm::ArrayRef<mlir::Operation *> deallocationPoints) {
  llvm::DenseSet<mlir::Block *> seen;
  // Insert the deallocation point blocks as "seen" so that the block
  // traversal will stop at them.
  for (mlir::Operation *deallocPoint : deallocationPoints)
    seen.insert(deallocPoint->getBlock());
  if (seen.contains(allocaBlock))
    return false;
  // Traverse the block successor graph starting by the alloca block.
  llvm::SmallVector<mlir::Block *> successors{allocaBlock};
  while (!successors.empty())
    for (mlir::Block *next : successors.pop_back_val()->getSuccessors()) {
      if (next == allocaBlock)
        return true;
      if (auto pair = seen.insert(next); pair.second)
        successors.push_back(next);
    }
  // The traversal did not reach the alloca block again.
  return false;
}
bool BlockCycleDetector::allocaIsInCycle(
    fir::AllocaOp alloca,
    llvm::ArrayRef<mlir::Operation *> deallocationPoints) {
  mlir::Block *allocaBlock = alloca->getBlock();
  auto analyzedPair = analyzed.try_emplace(allocaBlock, /*isInCycle=*/false);
  bool alreadyAnalyzed = !analyzedPair.second;
  bool &isInCycle = analyzedPair.first->second;
  // Fast exit if block was already analyzed and no cycle was found.
  if (alreadyAnalyzed && !isInCycle)
    return false;
  // If the analysis was not done generically for this block, run it and
  // save the result.
  if (!alreadyAnalyzed)
    isInCycle = allocaIsInCycleImpl(allocaBlock, /*deallocationPoints*/ {});
  if (!isInCycle)
    return false;
  // If the generic analysis found a block loop, see if the deallocation
  // point would be reached before reaching the block again. Do not
  // cache that analysis that is specific to the deallocation points
  // found for this alloca.
  return allocaIsInCycleImpl(allocaBlock, deallocationPoints);
}

static bool terminatorYieldsMemory(mlir::Operation &terminator) {
  return llvm::any_of(terminator.getResults(), [](mlir::OpResult res) {
    return fir::conformsWithPassByRef(res.getType());
  });
}

static bool isRegionTerminator(mlir::Operation &terminator) {
  // Using ReturnLike trait is tempting but it is not set on
  // all region terminator that matters (like omp::TerminatorOp that
  // has no results).
  // May be true for dead code. It is not a correctness issue and dead code can
  // be eliminated by running region simplification before this utility is
  // used.
  // May also be true for unreachable like terminators (e.g., after an abort
  // call related to Fortran STOP). This is also OK, the inserted deallocation
  // will simply never be reached. It is easier for the rest of the code here
  // to assume there is always at least one deallocation point, so keep
  // unreachable terminators.
  return !terminator.hasSuccessors();
}

mlir::Region *AllocaReplaceImpl::findDeallocationPointsAndOwner(
    fir::AllocaOp alloca,
    llvm::SmallVectorImpl<mlir::Operation *> &deallocationPoints) {
  // Step 1: Identify the operation and region owning the alloca.
  mlir::Region *owningRegion = alloca.getOwnerRegion();
  if (!owningRegion)
    return nullptr;
  mlir::Operation *owningOp = owningRegion->getParentOp();
  assert(owningOp && "region expected to be owned");
  // Step 2: Identify the exit points of the owning region, they are the default
  // deallocation points. TODO: detect and use lifetime markers to get earlier
  // deallocation points.
  bool isOpenACCMPRecipe = mlir::isa<mlir::accomp::RecipeInterface>(owningOp);
  for (mlir::Block &block : owningRegion->getBlocks())
    if (mlir::Operation *terminator = block.getTerminator();
        isRegionTerminator(*terminator)) {
      // FIXME: OpenACC and OpenMP privatization recipe are stand alone
      // operation meant to be later "inlined", the value they return may
      // be the address of a local alloca. It would be incorrect to insert
      // deallocation before the terminator (this would introduce use after
      // free once the recipe is inlined.
      // This probably require redesign or special handling on the OpenACC/MP
      // side.
      if (isOpenACCMPRecipe && terminatorYieldsMemory(*terminator))
        return nullptr;
      deallocationPoints.push_back(terminator);
    }
  // If no block terminators without successors have been found, this is
  // an odd region we cannot reason about (never seen yet in FIR and
  // mainstream dialects, but MLIR does not really prevent it).
  if (deallocationPoints.empty())
    return nullptr;

  // Step 3: detect block based loops between the allocation and deallocation
  // points, and add a deallocation point on the back edge to avoid memory
  // leaks.
  // The detection avoids doing region CFG analysis by assuming that there may
  // be cycles if deallocation points are not dominated by the alloca.
  // This leaves the cases where the deallocation points are in the same region
  // as the alloca (or nested inside it). In which cases there may be a back
  // edge between the alloca and the deallocation point via block successors. An
  // analysis is run to detect those cases.
  // When a loop is detected, the easiest solution to deallocate on the back
  // edge is to store the allocated memory address in a variable (that dominates
  // the loops) and to deallocate the address in that variable if it is set
  // before executing the allocation. This strategy still leads to correct
  // execution in the "false positive" cases.
  // Hence, the alloca is added as a deallocation point when there is no
  // dominance. Note that bringing lifetime markers above will reduce the
  // false positives.
  if (!allocDominatesDealloc(alloca, deallocationPoints) ||
      blockCycleDetector.allocaIsInCycle(alloca, deallocationPoints))
    deallocationPoints.push_back(alloca.getOperation());
  return owningRegion;
}

void AllocaReplaceImpl::genIndirectDeallocation(
    mlir::RewriterBase &rewriter, fir::AllocaOp alloca,
    llvm::ArrayRef<mlir::Operation *> deallocationPoints,
    mlir::Value replacement, mlir::Region &owningRegion) {
  mlir::Location loc = alloca.getLoc();
  auto replacementInsertPoint = rewriter.saveInsertionPoint();
  // Create C pointer variable in the entry block to store the alloc
  // and access it indirectly in the entry points that do not dominate.
  rewriter.setInsertionPointToStart(&owningRegion.front());
  mlir::Type heapType = fir::HeapType::get(alloca.getInType());
  mlir::Value ptrVar = rewriter.create<fir::AllocaOp>(loc, heapType);
  mlir::Value nullPtr = rewriter.create<fir::ZeroOp>(loc, heapType);
  rewriter.create<fir::StoreOp>(loc, nullPtr, ptrVar);
  // TODO: introducing a pointer compare op in FIR would help
  // generating less IR here.
  mlir::Type intPtrTy = fir::getIntPtrType(rewriter);
  mlir::Value c0 = rewriter.create<mlir::arith::ConstantOp>(
      loc, intPtrTy, rewriter.getIntegerAttr(intPtrTy, 0));

  // Store new storage address right after its creation.
  rewriter.restoreInsertionPoint(replacementInsertPoint);
  mlir::Value castReplacement =
      fir::factory::createConvert(rewriter, loc, heapType, replacement);
  rewriter.create<fir::StoreOp>(loc, castReplacement, ptrVar);

  // Generate conditional deallocation at every deallocation point.
  auto genConditionalDealloc = [&](mlir::Location loc) {
    mlir::Value ptrVal = rewriter.create<fir::LoadOp>(loc, ptrVar);
    mlir::Value ptrToInt =
        rewriter.create<fir::ConvertOp>(loc, intPtrTy, ptrVal);
    mlir::Value isAllocated = rewriter.create<mlir::arith::CmpIOp>(
        loc, mlir::arith::CmpIPredicate::ne, ptrToInt, c0);
    auto ifOp = rewriter.create<fir::IfOp>(loc, std::nullopt, isAllocated,
                                           /*withElseRegion=*/false);
    rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front());
    mlir::Value cast = fir::factory::createConvert(
        rewriter, loc, replacement.getType(), ptrVal);
    deallocGenerator(loc, rewriter, cast);
    // Currently there is no need to reset the pointer var because two
    // deallocation points can never be reached without going through the
    // alloca.
    rewriter.setInsertionPointAfter(ifOp);
  };
  for (mlir::Operation *deallocPoint : deallocationPoints) {
    rewriter.setInsertionPoint(deallocPoint);
    genConditionalDealloc(deallocPoint->getLoc());
  }
}

bool AllocaReplaceImpl::replace(mlir::RewriterBase &rewriter,
                                fir::AllocaOp alloca) {
  llvm::SmallVector<mlir::Operation *> deallocationPoints;
  mlir::Region *owningRegion =
      findDeallocationPointsAndOwner(alloca, deallocationPoints);
  if (!owningRegion)
    return false;
  rewriter.setInsertionPointAfter(alloca.getOperation());
  bool deallocPointsDominateAlloc =
      allocDominatesDealloc(alloca, deallocationPoints);
  if (mlir::Value replacement =
          allocaRewriter(rewriter, alloca, deallocPointsDominateAlloc)) {
    mlir::Value castReplacement = fir::factory::createConvert(
        rewriter, alloca.getLoc(), alloca.getType(), replacement);
    if (deallocPointsDominateAlloc)
      for (mlir::Operation *deallocPoint : deallocationPoints) {
        rewriter.setInsertionPoint(deallocPoint);
        deallocGenerator(deallocPoint->getLoc(), rewriter, replacement);
      }
    else
      genIndirectDeallocation(rewriter, alloca, deallocationPoints, replacement,
                              *owningRegion);
    rewriter.replaceOp(alloca, castReplacement);
  }
  return true;
}

bool fir::replaceAllocas(mlir::RewriterBase &rewriter,
                         mlir::Operation *parentOp,
                         MustRewriteCallBack mustReplace,
                         AllocaRewriterCallBack allocaRewriter,
                         DeallocCallBack deallocGenerator) {
  // If the parent operation is not an alloca owner, the code below would risk
  // modifying IR outside of parentOp.
  if (!fir::AllocaOp::ownsNestedAlloca(parentOp))
    return false;
  auto insertPoint = rewriter.saveInsertionPoint();
  bool replacedAllRequestedAlloca = true;
  AllocaReplaceImpl impl(allocaRewriter, deallocGenerator);
  parentOp->walk([&](fir::AllocaOp alloca) {
    if (mustReplace(alloca))
      replacedAllRequestedAlloca &= impl.replace(rewriter, alloca);
  });
  rewriter.restoreInsertionPoint(insertPoint);
  return replacedAllRequestedAlloca;
}