//===- ToyDialect.cpp - Toy IR Dialect registration in MLIR ---------------===//
//
// Copyright 2019 The MLIR Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// =============================================================================
//
// This file implements the dialect for the Toy IR: custom type parsing and
// operation verification.
//
//===----------------------------------------------------------------------===//

#include "toy/Dialect.h"

#include "mlir/IR/Builders.h"
#include "mlir/IR/StandardTypes.h"
#include "mlir/Support/STLExtras.h"
#include "llvm/ADT/iterator_range.h"
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Support/Regex.h"
#include "llvm/Support/raw_ostream.h"

using llvm::ArrayRef;
using llvm::raw_ostream;
using llvm::raw_string_ostream;
using llvm::SmallVector;
using llvm::StringRef;
using llvm::Twine;

namespace toy {
namespace detail {

/// This class holds the implementation of the ToyArrayType.
/// It is intended to be uniqued based on its content and owned by the context.
struct ToyArrayTypeStorage : public mlir::TypeStorage {
  /// This defines how we unique this type in the context: our key contains
  /// only the shape, a more complex type would have multiple entries in the
  /// tuple here.
  /// The element of the tuples usually matches 1-1 the arguments from the
  /// public `get()` method arguments from the facade.
  using KeyTy = std::tuple<ArrayRef<int64_t>>;
  static unsigned hashKey(const KeyTy &key) {
    return llvm::hash_combine(std::get<0>(key));
  }
  /// When the key hash hits an existing type, we compare the shape themselves
  /// to confirm we have the right type.
  bool operator==(const KeyTy &key) const { return key == KeyTy(getShape()); }

  /// This is a factory method to create our type storage. It is only
  /// invoked after looking up the type in the context using the key and not
  /// finding it.
  static ToyArrayTypeStorage *construct(mlir::TypeStorageAllocator &allocator,
                                        const KeyTy &key) {
    // Copy the shape array into the bumpptr allocator owned by the context.
    ArrayRef<int64_t> shape = allocator.copyInto(std::get<0>(key));

    // Allocate the instance for the ToyArrayTypeStorage itself
    auto *storage = allocator.allocate<ToyArrayTypeStorage>();
    // Initialize the instance using placement new.
    return new (storage) ToyArrayTypeStorage(shape);
  }

  ArrayRef<int64_t> getShape() const { return shape; }

private:
  ArrayRef<int64_t> shape;

