// Copyright 2019 The IREE Authors
//
// Licensed 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 "iree/compiler/Dialect/VM/Target/Bytecode/BytecodeEncoder.h"

#include "iree/compiler/Dialect/Util/IR/UtilTypes.h"
#include "iree/compiler/Dialect/VM/Analysis/RegisterAllocation.h"
#include "iree/compiler/Dialect/VM/IR/VMDialect.h"
#include "iree/compiler/Dialect/VM/IR/VMTypes.h"
#include "llvm/ADT/STLExtras.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Diagnostics.h"

namespace mlir::iree_compiler::IREE::VM {

namespace {

// v0 bytecode spec. This is in extreme flux and not guaranteed to be a stable
// representation. Always generate this from source in tooling and never check
// in any emitted files!
class V0BytecodeEncoder : public BytecodeEncoder {
public:
  V0BytecodeEncoder(llvm::DenseMap<Type, int> *typeTable,
                    RegisterAllocation *registerAllocation,
                    const OrdinalAnalysis *ordinalAnalysis)
      : typeTable_(typeTable), registerAllocation_(registerAllocation),
        ordinalAnalysis_(ordinalAnalysis) {}
  ~V0BytecodeEncoder() = default;

  LogicalResult beginBlock(Block *block) override {
    blockOffsets_[block] = bytecode_.size();
    return writeUint8(static_cast<uint8_t>(Opcode::Block));
  }

  LogicalResult endBlock(Block *block) override { return success(); }

  LogicalResult beginOp(Operation *op) override {
    currentOp_ = op;
    // TODO(benvanik): encode source information (start).
    return success();
  }

  LogicalResult endOp(Operation *op) override {
    // TODO(benvanik): encode source information (end).
    currentOp_ = nullptr;
    return success();
  }

  LogicalResult encodeI8(int value) override { return writeUint8(value); }

  LogicalResult encodeOpcode(StringRef name, int opcode) override {
    return writeUint8(opcode);
  }

  LogicalResult encodeSymbolOrdinal(SymbolTable &syms,
                                    StringRef name) override {
    auto *symbolOp = syms.lookup(name);
    if (!symbolOp) {
      return currentOp_->emitOpError() << "target symbol not found: " << name;
    }
    int32_t ordinal = ordinalAnalysis_->getOrdinal(symbolOp);
    if (isa<IREE::VM::ImportOp>(symbolOp)) {
      // Imported functions have their MSB set.
      ordinal |= 0x80000000u;
    }
    return writeInt32(ordinal);
  }

  LogicalResult encodeType(Value value) override {
    auto refPtrType = dyn_cast<IREE::VM::RefType>(value.getType());
    if (!refPtrType) {
      return currentOp_->emitOpError()
             << "type " << value.getType()
             << " is not supported as a serialized type kind";
    }
    return encodeType(refPtrType.getObjectType());
  }

  LogicalResult encodeType(Type type) override {
    // HACK: it'd be nice to remove the implicit ref wrapper hiding.
    if (auto refType = dyn_cast<IREE::VM::RefType>(type)) {
      if (isa<IREE::VM::ListType>(refType.getObjectType())) {
        type = refType.getObjectType();
      }
    }
    auto it = typeTable_->find(type);
    if (it == typeTable_->end()) {
      return currentOp_->emitOpError()
             << "type " << type
             << " cannot be encoded; not registered in type table";
    }
    int typeOrdinal = it->second;
    return writeUint32(typeOrdinal);
  }

  LogicalResult encodePrimitiveAttr(TypedAttr attr) override {
    unsigned int bitWidth = attr.getType().getIntOrFloatBitWidth();
    if (auto integerAttr = dyn_cast<IntegerAttr>(attr)) {
      uint64_t limitedValue =
          integerAttr.getValue().extractBitsAsZExtValue(bitWidth, 0);
      switch (bitWidth) {
      case 8:
        return writeUint8(static_cast<uint8_t>(limitedValue));
      case 16:
        return writeUint16(static_cast<uint16_t>(limitedValue));
      case 32:
        return writeUint32(static_cast<uint32_t>(limitedValue));
      case 64:
        return writeUint64(static_cast<uint64_t>(limitedValue));
      default:
        return currentOp_->emitOpError()
               << "attribute of bitwidth " << bitWidth << " not supported";
      }
    } else if (auto floatAttr = dyn_cast<FloatAttr>(attr)) {
      switch (bitWidth) {
      case 32: {
        union {
          float f32;
          uint32_t u32;
        } value;
        value.f32 = floatAttr.getValue().convertToFloat();
        return writeUint32(value.u32);
      }
      case 64: {
        union {
          double f64;
          uint64_t u64;
        } value;
        value.f64 = floatAttr.getValue().convertToDouble();
        return writeUint64(value.u64);
      }
      default:
        return currentOp_->emitOpError()
               << "attribute of bitwidth " << bitWidth << " not supported";
      }
    } else {
      return currentOp_->emitOpError()
             << "attribute type not supported for primitive serialization: "
             << attr;
    }
  }

