// Copyright 2019 Google LLC
//
// 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
//
//      https://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.

#include "third_party/mlir_edge/iree/compiler/IR/StructureOps.h"

#include "third_party/llvm/llvm/include/llvm/ADT/STLExtras.h"
#include "third_party/llvm/llvm/include/llvm/ADT/SmallString.h"
#include "third_party/llvm/llvm/include/llvm/ADT/SmallVector.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/Attributes.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/Builders.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/Diagnostics.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/OpImplementation.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/SymbolTable.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/Value.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/Support/LLVM.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/Support/STLExtras.h"
#include "third_party/mlir_edge/iree/compiler/IR/Types.h"

namespace mlir {
namespace iree_compiler {
namespace IREE {

//===----------------------------------------------------------------------===//
// Generic printers and parsers.
//===----------------------------------------------------------------------===//

// Parses an op that has no inputs and no outputs.
static ParseResult parseNoIOOp(OpAsmParser *parser, OperationState *state) {
  if (failed(parser->parseOptionalAttributeDict(state->attributes))) {
    return failure();
  }
  return success();
}

// Prints an op that has no inputs and no outputs.
static void printNoIOOp(Operation *op, OpAsmPrinter *printer) {
  *printer << op->getName();
  printer->printOptionalAttrDict(op->getAttrs());
}

//===----------------------------------------------------------------------===//
// iree.module
//===----------------------------------------------------------------------===//

void ModuleOp::build(Builder *builder, OperationState *state) {
  ensureTerminator(*state->addRegion(), *builder, state->location);
}

static ParseResult parseModuleOp(OpAsmParser *parser, OperationState *state) {
  Region *body = state->addRegion();
  if (parser->parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{})) {
    return failure();
  }
  if (parser->parseOptionalAttributeDict(state->attributes)) {
    return failure();
  }
  ModuleOp::ensureTerminator(*body, parser->getBuilder(), state->location);
  return success();
}

static void printModuleOp(OpAsmPrinter *printer, Operation *op) {
  *printer << op->getName();
  printer->printRegion(op->getRegion(0), /*printEntryBlockArgs=*/false,
                       /*printBlockTerminators=*/false);
  printer->printOptionalAttrDict(op->getAttrs());
}

//===----------------------------------------------------------------------===//
// iree.multi_arch_executable
//===----------------------------------------------------------------------===//

void MultiArchExecutableOp::build(Builder *builder, OperationState *state,
                                  StringRef name) {
  state->addAttribute(SymbolTable::getSymbolAttrName(),
                      builder->getStringAttr(name));
  ensureTerminator(*state->addRegion(), *builder, state->location);
}

static ParseResult parseMultiArchExecutableOp(OpAsmParser *parser,
                                              OperationState *state) {
  auto &builder = parser->getBuilder();

  // Parse the name as a symbol reference attr and then convert to a string.
  SymbolRefAttr nameAttr;
  if (failed(parser->parseAttribute(nameAttr, SymbolTable::getSymbolAttrName(),
                                    state->attributes))) {
    return failure();
  }
  state->attributes.back().second = builder.getStringAttr(nameAttr.getValue());

  if (succeeded(parser->parseOptionalLSquare())) {
    IntegerAttr ordinalAttr;
    if (failed(parser->parseAttribute(ordinalAttr, builder.getIntegerType(32),
                                      "iree.ordinal", state->attributes)) ||
        failed(parser->parseRSquare())) {
      return failure();
    }
  }

  if (failed(parser->parseLParen()) || failed(parser->parseRParen())) {
    return failure();
  }

  Region *body = state->addRegion();
  if (failed(parser->parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{}))) {
    return failure();
  }
  if (succeeded(parser->parseOptionalKeyword("attributes"))) {
    if (failed(parser->parseOptionalAttributeDict(state->attributes))) {
      return failure();
    }
  }

  MultiArchExecutableOp::ensureTerminator(*body, builder, state->location);