  /// Constructor is only invoked from the `construct()` method above.
  ToyArrayTypeStorage(ArrayRef<int64_t> shape) : shape(shape) {}
};

} // namespace detail

mlir::Type ToyArrayType::getElementType() {
  return mlir::FloatType::getF64(getContext());
}

ToyArrayType ToyArrayType::get(mlir::MLIRContext *context,
                               ArrayRef<int64_t> shape) {
  return Base::get(context, ToyTypeKind::TOY_ARRAY, shape);
}

ArrayRef<int64_t> ToyArrayType::getShape() { return getImpl()->getShape(); }

mlir::MemRefType ToyArrayType::toMemref() {
  auto memRefType = mlir::MemRefType::get(getShape(), getElementType(), {}, 0);
  return memRefType;
}

/// Dialect creation, the instance will be owned by the context. This is the
/// point of registration of custom types and operations for the dialect.
ToyDialect::ToyDialect(mlir::MLIRContext *ctx) : mlir::Dialect("toy", ctx) {
  addOperations<ConstantOp, GenericCallOp, PrintOp, TransposeOp, ReshapeOp,
                MulOp, AddOp, ReturnOp, AllocOp, TypeCastOp>();
  addTypes<ToyArrayType>();
}

/// Parse a type registered to this dialect, we expect only Toy arrays.
mlir::Type ToyDialect::parseType(StringRef tyData, mlir::Location loc) const {
  // Sanity check: we only support array or array<...>
  if (!tyData.startswith("array")) {
    getContext()->emitError(loc, "Invalid Toy type '" + tyData +
                                     "', array expected");
    return nullptr;
  }
  // Drop the "array" prefix from the type name, we expect either an empty
  // string or just the shape.
  tyData = tyData.drop_front(StringRef("array").size());
  // This is the generic array case without shape, early return it.
  if (tyData.empty())
    return ToyArrayType::get(getContext());

  // Use a regex to parse the shape (for efficient we should store this regex in
  // the dialect itself).
  SmallVector<StringRef, 4> matches;
  auto shapeRegex = llvm::Regex("^<([0-9]+)(, ([0-9]+))*>$");
  if (!shapeRegex.match(tyData, &matches)) {
    getContext()->emitError(loc, "Invalid toy array shape '" + tyData + "'");
    return nullptr;
  }
  SmallVector<int64_t, 4> shape;
  // Iterate through the captures, skip the first one which is the full string.
  for (auto dimStr :
       llvm::make_range(std::next(matches.begin()), matches.end())) {
    if (dimStr.startswith(","))
      continue; // POSIX misses non-capturing groups.
    if (dimStr.empty())
      continue; // '*' makes it an optional group capture
    // Convert the capture to an integer
    unsigned long long dim;
    if (getAsUnsignedInteger(dimStr, /* Radix = */ 10, dim)) {
      getContext()->emitError(
          loc, "Couldn't parse dimension as integer, matched: " + dimStr);
      return mlir::Type();
    }
    shape.push_back(dim);
  }
  // Finally we collected all the dimensions in the shape,
  // create the array type.
  return ToyArrayType::get(getContext(), shape);
}

/// Print a Toy array type, for example `array<2, 3, 4>`
void ToyDialect::printType(mlir::Type type, raw_ostream &os) const {
  auto arrayTy = type.dyn_cast<ToyArrayType>();
  if (!arrayTy) {
    os << "unknown toy type";
    return;
  }
  os << "array";
  if (!arrayTy.getShape().empty()) {
    os << "<";
    mlir::interleaveComma(arrayTy.getShape(), os);
    os << ">";
  }
}

////////////////////////////////////////////////////////////////////////////////
//////////////////// Custom Operations for the Dialect /////////////////////////
////////////////////////////////////////////////////////////////////////////////

/// Helper to verify that the result of an operation is a Toy array type.
template <typename T> static mlir::LogicalResult verifyToyReturnArray(T *op) {
  if (!op->getResult()->getType().template isa<ToyArrayType>()) {
    std::string msg;
    raw_string_ostream os(msg);
    os << "expects a Toy Array for its argument, got "
       << op->getResult()->getType();
    return op->emitOpError(os.str());
  }
  return mlir::success();
}

/// Helper to verify that the two operands of a binary operation are Toy
/// arrays..
template <typename T> static mlir::LogicalResult verifyToyBinOperands(T *op) {
  if (!op->getOperand(0)->getType().template isa<ToyArrayType>()) {
    std::string msg;
    raw_string_ostream os(msg);
    os << "expects a Toy Array for its LHS, got "
       << op->getOperand(0)->getType();
    return op->emitOpError(os.str());
  }
  if (!op->getOperand(1)->getType().template isa<ToyArrayType>()) {
    std::string msg;
    raw_string_ostream os(msg);
    os << "expects a Toy Array for its LHS, got "
       << op->getOperand(0)->getType();
    return op->emitOpError(os.str());
  }
  return mlir::success();
}

/// Build a constant operation.
/// The builder is passed as an argument, so is the state that this method is
/// expected to fill in order to build the operation.
void ConstantOp::build(mlir::Builder *builder, mlir::OperationState *state,
                       ArrayRef<int64_t> shape, mlir::DenseElementsAttr value) {
  state->types.push_back(ToyArrayType::get(builder->getContext(), shape));
  auto dataAttribute = builder->getNamedAttr("value", value);
  state->attributes.push_back(dataAttribute);
}

/// Build a constant operation.
/// The builder is passed as an argument, so is the state that this method is
/// expected to fill in order to build the operation.
void ConstantOp::build(mlir::Builder *builder, mlir::OperationState *state,
                       mlir::FloatAttr value) {
  // Broadcast and forward to the other build factory
  mlir::Type elementType = mlir::FloatType::getF64(builder->getContext());
  auto dataType = builder->getTensorType({1}, elementType);
  auto dataAttribute = builder->getDenseElementsAttr(dataType, {value})
                           .cast<mlir::DenseElementsAttr>();

  ConstantOp::build(builder, state, {1}, dataAttribute);
}

/// Verifier for constant operation.
mlir::LogicalResult ConstantOp::verify() {
  // Ensure that the return type is a Toy array
  if (failed(verifyToyReturnArray(this)))
    return mlir::failure();

  // We expect the constant itself to be stored as an attribute.
  auto dataAttr = getAttr("value").dyn_cast<mlir::DenseElementsAttr>();
  if (!dataAttr) {
    return emitOpError(
        "missing valid `value` DenseElementsAttribute on toy.constant()");
  }
  auto attrType = dataAttr.getType().dyn_cast<mlir::TensorType>();
  if (!attrType) {
    return emitOpError(
        "missing valid `value` DenseElementsAttribute on toy.constant()");
  }

  // If the return type of the constant is not a generic array, the shape must
  // match the shape of the attribute holding the data.
  auto resultType = getResult()->getType().cast<ToyArrayType>();
  if (!resultType.isGeneric()) {
    if (attrType.getRank() != resultType.getRank()) {
      return emitOpError("The rank of the toy.constant return type must match "
                         "the one of the attached value attribute: " +
                         Twine(attrType.getRank()) +
                         " != " + Twine(resultType.getRank()));
    }
    for (int dim = 0; dim < attrType.getRank(); ++dim) {
      if (attrType.getShape()[dim] != resultType.getShape()[dim]) {
        std::string msg;
        raw_string_ostream os(msg);
        return emitOpError(
            "Shape mismatch between toy.constant return type and its "
            "attribute at dimension " +
            Twine(dim) + ": " + Twine(attrType.getShape()[dim]) +
            " != " + Twine(resultType.getShape()[dim]));
      }
    }
  }
  return mlir::success();
}

void GenericCallOp::build(mlir::Builder *builder, mlir::OperationState *state,
                          StringRef callee, ArrayRef<mlir::Value *> arguments) {
  // Generic call always returns a generic ToyArray initially
  state->types.push_back(ToyArrayType::get(builder->getContext()));
  state->operands.assign(arguments.begin(), arguments.end());
  auto calleeAttr = builder->getStringAttr(callee);
  state->attributes.push_back(builder->getNamedAttr("callee", calleeAttr));
}

mlir::LogicalResult GenericCallOp::verify() {
  // Verify that every operand is a Toy Array
  for (int opId = 0, num = getNumOperands(); opId < num; ++opId) {
    if (!getOperand(opId)->getType().template isa<ToyArrayType>()) {
      std::string msg;
      raw_string_ostream os(msg);
      os << "expects a Toy Array for its " << opId << " operand, got "
         << getOperand(opId)->getType();
      return emitOpError(os.str());
    }
  }
  return mlir::success();
}

/// Return the name of the callee.
StringRef GenericCallOp::getCalleeName() {
  return getAttr("callee").cast<mlir::StringAttr>().getValue();
}

template <typename T> static mlir::LogicalResult verifyToySingleOperand(T *op) {
  if (!op->getOperand()->getType().template isa<ToyArrayType>()) {
    std::string msg;
    raw_string_ostream os(msg);
    os << "expects a Toy Array for its argument, got "
       << op->getOperand()->getType();
    return op->emitOpError(os.str());
  }
  return mlir::success();
}

void ReturnOp::build(mlir::Builder *builder, mlir::OperationState *state,
                     mlir::Value *value) {
  // Return does not return any value and has an optional single argument
  if (value)
    state->operands.push_back(value);
}

mlir::LogicalResult ReturnOp::verify() {
  if (getNumOperands() > 1)
    return emitOpError("expects zero or one operand, got " +
                       Twine(getNumOperands()));
  if (hasOperand() && failed(verifyToySingleOperand(this)))
    return mlir::failure();
  return mlir::success();
}

void PrintOp::build(mlir::Builder *builder, mlir::OperationState *state,
                    mlir::Value *value) {
  // Print does not return any value and has a single argument
  state->operands.push_back(value);
}

mlir::LogicalResult PrintOp::verify() {
  if (failed(verifyToySingleOperand(this)))
    return mlir::failure();
  return mlir::success();
}

void TransposeOp::build(mlir::Builder *builder, mlir::OperationState *state,
                        mlir::Value *value) {
  state->types.push_back(ToyArrayType::get(builder->getContext()));
  state->operands.push_back(value);
}

mlir::LogicalResult TransposeOp::verify() {
  if (failed(verifyToySingleOperand(this)))
    return mlir::failure();
  return mlir::success();
}

void ReshapeOp::build(mlir::Builder *builder, mlir::OperationState *state,
                      mlir::Value *value, ToyArrayType reshapedType) {
  state->types.push_back(reshapedType);
  state->operands.push_back(value);
}

mlir::LogicalResult ReshapeOp::verify() {
  if (failed(verifyToySingleOperand(this)))
    return mlir::failure();
  auto retTy = getResult()->getType().dyn_cast<ToyArrayType>();
  if (!retTy)
    return emitOpError("toy.reshape is expected to produce a Toy array");
  if (retTy.isGeneric())
    return emitOpError("toy.reshape is expected to produce a shaped Toy array, "
                       "got a generic one.");
  return mlir::success();
}

void AddOp::build(mlir::Builder *builder, mlir::OperationState *state,
                  mlir::Value *lhs, mlir::Value *rhs) {
  state->types.push_back(ToyArrayType::get(builder->getContext()));
  state->operands.push_back(lhs);
  state->operands.push_back(rhs);
}

mlir::LogicalResult AddOp::verify() {
  if (failed(verifyToyBinOperands(this)))
    return mlir::failure();
  return mlir::success();
}

void MulOp::build(mlir::Builder *builder, mlir::OperationState *state,
                  mlir::Value *lhs, mlir::Value *rhs) {
  state->types.push_back(ToyArrayType::get(builder->getContext()));
  state->operands.push_back(lhs);
  state->operands.push_back(rhs);
}

mlir::LogicalResult MulOp::verify() {
  if (failed(verifyToyBinOperands(this)))
    return mlir::failure();
  return mlir::success();
}

void AllocOp::build(mlir::Builder *builder, mlir::OperationState *state,
                    mlir::Type retType) {
  state->types.push_back(retType);
}

void TypeCastOp::build(mlir::Builder *builder, mlir::OperationState *state,
                       mlir::Value *value, mlir::Type destTy) {
  state->operands.push_back(value);
  state->types.push_back(destTy);
}

} // namespace toy