//===- FunctionSupport.cpp - Utility types for function-like ops ----------===//
//
// 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 "mlir/IR/FunctionSupport.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/BitVector.h"

using namespace mlir;

/// Helper to call a callback once on each index in the range
/// [0, `totalIndices`), *except* for the indices given in `indices`.
/// `indices` is allowed to have duplicates and can be in any order.
inline void iterateIndicesExcept(unsigned totalIndices,
                                 ArrayRef<unsigned> indices,
                                 function_ref<void(unsigned)> callback) {
  llvm::BitVector skipIndices(totalIndices);
  for (unsigned i : indices)
    skipIndices.set(i);

  for (unsigned i = 0; i < totalIndices; ++i)
    if (!skipIndices.test(i))
      callback(i);
}

//===----------------------------------------------------------------------===//
// Function Arguments and Results.
//===----------------------------------------------------------------------===//

static bool isEmptyAttrDict(Attribute attr) {
  return attr.cast<DictionaryAttr>().empty();
}

DictionaryAttr mlir::function_like_impl::getArgAttrDict(Operation *op,
                                                        unsigned index) {
  ArrayAttr attrs = op->getAttrOfType<ArrayAttr>(getArgDictAttrName());
  DictionaryAttr argAttrs =
      attrs ? attrs[index].cast<DictionaryAttr>() : DictionaryAttr();
  return (argAttrs && !argAttrs.empty()) ? argAttrs : DictionaryAttr();
}

DictionaryAttr mlir::function_like_impl::getResultAttrDict(Operation *op,
                                                           unsigned index) {
  ArrayAttr attrs = op->getAttrOfType<ArrayAttr>(getResultDictAttrName());
  DictionaryAttr resAttrs =
      attrs ? attrs[index].cast<DictionaryAttr>() : DictionaryAttr();
  return (resAttrs && !resAttrs.empty()) ? resAttrs : DictionaryAttr();
}

void mlir::function_like_impl::detail::setArgResAttrDict(
    Operation *op, StringRef attrName, unsigned numTotalIndices, unsigned index,
    DictionaryAttr attrs) {
  ArrayAttr allAttrs = op->getAttrOfType<ArrayAttr>(attrName);
  if (!allAttrs) {
    if (attrs.empty())
      return;

    // If this attribute is not empty, we need to create a new attribute array.
    SmallVector<Attribute, 8> newAttrs(numTotalIndices,
                                       DictionaryAttr::get(op->getContext()));
    newAttrs[index] = attrs;
    op->setAttr(attrName, ArrayAttr::get(op->getContext(), newAttrs));
    return;
  }
  // Check to see if the attribute is different from what we already have.
  if (allAttrs[index] == attrs)
    return;

  // If it is, check to see if the attribute array would now contain only empty
  // dictionaries.
  ArrayRef<Attribute> rawAttrArray = allAttrs.getValue();
  if (attrs.empty() &&
      llvm::all_of(rawAttrArray.take_front(index), isEmptyAttrDict) &&
      llvm::all_of(rawAttrArray.drop_front(index + 1), isEmptyAttrDict)) {
    op->removeAttr(attrName);
    return;
  }

  // Otherwise, create a new attribute array with the updated dictionary.
  SmallVector<Attribute, 8> newAttrs(rawAttrArray.begin(), rawAttrArray.end());
  newAttrs[index] = attrs;
  op->setAttr(attrName, ArrayAttr::get(op->getContext(), newAttrs));
}

/// Set all of the argument or result attribute dictionaries for a function.
static void setAllArgResAttrDicts(Operation *op, StringRef attrName,
                                  ArrayRef<Attribute> attrs) {
  if (llvm::all_of(attrs, isEmptyAttrDict))
    op->removeAttr(attrName);
  else
    op->setAttr(attrName, ArrayAttr::get(op->getContext(), attrs));
}

void mlir::function_like_impl::setAllArgAttrDicts(
    Operation *op, ArrayRef<DictionaryAttr> attrs) {
  setAllArgAttrDicts(op, ArrayRef<Attribute>(attrs.data(), attrs.size()));
}
void mlir::function_like_impl::setAllArgAttrDicts(Operation *op,
                                                  ArrayRef<Attribute> attrs) {
  auto wrappedAttrs = llvm::map_range(attrs, [op](Attribute attr) -> Attribute {
    return !attr ? DictionaryAttr::get(op->getContext()) : attr;
  });
  setAllArgResAttrDicts(op, getArgDictAttrName(),
                        llvm::to_vector<8>(wrappedAttrs));
}