  return success();
}

static void printMultiArchExecutableOp(OpAsmPrinter *printer,
                                       MultiArchExecutableOp op) {
  *printer << op.getOperationName() << " @" << op.sym_name();
  if (auto ordinalAttr =
          op.getAttr("iree.ordinal").dyn_cast_or_null<IntegerAttr>()) {
    *printer << "[" << ordinalAttr.getInt() << "]";
  }
  *printer << "()";

  printer->printRegion(op.body(), /*printEntryBlockArgs=*/false,
                       /*printBlockTerminators=*/false);

  // Print out executable attributes, if present.
  SmallVector<StringRef, 2> ignoredAttrs = {
      SymbolTable::getSymbolAttrName(),
      "iree.ordinal",
  };
  SmallVector<NamedAttribute, 4> attrs(
      llvm::make_filter_range(op.getAttrs(), [&](const NamedAttribute &attr) {
        return llvm::count(ignoredAttrs, attr.first) == 0;
      }));
  if (!attrs.empty()) {
    *printer << "\n    attributes ";
    printer->printOptionalAttrDict(attrs);
  }
}

//===----------------------------------------------------------------------===//
// iree.executable
//===----------------------------------------------------------------------===//

void ExecutableOp::build(Builder *builder, OperationState *state,
                         IREE::ExecutableFormat format) {
  state->addAttribute(
      "format", builder->getI32IntegerAttr(static_cast<uint32_t>(format)));
  ensureTerminator(*state->addRegion(), *builder, state->location);
}

static ParseResult parseExecutableOp(OpAsmParser *parser,
                                     OperationState *state) {
  auto &builder = parser->getBuilder();

  if (succeeded(parser->parseOptionalLSquare())) {
    IntegerAttr ordinalAttr;
    if (failed(parser->parseAttribute(ordinalAttr, builder.getIntegerType(32),
                                      "iree.ordinal", state->attributes)) ||
        failed(parser->parseRSquare())) {
      return failure();
    }
  }

  IntegerAttr executableOrdinalAttr;
  StringAttr formatAttr;
  llvm::SMLoc formatLoc;
  if (failed(parser->parseLParen()) ||
      failed(parser->getCurrentLocation(&formatLoc)) ||
      failed(parser->parseAttribute(formatAttr, "format", state->attributes))) {
    return failure();
  }
  auto format = symbolizeExecutableFormat(formatAttr.getValue());
  if (!format.hasValue()) {
    return parser->emitError(formatLoc)
           << "Unknown executable format " << formatAttr.getValue();
  }
  state->attributes.back().second =
      builder.getI32IntegerAttr(static_cast<int32_t>(format.getValue()));

  Region *body = state->addRegion();
  if (failed(parser->parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{}))) {
    return failure();
  }
  if (succeeded(parser->parseOptionalKeyword("attributes"))) {
    if (failed(parser->parseOptionalAttributeDict(state->attributes))) {
      return failure();
    }
  }

  ExecutableOp::ensureTerminator(*body, parser->getBuilder(), state->location);

  return success();
}

static void printExecutableOp(OpAsmPrinter *printer, ExecutableOp op) {
  *printer << op.getOperationName();
  if (auto ordinalAttr =
          op.getAttr("iree.ordinal").dyn_cast_or_null<IntegerAttr>()) {
    *printer << "[" << ordinalAttr.getInt() << "]";
  }
  *printer << "(";
  auto format = symbolizeExecutableFormat(op.format());
  if (format.hasValue()) {
    *printer << stringifyExecutableFormat(format.getValue());
  } else {
    *printer << "INVALID FORMAT";
  }
  *printer << ")";

  printer->printRegion(op.body(), /*printEntryBlockArgs=*/false,
                       /*printBlockTerminators=*/false);

  // Print out executable attributes, if present.
  SmallVector<StringRef, 2> ignoredAttrs = {
      "iree.ordinal",
      "format",
  };
  SmallVector<NamedAttribute, 4> attrs(
      llvm::make_filter_range(op.getAttrs(), [&](const NamedAttribute &attr) {
        return llvm::count(ignoredAttrs, attr.first) == 0;
      }));
  if (!attrs.empty()) {
    *printer << "\n      attributes ";
    printer->printOptionalAttrDict(attrs);
  }
}

#define GET_OP_CLASSES
#include "third_party/mlir_edge/iree/compiler/IR/StructureOps.cpp.inc"

}  // namespace IREE
}  // namespace iree_compiler
}  // namespace mlir