  LogicalResult encodePrimitiveArrayAttr(DenseElementsAttr value) override {
    if (value.getNumElements() > UINT16_MAX || failed(ensureAlignment(2)) ||
        failed(writeUint16(value.getNumElements()))) {
      return currentOp_->emitOpError() << "integer array size out of bounds";
    }
    for (auto el : value.getValues<TypedAttr>()) {
      if (failed(encodePrimitiveAttr(el))) {
        return currentOp_->emitOpError() << "failed to encode element " << el;
      }
    }
    return success();
  }

  LogicalResult encodeStrAttr(StringAttr value) override {
    if (!value) {
      return writeUint16(0);
    }
    auto stringValue = value.getValue();
    if (stringValue.size() > UINT16_MAX) {
      return currentOp_->emitOpError()
             << "string attribute too large for 16-bit p-string (needs "
             << stringValue.size() << " bytes)";
    }
    return failure(failed(writeUint16(stringValue.size())) ||
                   failed(writeBytes(stringValue.data(), stringValue.size())));
  }

  LogicalResult encodeBranch(Block *targetBlock,
                             Operation::operand_range operands,
                             int successorIndex) override {
    // Reserve space for the block offset. It will get fixed up when we are all
    // done and know all of the block offsets.
    blockOffsetFixups_.emplace_back(targetBlock, bytecode_.size());
    bytecode_.resize(bytecode_.size() + sizeof(int32_t));

    // Compute required remappings - we only need to emit them when the source
    // and dest registers differ. Hopefully the allocator did a good job and
    // this list is small :)
    //
    // 64-bit registers, which are split in 2, need to be remapped as each part.
    // This is something that could be improved in the bytecode format with
    // dedicated remapping commands or something.
    auto srcDstRegs = registerAllocation_->remapSuccessorRegisters(
        currentOp_, successorIndex);
    uint16_t registerParts = 0;
    for (auto srcDstReg : srcDstRegs) {
      registerParts +=
          srcDstReg.first.isValue() && srcDstReg.first.byteWidth() == 8 ? 2 : 1;
    }
    if (failed(ensureAlignment(2)) || failed(writeUint16(registerParts))) {
      return failure();
    }
    for (auto srcDstReg : srcDstRegs) {
      if (failed(writeUint16(srcDstReg.first.encode())) ||
          failed(writeUint16(srcDstReg.second.encode()))) {
        return failure();
      }
      if (srcDstReg.first.isValue() && srcDstReg.first.byteWidth() == 8) {
        if (failed(writeUint16(srcDstReg.first.encodeHi())) ||
            failed(writeUint16(srcDstReg.second.encodeHi()))) {
          return failure();
        }
      }
    }

    return success();
  }

  LogicalResult encodeBranchTarget(Block *targetBlock) override {
    // Reserve space for the block offset. It will get fixed up when we are all
    // done and know all of the block offsets.
    blockOffsetFixups_.emplace_back(targetBlock, bytecode_.size());
    bytecode_.resize(bytecode_.size() + sizeof(int32_t));
    return success();
  }

  LogicalResult encodeBranchTable(SuccessorRange caseSuccessors,
                                  OperandRangeRange caseOperands,
                                  int baseSuccessorIndex) override {
    if (failed(writeUint16(caseSuccessors.size()))) {
      return failure();
    }
    for (auto [successor, operands] :
         llvm::zip_equal(caseSuccessors, caseOperands)) {
      if (failed(encodeBranch(successor, operands, ++baseSuccessorIndex))) {
        return failure();
      }
    }
    return success();
  }