void mlir::function_like_impl::setAllResultAttrDicts(
    Operation *op, ArrayRef<DictionaryAttr> attrs) {
  setAllResultAttrDicts(op, ArrayRef<Attribute>(attrs.data(), attrs.size()));
}
void mlir::function_like_impl::setAllResultAttrDicts(
    Operation *op, ArrayRef<Attribute> attrs) {
  auto wrappedAttrs = llvm::map_range(attrs, [op](Attribute attr) -> Attribute {
    return !attr ? DictionaryAttr::get(op->getContext()) : attr;
  });
  setAllArgResAttrDicts(op, getResultDictAttrName(),
                        llvm::to_vector<8>(wrappedAttrs));
}

void mlir::function_like_impl::eraseFunctionArguments(
    Operation *op, ArrayRef<unsigned> argIndices, unsigned originalNumArgs,
    Type newType) {
  // There are 3 things that need to be updated:
  // - Function type.
  // - Arg attrs.
  // - Block arguments of entry block.
  Block &entry = op->getRegion(0).front();

  // Update the argument attributes of the function.
  if (auto argAttrs = op->getAttrOfType<ArrayAttr>(getArgDictAttrName())) {
    SmallVector<DictionaryAttr, 4> newArgAttrs;
    newArgAttrs.reserve(argAttrs.size());
    iterateIndicesExcept(originalNumArgs, argIndices, [&](unsigned i) {
      newArgAttrs.emplace_back(argAttrs[i].cast<DictionaryAttr>());
    });
    setAllArgAttrDicts(op, newArgAttrs);
  }

  // Update the function type and any entry block arguments.
  op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
  entry.eraseArguments(argIndices);
}

void mlir::function_like_impl::eraseFunctionResults(
    Operation *op, ArrayRef<unsigned> resultIndices,
    unsigned originalNumResults, Type newType) {
  // There are 2 things that need to be updated:
  // - Function type.
  // - Result attrs.

  // Update the result attributes of the function.
  if (auto resAttrs = op->getAttrOfType<ArrayAttr>(getResultDictAttrName())) {
    SmallVector<DictionaryAttr, 4> newResultAttrs;
    newResultAttrs.reserve(resAttrs.size());
    iterateIndicesExcept(originalNumResults, resultIndices, [&](unsigned i) {
      newResultAttrs.emplace_back(resAttrs[i].cast<DictionaryAttr>());
    });
    setAllResultAttrDicts(op, newResultAttrs);
  }

  // Update the function type.
  op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
}

//===----------------------------------------------------------------------===//
// Function type signature.
//===----------------------------------------------------------------------===//

FunctionType mlir::function_like_impl::getFunctionType(Operation *op) {
  assert(op->hasTrait<OpTrait::FunctionLike>());
  return op->getAttrOfType<TypeAttr>(getTypeAttrName())
      .getValue()
      .cast<FunctionType>();
}

void mlir::function_like_impl::setFunctionType(Operation *op,
                                               FunctionType newType) {
  assert(op->hasTrait<OpTrait::FunctionLike>());
  FunctionType oldType = getFunctionType(op);
  op->setAttr(getTypeAttrName(), TypeAttr::get(newType));

  // Functor used to update the argument and result attributes of the function.
  auto updateAttrFn = [&](StringRef attrName, unsigned oldCount,
                          unsigned newCount, auto setAttrFn) {
    if (oldCount == newCount)
      return;
    // The new type has no arguments/results, just drop the attribute.
    if (newCount == 0) {
      op->removeAttr(attrName);
      return;
    }
    ArrayAttr attrs = op->getAttrOfType<ArrayAttr>(attrName);
    if (!attrs)
      return;

    // The new type has less arguments/results, take the first N attributes.
    if (newCount < oldCount)
      return setAttrFn(op, attrs.getValue().take_front(newCount));

    // Otherwise, the new type has more arguments/results. Initialize the new
    // arguments/results with empty attributes.
    SmallVector<Attribute> newAttrs(attrs.begin(), attrs.end());
    newAttrs.resize(newCount);
    setAttrFn(op, newAttrs);
  };

  // Update the argument and result attributes.
  updateAttrFn(function_like_impl::getArgDictAttrName(), oldType.getNumInputs(),
               newType.getNumInputs(), [&](Operation *op, auto &&attrs) {
                 setAllArgAttrDicts(op, attrs);
               });
  updateAttrFn(
      function_like_impl::getResultDictAttrName(), oldType.getNumResults(),
      newType.getNumResults(),
      [&](Operation *op, auto &&attrs) { setAllResultAttrDicts(op, attrs); });
}

//===----------------------------------------------------------------------===//
// Function body.
//===----------------------------------------------------------------------===//

Region &mlir::function_like_impl::getFunctionBody(Operation *op) {
  assert(op->hasTrait<OpTrait::FunctionLike>());
  return op->getRegion(0);
}