//===- OMPDescriptorMapInfoGen.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
//
//===----------------------------------------------------------------------===//

//===----------------------------------------------------------------------===//
/// \file
/// An OpenMP dialect related pass for FIR/HLFIR which expands MapInfoOp's
/// containing descriptor related types (fir::BoxType's) into multiple
/// MapInfoOp's containing the parent descriptor and pointer member components
/// for individual mapping, treating the descriptor type as a record type for
/// later lowering in the OpenMP dialect.
//===----------------------------------------------------------------------===//

#include "flang/Optimizer/Builder/FIRBuilder.h"
#include "flang/Optimizer/Dialect/FIRType.h"
#include "flang/Optimizer/Dialect/Support/KindMapping.h"
#include "flang/Optimizer/Transforms/Passes.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
#include "mlir/IR/BuiltinDialect.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/SmallPtrSet.h"
#include <iterator>

namespace fir {
#define GEN_PASS_DEF_OMPDESCRIPTORMAPINFOGENPASS
#include "flang/Optimizer/Transforms/Passes.h.inc"
} // namespace fir

namespace {
class OMPDescriptorMapInfoGenPass
    : public fir::impl::OMPDescriptorMapInfoGenPassBase<
          OMPDescriptorMapInfoGenPass> {

  void genDescriptorMemberMaps(mlir::omp::MapInfoOp op,
                               fir::FirOpBuilder &builder,
                               mlir::Operation *target) {
    mlir::Location loc = builder.getUnknownLoc();
    mlir::Value descriptor = op.getVarPtr();

    // If we enter this function, but the mapped type itself is not the
    // descriptor, then it's likely the address of the descriptor so we
    // must retrieve the descriptor SSA.
    if (!fir::isTypeWithDescriptor(op.getVarType())) {
      if (auto addrOp = mlir::dyn_cast_if_present<fir::BoxAddrOp>(
              op.getVarPtr().getDefiningOp())) {
        descriptor = addrOp.getVal();
      }
    }

    // The fir::BoxOffsetOp only works with !fir.ref<!fir.box<...>> types, as
    // allowing it to access non-reference box operations can cause some
    // problematic SSA IR. However, in the case of assumed shape's the type
    // is not a !fir.ref, in these cases to retrieve the appropriate
    // !fir.ref<!fir.box<...>> to access the data we need to map we must
    // perform an alloca and then store to it and retrieve the data from the new
    // alloca.
    if (mlir::isa<fir::BaseBoxType>(descriptor.getType())) {
      mlir::OpBuilder::InsertPoint insPt = builder.saveInsertionPoint();
      builder.setInsertionPointToStart(builder.getAllocaBlock());
      auto alloca = builder.create<fir::AllocaOp>(loc, descriptor.getType());
      builder.restoreInsertionPoint(insPt);
      builder.create<fir::StoreOp>(loc, descriptor, alloca);
      descriptor = alloca;
    }

    mlir::Value baseAddrAddr = builder.create<fir::BoxOffsetOp>(
        loc, descriptor, fir::BoxFieldAttr::base_addr);

    // Member of the descriptor pointing at the allocated data
    mlir::Value baseAddr = builder.create<mlir::omp::MapInfoOp>(
        loc, baseAddrAddr.getType(), descriptor,
        llvm::cast<mlir::omp::PointerLikeType>(
            fir::unwrapRefType(baseAddrAddr.getType()))
            .getElementType(),
        baseAddrAddr, mlir::SmallVector<mlir::Value>{}, op.getBounds(),
        builder.getIntegerAttr(builder.getIntegerType(64, false),
                               op.getMapType().value()),
        builder.getAttr<mlir::omp::VariableCaptureKindAttr>(
            mlir::omp::VariableCaptureKind::ByRef),
        builder.getStringAttr("") /*name*/);

    // TODO: map the addendum segment of the descriptor, similarly to the
    // above base address/data pointer member.

    if (auto mapClauseOwner =
            llvm::dyn_cast<mlir::omp::MapClauseOwningOpInterface>(target)) {
      llvm::SmallVector<mlir::Value> newMapOps;
      mlir::OperandRange mapOperandsArr = mapClauseOwner.getMapOperands();

      for (size_t i = 0; i < mapOperandsArr.size(); ++i) {
        if (mapOperandsArr[i] == op) {
          // Push new implicit maps generated for the descriptor.
          newMapOps.push_back(baseAddr);

          // for TargetOp's which have IsolatedFromAbove we must align the
          // new additional map operand with an appropriate BlockArgument,
          // as the printing and later processing currently requires a 1:1
          // mapping of BlockArgs to MapInfoOp's at the same placement in
          // each array (BlockArgs and MapOperands).
          if (auto targetOp = llvm::dyn_cast<mlir::omp::TargetOp>(target))
            targetOp.getRegion().insertArgument(i, baseAddr.getType(), loc);
        }
        newMapOps.push_back(mapOperandsArr[i]);
      }
      mapClauseOwner.getMapOperandsMutable().assign(newMapOps);
    }

    mlir::Value newDescParentMapOp = builder.create<mlir::omp::MapInfoOp>(
        op->getLoc(), op.getResult().getType(), descriptor,
        fir::unwrapRefType(descriptor.getType()), mlir::Value{},
        mlir::SmallVector<mlir::Value>{baseAddr},
        mlir::SmallVector<mlir::Value>{},
        builder.getIntegerAttr(builder.getIntegerType(64, false),
                               op.getMapType().value()),
        op.getMapCaptureTypeAttr(), op.getNameAttr());
    op.replaceAllUsesWith(newDescParentMapOp);
    op->erase();
  }

  // This pass executes on mlir::ModuleOp's finding omp::MapInfoOp's containing
  // descriptor based types (allocatables, pointers, assumed shape etc.) and
  // expanding them into multiple omp::MapInfoOp's for each pointer member
  // contained within the descriptor.
  void runOnOperation() override {
    mlir::func::FuncOp func = getOperation();
    mlir::ModuleOp module = func->getParentOfType<mlir::ModuleOp>();
    fir::KindMapping kindMap = fir::getKindMapping(module);
    fir::FirOpBuilder builder{module, std::move(kindMap)};

    func->walk([&](mlir::omp::MapInfoOp op) {
      if (fir::isTypeWithDescriptor(op.getVarType()) ||
          mlir::isa_and_present<fir::BoxAddrOp>(
              op.getVarPtr().getDefiningOp())) {
        builder.setInsertionPoint(op);
        // TODO: Currently only supports a single user for the MapInfoOp, this
        // is fine for the moment as the Fortran Frontend will generate a
        // new MapInfoOp per Target operation for the moment. However, when/if
        // we optimise/cleanup the IR, it likely isn't too difficult to
        // extend this function, it would require some modification to create a
        // single new MapInfoOp per new MapInfoOp generated and share it across
        // all users appropriately, making sure to only add a single member link
        // per new generation for the original originating descriptor MapInfoOp.
        assert(llvm::hasSingleElement(op->getUsers()) &&
               "OMPDescriptorMapInfoGen currently only supports single users "
               "of a MapInfoOp");
        genDescriptorMemberMaps(op, builder, *op->getUsers().begin());
      }
    });
  }
};

} // namespace

namespace fir {
std::unique_ptr<mlir::Pass> createOMPDescriptorMapInfoGenPass() {
  return std::make_unique<OMPDescriptorMapInfoGenPass>();
}
} // namespace fir