//===- FunctionImplementation.cpp - Utilities 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/FunctionImplementation.h"

#include "mlir/IR/Builders.h"

#include "mlir/IR/FunctionSupport.h"

#include "mlir/IR/SymbolTable.h"

using namespace mlir;

static ParseResult

parseArgumentList(OpAsmParser &parser, bool allowVariadic,

                  SmallVectorImpl<Type> &argTypes,

                  SmallVectorImpl<OpAsmParser::OperandType> &argNames,

                  SmallVectorImpl<NamedAttrList> &argAttrs, bool &isVariadic) {

  if (parser.parseLParen())

    return failure();

  // The argument list either has to consistently have ssa-id's followed by

  // types, or just be a type list.  It isn't ok to sometimes have SSA ID's and

  // sometimes not.

  auto parseArgument = [&]() -> ParseResult {

    llvm::SMLoc loc = parser.getCurrentLocation();

    // Parse argument name if present.

    OpAsmParser::OperandType argument;

    Type argumentType;

    if (succeeded(parser.parseOptionalRegionArgument(argument)) &&

        !argument.name.empty()) {

      // Reject this if the preceding argument was missing a name.

      if (argNames.empty() && !argTypes.empty())

        return parser.emitError(loc, "expected type instead of SSA identifier");

      argNames.push_back(argument);

      if (parser.parseColonType(argumentType))

        return failure();

    } else if (allowVariadic && succeeded(parser.parseOptionalEllipsis())) {

      isVariadic = true;

      return success();

    } else if (!argNames.empty()) {

      // Reject this if the preceding argument had a name.

      return parser.emitError(loc, "expected SSA identifier");

    } else if (parser.parseType(argumentType)) {

      return failure();

    }

    // Add the argument type.

    argTypes.push_back(argumentType);

    // Parse any argument attributes.

    NamedAttrList attrs;

    if (parser.parseOptionalAttrDict(attrs))

      return failure();

    argAttrs.push_back(attrs);

    return success();

  };

  // Parse the function arguments.

  isVariadic = false;

  if (failed(parser.parseOptionalRParen())) {

    do {

      unsigned numTypedArguments = argTypes.size();

      if (parseArgument())

        return failure();

      llvm::SMLoc loc = parser.getCurrentLocation();

      if (argTypes.size() == numTypedArguments &&

          succeeded(parser.parseOptionalComma()))

        return parser.emitError(

            loc, "variadic arguments must be in the end of the argument list");

    } while (succeeded(parser.parseOptionalComma()));

    parser.parseRParen();

  }

  return success();

}

/// Parse a function result list.

///

///   function-result-list ::= function-result-list-parens

///                          | non-function-type

///   function-result-list-parens ::= `(` `)`

///                                 | `(` function-result-list-no-parens `)`

///   function-result-list-no-parens ::= function-result (`,` function-result)*

///   function-result ::= type attribute-dict?

///

static ParseResult

parseFunctionResultList(OpAsmParser &parser, SmallVectorImpl<Type> &resultTypes,

                        SmallVectorImpl<NamedAttrList> &resultAttrs) {

  if (failed(parser.parseOptionalLParen())) {

    // We already know that there is no `(`, so parse a type.

    // Because there is no `(`, it cannot be a function type.

    Type ty;

    if (parser.parseType(ty))

      return failure();

    resultTypes.push_back(ty);

    resultAttrs.emplace_back();

    return success();

  }

  // Special case for an empty set of parens.

  if (succeeded(parser.parseOptionalRParen()))

    return success();

  // Parse individual function results.

  do {

    resultTypes.emplace_back();

    resultAttrs.emplace_back();

    if (parser.parseType(resultTypes.back()) ||

        parser.parseOptionalAttrDict(resultAttrs.back())) {

      return failure();

    }

  } while (succeeded(parser.parseOptionalComma()));

  return parser.parseRParen();

}

/// Parses a function signature using `parser`. The `allowVariadic` argument

/// indicates whether functions with variadic arguments are supported. The

/// trailing arguments are populated by this function with names, types and

/// attributes of the arguments and those of the results.

