// 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 "compiler/Serialization/VMFunctionBuilder.h"

#include "flatbuffers/flatbuffers.h"
#include "compiler/IR/Dialect.h"
#include "compiler/IR/Types.h"
#include "compiler/Serialization/BytecodeTables.h"
#include "compiler/Utils/Macros.h"
#include "schemas/bytecode/bytecode_v0.h"
#include "schemas/type_def_generated.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Location.h"
#include "mlir/IR/Module.h"

namespace mlir {
namespace iree_compiler {

namespace {

LogicalResult WriteGenericIreeOp(Block *block, Operation *op,
                                 BytecodeWriter *writer) {
  // Strip the dialect name from the op name and lookup the opcode.
  // TODO(benvanik): adjust for supporting sequencer opcodes.

  auto opName = op->getName().getStringRef();
  auto dialect = op->getDialect();
  if (!dialect) {
    return op->emitOpError() << "Op does not belong to a registered dialect";
  }

  auto dialectNamespace = dialect->getNamespace();
  std::unique_ptr<OpcodeInfo> operandInfo;
  auto strippedOpName = opName.substr(opName.find('.') + 1).str();
  if (dialectNamespace == "iree_ll_seq") {
    auto opcode = GetSequencerOpcodeByName(strippedOpName);
    if (!opcode.hasValue()) {
      return op->emitOpError()
             << "No sequencer opcode found for op; is it a pseudo op?";
    }
    RETURN_IF_FAILURE(writer->WriteOpcode(opcode.getValue()));
    operandInfo =
        std::make_unique<OpcodeInfo>(GetSequencerOpcodeInfo(opcode.getValue()));
  } else if (dialectNamespace == "iree_ll_interp" ||
             // TODO(gcmn) remove special case for IREE dialect?
             dialectNamespace == IREEDialect::getDialectNamespace()) {
    auto opcode = GetInterpreterOpcodeByName(strippedOpName);
    if (!opcode.hasValue()) {
      return op->emitOpError()
             << "No interpreter opcode found for op; is it a pseudo op?";
    }
    RETURN_IF_FAILURE(writer->WriteOpcode(opcode.getValue()));
    operandInfo = std::make_unique<OpcodeInfo>(
        GetInterpreterOpcodeInfo(opcode.getValue()));
  } else {
    return op->emitOpError()
           << "Op belongs to unknown dialect " << dialectNamespace.str();
  }
  // Write inputs and outputs based on the bytecode encoding.
  int operandIndex = 0;
  int resultIndex = 0;
  for (int i = 0; i < llvm::array_lengthof(operandInfo->operands); ++i) {
    auto op_encoding = operandInfo->operands[i];
    if (op_encoding == iree::OperandEncoding::kNone) break;
    switch (op_encoding) {
      case iree::OperandEncoding::kInputSlot:
      case iree::OperandEncoding::kOutputSlot: {
        auto *value = op->getOperand(operandIndex++);
        RETURN_IF_FAILURE(writer->WriteLocal(value));
        break;
      }
      case iree::OperandEncoding::kVariadicInputSlots:
      case iree::OperandEncoding::kVariadicOutputSlots: {
        int count = op->getNumOperands() - operandIndex;
        RETURN_IF_FAILURE(writer->WriteCount(count));
        for (; count; --count) {
          auto *value = op->getOperand(operandIndex++);
          RETURN_IF_FAILURE(writer->WriteLocal(value));
        }
        break;
      }
      case iree::OperandEncoding::kResultSlot: {
        auto *value = op->getResult(resultIndex++);
        RETURN_IF_FAILURE(writer->WriteLocal(value));
        break;
      }
      case iree::OperandEncoding::kVariadicResultSlots: {
        int count = op->getNumResults() - resultIndex;
        RETURN_IF_FAILURE(writer->WriteCount(count));
        for (; count; --count) {
          auto *value = op->getResult(resultIndex++);
          RETURN_IF_FAILURE(writer->WriteLocal(value));
        }
        break;
      }
      case iree::OperandEncoding::kConstant:
      case iree::OperandEncoding::kFunctionOrdinal:
      case iree::OperandEncoding::kBlockOffset:
      case iree::OperandEncoding::kTypeIndex:
      case iree::OperandEncoding::kIndex:
      case iree::OperandEncoding::kIndexList:
      case iree::OperandEncoding::kCmpIPredicate:
      case iree::OperandEncoding::kCmpFPredicate:
        return op->emitOpError()
               << "Operand encoding " << static_cast<char>(op_encoding)
               << " not supported by generic writer for " << opName.str();
        return failure();
      default:
        return op->emitOpError()
               << "Operand encoding " << static_cast<char>(op_encoding) << " ("
               << static_cast<int>(op_encoding) << ") not recognized (typo?)";
    }
  }

  return success();
}

}  // namespace

VMFunctionBuilder::VMFunctionBuilder(FuncOp function,
                                     VMFunctionTableBuilder *functionTable,
                                     ::flatbuffers::FlatBufferBuilder *fbb)
    : context_(function.getContext()),
      function_(function),
      functionTable_(functionTable),
      fbb_(fbb) {}

void VMFunctionBuilder::RegisterCustomWriter(StringRef operationName,
                                             CustomWriterFn writerFn) {
  customWriters_.insert({operationName, writerFn});
}

LogicalResult VMFunctionBuilder::ConvertBytecode() {
  BytecodeWriter writer;
  sourceMap_ = {};

  RETURN_IF_FAILURE(BeginFunction(function_, &writer));
  for (auto &block : function_.getBlocks()) {
    RETURN_IF_FAILURE(BeginBlock(&block, &writer));
    for (auto &op : block.getOperations()) {
      if (failed(WriteOperation(&block, &op, &writer))) {
        op.emitError() << "Unable to serialize operation";
        return failure();
      }
    }
    RETURN_IF_FAILURE(EndBlock(&block, block.getTerminator(), &writer));
  }
  RETURN_IF_FAILURE(EndFunction(function_, &writer));

  int localCount = writer.local_count();
  auto bodyBytes = writer.Finish();
  auto bodyOffset = fbb_->CreateVector(
      reinterpret_cast<const int8_t *>(bodyBytes.data()), bodyBytes.size());
  iree::BytecodeDefBuilder bdb(*fbb_);
  bdb.add_local_count(localCount);
  bdb.add_contents(bodyOffset);
  bytecodeDef_ = bdb.Finish();

  return success();
}

::flatbuffers::Offset<iree::FunctionDef> VMFunctionBuilder::Finish() {
  using TypeDefVector =
      ::flatbuffers::Vector<::flatbuffers::Offset<iree::TypeDef>>;

  const auto &functionType = function_.getType();
  std::vector<::flatbuffers::Offset<iree::TypeDef>> inputs;
  for (const auto &type : functionType.getInputs()) {
    auto typeOffset = SerializeType(type, fbb_);
    if (typeOffset.IsNull()) return {};
    inputs.push_back(typeOffset);
  }
  ::flatbuffers::Offset<TypeDefVector> inputsOffset;
  if (!inputs.empty()) {
    inputsOffset = fbb_->CreateVector(inputs);
  }

  std::vector<::flatbuffers::Offset<iree::TypeDef>> results;
  for (const auto &type : functionType.getResults()) {
    auto typeOffset = SerializeType(type, fbb_);
    if (typeOffset.IsNull()) return {};
    results.push_back(typeOffset);
  }
  ::flatbuffers::Offset<TypeDefVector> resultsOffset;
  if (!results.empty()) {
    resultsOffset = fbb_->CreateVector(results);
  }
  iree::FunctionTypeDefBuilder ftb(*fbb_);
  ftb.add_inputs(inputsOffset);
  ftb.add_results(resultsOffset);
  auto functionTypeOffset = ftb.Finish();

  // TODO(benvanik): strip names of internal functions.
  auto nameOffset = fbb_->CreateString(function_.getName().str());
  iree::FunctionDefBuilder fdb(*fbb_);
  fdb.add_name(nameOffset);
  fdb.add_type(functionTypeOffset);
  fdb.add_bytecode(bytecodeDef_);
  return fdb.Finish();
}

LogicalResult VMFunctionBuilder::BeginFunction(FuncOp function,
                                               BytecodeWriter *writer) {
  // Assign value slots for all arguments and results.
  // Keeping them at the front will make it easier to find during debugging
  // and makes spans easier to compute at runtime.
  for (auto argument : function.getArguments()) {
    RETURN_IF_FAILURE(writer->PrepareLocal(argument));
  }
  return success();
}

LogicalResult VMFunctionBuilder::EndFunction(FuncOp function,
                                             BytecodeWriter *writer) {
  RETURN_IF_FAILURE(writer->FixupOffsets());
  return success();
}

LogicalResult VMFunctionBuilder::BeginBlock(Block *block,
                                            BytecodeWriter *writer) {
  RETURN_IF_FAILURE(writer->MarkBlockOffset(block));
  return success();
}

LogicalResult VMFunctionBuilder::EndBlock(Block *block, Operation *op,
                                          BytecodeWriter *writer) {
  return success();
}

LogicalResult VMFunctionBuilder::WriteOperation(Block *block, Operation *baseOp,
                                                BytecodeWriter *writer) {
  if (!baseOp->getLoc().isa<UnknownLoc>()) {
    sourceMap_.locations.push_back({writer->offset(), baseOp->getLoc()});
  }

  // Check registered writers first to allow overrides.
  auto writerIt = customWriters_.find(baseOp->getName().getStringRef());
  if (writerIt != customWriters_.end()) {
    return writerIt->second(baseOp, writer);
  }

  // Fallback to using the generic writer.
  if (baseOp->getAbstractOperation()->dialect.getNamespace().startswith(
          "iree")) {
    RETURN_IF_FAILURE(WriteGenericIreeOp(block, baseOp, writer));
  } else {
    return baseOp->emitError()
           << "Unsupported op " << baseOp->getName().getStringRef().str()
           << "; incorrectly outlined or not yet implemented";
  }
  return success();
}

::flatbuffers::Offset<iree::TypeDef> VMFunctionBuilder::SerializeType(
    Type type, ::flatbuffers::FlatBufferBuilder *fbb) {
  ::flatbuffers::Offset<void> typeDefUnion;
  iree::TypeDefUnion typeUnionType;
  if (auto memRefType = type.dyn_cast<MemRefType>()) {
    auto memRefTypeOffset = SerializeMemRefType(memRefType, fbb_);
    if (memRefTypeOffset.IsNull()) return {};
    typeDefUnion = memRefTypeOffset.Union();
    typeUnionType = iree::TypeDefUnion::MemRefTypeDef;
  } else if (auto deviceType = type.dyn_cast<DeviceType>()) {
    typeDefUnion = iree::CreateDeviceTypeDef(*fbb).Union();
    typeUnionType = iree::TypeDefUnion::DeviceTypeDef;
  } else if (auto commandBufferType = type.dyn_cast<CommandBufferType>()) {
    typeDefUnion = iree::CreateCommandBufferTypeDef(*fbb).Union();
    typeUnionType = iree::TypeDefUnion::CommandBufferTypeDef;
  } else if (auto eventType = type.dyn_cast<EventType>()) {
    typeDefUnion = iree::CreateEventTypeDef(*fbb).Union();
    typeUnionType = iree::TypeDefUnion::EventTypeDef;
  } else if (auto semaphoreType = type.dyn_cast<SemaphoreType>()) {
    typeDefUnion = iree::CreateSemaphoreTypeDef(*fbb).Union();
    typeUnionType = iree::TypeDefUnion::SemaphoreTypeDef;
  } else if (auto fenceType = type.dyn_cast<FenceType>()) {
    typeDefUnion = iree::CreateFenceTypeDef(*fbb).Union();
    typeUnionType = iree::TypeDefUnion::FenceTypeDef;
  } else {
    function_.emitError() << "Function " << function_.getName().str()
                          << " has unsupported I/O with type " << type;
    return {};
  }

  iree::TypeDefBuilder tdb(*fbb);
  tdb.add_type_union_type(typeUnionType);
  tdb.add_type_union(typeDefUnion);
  return tdb.Finish();
}

::flatbuffers::Offset<iree::MemRefTypeDef>
VMFunctionBuilder::SerializeMemRefType(const MemRefType &type,
                                       ::flatbuffers::FlatBufferBuilder *fbb) {
  auto elementTypeOffset = SerializeElementType(type.getElementType(), fbb);
  if (elementTypeOffset.IsNull()) return {};
  std::vector<int> shape;
  for (int dim : type.getShape()) {
    shape.push_back(dim);
  }
  auto shapeOffset = fbb->CreateVector(shape);
  iree::MemRefTypeDefBuilder tb(*fbb);
  tb.add_element_type(elementTypeOffset);
  tb.add_shape(shapeOffset);
  tb.add_memory_space(type.getMemorySpace());
  return tb.Finish();
}

::flatbuffers::Offset<iree::ElementTypeDef>
VMFunctionBuilder::SerializeElementType(const Type &genericType,
                                        ::flatbuffers::FlatBufferBuilder *fbb) {
  ::flatbuffers::Offset<void> typeDefUnion;
  iree::ElementTypeDefUnion typeUnionType;
  if (auto type = genericType.dyn_cast<FloatType>()) {
    iree::FloatTypeDefBuilder tb(*fbb);
    tb.add_width(type.getWidth());
    typeDefUnion = tb.Finish().Union();
    typeUnionType = iree::ElementTypeDefUnion::FloatTypeDef;
  } else if (auto type = genericType.dyn_cast<IntegerType>()) {
    iree::IntegerTypeDefBuilder tb(*fbb);
    tb.add_width(type.getWidth());
    typeDefUnion = tb.Finish().Union();
    typeUnionType = iree::ElementTypeDefUnion::IntegerTypeDef;
  } else if (auto type = genericType.dyn_cast<OpaqueType>()) {
    auto dialectOffset = fbb->CreateString(type.getDialectNamespace().c_str());
    auto typeDataOffset = fbb->CreateString(type.getTypeData().data());
    iree::UnknownTypeDefBuilder tb(*fbb);
    tb.add_dialect(dialectOffset);
    tb.add_type_data(typeDataOffset);
    typeDefUnion = tb.Finish().Union();
    typeUnionType = iree::ElementTypeDefUnion::UnknownTypeDef;
  } else {
    function_.emitError()
        << "Unimplemented type encoding: " << genericType
        << "; ensure IREE lowering passes are converting types to the IREE "
           "set";
    return {};
  }

  iree::ElementTypeDefBuilder tdb(*fbb);
  tdb.add_type_union_type(typeUnionType);
  tdb.add_type_union(typeDefUnion);
  return tdb.Finish();
}

}  // namespace iree_compiler
}  // namespace mlir