  LogicalResult encodeOperand(Value value, int ordinal) override {
    uint16_t reg =
        registerAllocation_->mapUseToRegister(value, currentOp_, ordinal)
            .encode();
    return writeUint16(reg);
  }

  LogicalResult encodeOperands(Operation::operand_range values) override {
    if (failed(ensureAlignment(2)) ||
        failed(writeUint16(std::distance(values.begin(), values.end())))) {
      return failure();
    }
    for (auto it : llvm::enumerate(values)) {
      uint16_t reg = registerAllocation_
                         ->mapUseToRegister(it.value(), currentOp_, it.index())
                         .encode();
      if (failed(writeUint16(reg))) {
        return failure();
      }
    }
    return success();
  }

  LogicalResult
  encodeOperands(ArrayRef<std::pair<Value, int>> valuesWithIndices) override {
    if (failed(ensureAlignment(2)) ||
        failed(writeUint16(valuesWithIndices.size()))) {
      return failure();
    }
    for (auto [value, operandIndex] : valuesWithIndices) {
      uint16_t reg =
          registerAllocation_->mapUseToRegister(value, currentOp_, operandIndex)
              .encode();
      if (failed(writeUint16(reg))) {
        return failure();
      }
    }
    return success();
  }

  LogicalResult encodeResult(Value value) override {
    uint16_t reg = registerAllocation_->mapToRegister(value).encode();
    return writeUint16(reg);
  }

  LogicalResult encodeResults(Operation::result_range values) override {
    if (failed(ensureAlignment(2)) ||
        failed(writeUint16(std::distance(values.begin(), values.end())))) {
      return failure();
    }
    for (auto value : values) {
      uint16_t reg = registerAllocation_->mapToRegister(value).encode();
      if (failed(writeUint16(reg))) {
        return failure();
      }
    }
    return success();
  }

  LogicalResult encodeBlockArgResults(Block *targetBlock) override {
    auto blockArgs = targetBlock->getArguments();
    if (failed(ensureAlignment(2)) || failed(writeUint16(blockArgs.size()))) {
      return failure();
    }
    for (auto arg : blockArgs) {
      uint16_t reg = registerAllocation_->mapToRegister(arg).encode();
      if (failed(writeUint16(reg))) {
        return failure();
      }
    }
    return success();
  }

  std::optional<std::vector<uint8_t>> finish() {
    if (failed(fixupOffsets())) {
      return std::nullopt;
    }
    return std::move(bytecode_);
  }

  size_t getOffset() const { return bytecode_.size(); }

  const RegisterAllocation *getRegisterAllocation() const override {
    return registerAllocation_;
  }

  LogicalResult ensureAlignment(size_t alignment) {
    size_t paddedSize = (bytecode_.size() + (alignment - 1)) & ~(alignment - 1);
    size_t padding = paddedSize - bytecode_.size();
    if (padding == 0) {
      return success();
    }
    static const uint8_t kZeros[32] = {0};
    if (padding > sizeof(kZeros)) {
      return failure();
    }
    return writeBytes(kZeros, padding);
  }

private:
  // TODO(benvanik): replace this with something not using an ever-expanding
  // vector. I'm sure LLVM has something.

  MutableArrayRef<uint8_t> reserveBytes(size_t dataLength) {
    int offset = bytecode_.size();
    bytecode_.resize(offset + dataLength);
    return MutableArrayRef<uint8_t>(
        reinterpret_cast<uint8_t *>(bytecode_.data()) + offset, dataLength);
  }

  LogicalResult writeBytes(const void *data, size_t dataLength) {
    auto dst = reserveBytes(dataLength);
    if (dataLength != dst.size()) {
      return failure();
    }
    std::memcpy(dst.data(), data, dst.size());
    return success();
  }

  LogicalResult writeUint8(uint8_t value) {
    return writeBytes(&value, sizeof(value));
  }

  LogicalResult writeUint16(uint16_t value) {
    return writeBytes(&value, sizeof(value));
  }

  LogicalResult writeInt32(int32_t value) {
    return writeBytes(&value, sizeof(value));
  }

  LogicalResult writeUint32(uint32_t value) {
    return writeBytes(&value, sizeof(value));
  }

  LogicalResult writeUint64(uint64_t value) {
    return writeBytes(&value, sizeof(value));
  }