ParseResult mlir::impl::parseFunctionSignature(

    OpAsmParser &parser, bool allowVariadic,

    SmallVectorImpl<OpAsmParser::OperandType> &argNames,

    SmallVectorImpl<Type> &argTypes, SmallVectorImpl<NamedAttrList> &argAttrs,

    bool &isVariadic, SmallVectorImpl<Type> &resultTypes,

    SmallVectorImpl<NamedAttrList> &resultAttrs) {

  if (parseArgumentList(parser, allowVariadic, argTypes, argNames, argAttrs,

                        isVariadic))

    return failure();

  if (succeeded(parser.parseOptionalArrow()))

    return parseFunctionResultList(parser, resultTypes, resultAttrs);

  return success();

}

void mlir::impl::addArgAndResultAttrs(Builder &builder, OperationState &result,

                                      ArrayRef<NamedAttrList> argAttrs,

                                      ArrayRef<NamedAttrList> resultAttrs) {

  // Add the attributes to the function arguments.

  SmallString<8> attrNameBuf;

  for (unsigned i = 0, e = argAttrs.size(); i != e; ++i)

    if (!argAttrs[i].empty())

      result.addAttribute(getArgAttrName(i, attrNameBuf),

                          builder.getDictionaryAttr(argAttrs[i]));

  // Add the attributes to the function results.

  for (unsigned i = 0, e = resultAttrs.size(); i != e; ++i)

    if (!resultAttrs[i].empty())

      result.addAttribute(getResultAttrName(i, attrNameBuf),

                          builder.getDictionaryAttr(resultAttrs[i]));

}

/// Parser implementation for function-like operations.  Uses `funcTypeBuilder`

/// to construct the custom function type given lists of input and output types.

ParseResult

mlir::impl::parseFunctionLikeOp(OpAsmParser &parser, OperationState &result,

                                bool allowVariadic,

                                mlir::impl::FuncTypeBuilder funcTypeBuilder) {

  SmallVector<OpAsmParser::OperandType, 4> entryArgs;

  SmallVector<NamedAttrList, 4> argAttrs;

  SmallVector<NamedAttrList, 4> resultAttrs;

  SmallVector<Type, 4> argTypes;

  SmallVector<Type, 4> resultTypes;

  auto &builder = parser.getBuilder();

  // Parse the name as a symbol.

  StringAttr nameAttr;

  if (parser.parseSymbolName(nameAttr, ::mlir::SymbolTable::getSymbolAttrName(),

                             result.attributes))

    return failure();

  // Parse the function signature.

  auto signatureLocation = parser.getCurrentLocation();

  bool isVariadic = false;

  if (parseFunctionSignature(parser, allowVariadic, entryArgs, argTypes,

                             argAttrs, isVariadic, resultTypes, resultAttrs))

    return failure();

  std::string errorMessage;

  if (auto type = funcTypeBuilder(builder, argTypes, resultTypes,

                                  impl::VariadicFlag(isVariadic), errorMessage))

    result.addAttribute(getTypeAttrName(), TypeAttr::get(type));

  else

    return parser.emitError(signatureLocation)

           << "failed to construct function type"

           << (errorMessage.empty() ? "" : ": ") << errorMessage;

  // If function attributes are present, parse them.

  if (parser.parseOptionalAttrDictWithKeyword(result.attributes))

    return failure();

  // Add the attributes to the function arguments.

  assert(argAttrs.size() == argTypes.size());

  assert(resultAttrs.size() == resultTypes.size());

  addArgAndResultAttrs(builder, result, argAttrs, resultAttrs);

  // Parse the optional function body.

  auto *body = result.addRegion();

  return parser.parseOptionalRegion(

      *body, entryArgs, entryArgs.empty() ? ArrayRef<Type>() : argTypes);

}

// Print a function result list.

static void printFunctionResultList(OpAsmPrinter &p, ArrayRef<Type> types,

                                    ArrayRef<ArrayRef<NamedAttribute>> attrs) {

  assert(!types.empty() && "Should not be called for empty result list.");

  auto &os = p.getStream();

  bool needsParens =

      types.size() > 1 || types[0].isa<FunctionType>() || !attrs[0].empty();

  if (needsParens)

    os << '(';

  llvm::interleaveComma(

      llvm::zip(types, attrs), os,

      [&](const std::tuple<Type, ArrayRef<NamedAttribute>> &t) {

        p.printType(std::get<0>(t));

        p.printOptionalAttrDict(std::get<1>(t));

      });

  if (needsParens)

    os << ')';

}

/// Print the signature of the function-like operation `op`.  Assumes `op` has

/// the FunctionLike trait and passed the verification.

void mlir::impl::printFunctionSignature(OpAsmPrinter &p, Operation *op,

                                        ArrayRef<Type> argTypes,

                                        bool isVariadic,

                                        ArrayRef<Type> resultTypes) {

  Region &body = op->getRegion(0);

  bool isExternal = body.empty();

  p << '(';

  for (unsigned i = 0, e = argTypes.size(); i < e; ++i) {

    if (i > 0)

      p << ", ";

    if (!isExternal) {

      p.printOperand(body.front().getArgument(i));

      p << ": ";

    }

    p.printType(argTypes[i]);

    p.printOptionalAttrDict(::mlir::impl::getArgAttrs(op, i));

  }

  if (isVariadic) {

    if (!argTypes.empty())

      p << ", ";

    p << "...";

  }

  p << ')';

  if (!resultTypes.empty()) {

    p.getStream() << " -> ";

    SmallVector<ArrayRef<NamedAttribute>, 4> resultAttrs;

    for (int i = 0, e = resultTypes.size(); i < e; ++i)

      resultAttrs.push_back(::mlir::impl::getResultAttrs(op, i));

    printFunctionResultList(p, resultTypes, resultAttrs);

  }

}

/// Prints the list of function prefixed with the "attributes" keyword. The

/// attributes with names listed in "elided" as well as those used by the

/// function-like operation internally are not printed. Nothing is printed

/// if all attributes are elided. Assumes `op` has the `FunctionLike` trait and

/// passed the verification.

void mlir::impl::printFunctionAttributes(OpAsmPrinter &p, Operation *op,

                                         unsigned numInputs,

                                         unsigned numResults,

                                         ArrayRef<StringRef> elided) {

  // Print out function attributes, if present.

  SmallVector<StringRef, 2> ignoredAttrs = {

      ::mlir::SymbolTable::getSymbolAttrName(), getTypeAttrName()};

  ignoredAttrs.append(elided.begin(), elided.end());

  SmallString<8> attrNameBuf;

  // Ignore any argument attributes.

  std::vector<SmallString<8>> argAttrStorage;

  for (unsigned i = 0; i != numInputs; ++i)

    if (op->getAttr(getArgAttrName(i, attrNameBuf)))

      argAttrStorage.emplace_back(attrNameBuf);

  ignoredAttrs.append(argAttrStorage.begin(), argAttrStorage.end());

  // Ignore any result attributes.

  std::vector<SmallString<8>> resultAttrStorage;

  for (unsigned i = 0; i != numResults; ++i)

    if (op->getAttr(getResultAttrName(i, attrNameBuf)))

      resultAttrStorage.emplace_back(attrNameBuf);

  ignoredAttrs.append(resultAttrStorage.begin(), resultAttrStorage.end());

  p.printOptionalAttrDictWithKeyword(op->getAttrs(), ignoredAttrs);

}

/// Printer implementation for function-like operations.  Accepts lists of

/// argument and result types to use while printing.

void mlir::impl::printFunctionLikeOp(OpAsmPrinter &p, Operation *op,

                                     ArrayRef<Type> argTypes, bool isVariadic,

                                     ArrayRef<Type> resultTypes) {

  // Print the operation and the function name.

  auto funcName =

      op->getAttrOfType<StringAttr>(::mlir::SymbolTable::getSymbolAttrName())

          .getValue();

  p << op->getName() << ' ';

  p.printSymbolName(funcName);

  printFunctionSignature(p, op, argTypes, isVariadic, resultTypes);

  printFunctionAttributes(p, op, argTypes.size(), resultTypes.size());

  // Print the body if this is not an external function.

  Region &body = op->getRegion(0);

  if (!body.empty())

    p.printRegion(body, /*printEntryBlockArgs=*/false,

                  /*printBlockTerminators=*/true);

}