  LogicalResult fixupOffsets() {
    for (const auto &fixup : blockOffsetFixups_) {
      auto blockOffset = blockOffsets_.find(fixup.first);
      if (blockOffset == blockOffsets_.end()) {
        emitError(fixup.first->front().getLoc()) << "Block offset not found";
        return failure();
      }
      std::memcpy(bytecode_.data() + fixup.second, &blockOffset->second,
                  sizeof(int32_t));
    }
    blockOffsetFixups_.clear();
    return success();
  }

  llvm::DenseMap<Type, int> *typeTable_;
  RegisterAllocation *registerAllocation_;
  const OrdinalAnalysis *ordinalAnalysis_;

  Operation *currentOp_ = nullptr;

  std::vector<uint8_t> bytecode_;
  llvm::DenseMap<Block *, size_t> blockOffsets_;
  std::vector<std::pair<Block *, size_t>> blockOffsetFixups_;
};

} // namespace

// static
std::optional<EncodedBytecodeFunction> BytecodeEncoder::encodeFunction(
    IREE::VM::FuncOp funcOp, llvm::DenseMap<Type, int> &typeTable,
    SymbolTable &symbolTable, const OrdinalAnalysis &ordinalAnalysis,
    DebugDatabaseBuilder &debugDatabase) {
  EncodedBytecodeFunction result;

  // Perform register allocation first so that we can quickly lookup values as
  // we encode the operands/results.
  RegisterAllocation registerAllocation;
  if (failed(registerAllocation.recalculate(funcOp))) {
    funcOp.emitError() << "register allocation failed";
    return std::nullopt;
  }

  FunctionSourceMap sourceMap;
  sourceMap.localName = funcOp.getName().str();

  V0BytecodeEncoder encoder(&typeTable, &registerAllocation, &ordinalAnalysis);
  for (auto &block : funcOp.getBlocks()) {
    size_t blockStart = encoder.getOffset();

    if (failed(encoder.beginBlock(&block))) {
      funcOp.emitError() << "failed to begin block";
      return std::nullopt;
    }

    for (auto &op : block.getOperations()) {
      auto serializableOp = dyn_cast<IREE::VM::VMSerializableOp>(op);
      if (!serializableOp) {
        op.emitOpError() << "is not serializable";
        return std::nullopt;
      }
      const int32_t bytecodeOffset = static_cast<int32_t>(encoder.getOffset());
      if (failed(encoder.beginOp(&op))) {
        op.emitOpError() << "failed to encode";
        return std::nullopt;
      }
      auto encodeResult = serializableOp.encode(symbolTable, encoder);
      if (failed(encodeResult)) {
        op.emitOpError() << "failed to encode";
        return std::nullopt;
      }
      if (failed(encoder.endOp(&op))) {
        op.emitOpError() << "failed to encode";
        return std::nullopt;
      }
      if (*encodeResult) {
        sourceMap.locations.push_back({bytecodeOffset, op.getLoc()});
      }
    }

    if (failed(encoder.endBlock(&block))) {
      funcOp.emitError() << "failed to end block";
      return std::nullopt;
    }

    // From isa.h: IREE_VM_ISA_PC_BLOCK_MAX
    static const size_t kVMMaxBlockSize = 0x00FFFFFFu;
    size_t blockLength = encoder.getOffset() - blockStart;
    if (blockLength > kVMMaxBlockSize) {
      funcOp.emitError() << "encoded block too large; VM currently restricts "
                            "bytecode function block bodies to "
                         << kVMMaxBlockSize << "B but this function encodes to "
                         << blockLength << "B";
      return std::nullopt;
    }
  }

  debugDatabase.addFunctionSourceMap(funcOp, sourceMap);

  size_t finalLength = encoder.getOffset();
  if (failed(encoder.ensureAlignment(8))) {
    funcOp.emitError() << "failed to pad function";
    return std::nullopt;
  }
  auto bytecodeData = encoder.finish();
  if (!bytecodeData.has_value()) {
    funcOp.emitError() << "failed to fixup and finish encoding";
    return std::nullopt;
  }
  result.bytecodeData = bytecodeData.value();
  result.bytecodeLength = finalLength;
  result.blockCount = funcOp.getBlocks().size();
  result.i32RegisterCount = registerAllocation.getMaxI32RegisterOrdinal() + 1;
  result.refRegisterCount = registerAllocation.getMaxRefRegisterOrdinal() + 1;
  return result;
}

} // namespace mlir::iree_compiler::IREE::VM
