// 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/Util/IR/UtilOps.h"
#include "iree/compiler/Dialect/Util/IR/UtilTypes.h"
#include "iree/compiler/Dialect/VM/Conversion/TypeConverter.h"
#include "iree/compiler/Dialect/VM/IR/VMOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/Matchers.h"
#include "mlir/Transforms/DialectConversion.h"

namespace mlir::iree_compiler {

namespace {

static Value castToI32(Value value, OpBuilder &builder) {
  if (value.getType().isInteger(32)) {
    return value;
  }
  return builder.createOrFold<IREE::VM::TruncI64I32Op>(
      value.getLoc(), builder.getI32Type(), value);
}

static Value castToIndex(Value value, OpBuilder &builder) {
  if (value.getType().isIndex()) {
    return value;
  }
  return builder.createOrFold<arith::IndexCastOp>(
      value.getLoc(), builder.getIndexType(), value);
}

class ListCreateOpConversion
    : public OpConversionPattern<IREE::Util::ListCreateOp> {
  using Base::Base;
  LogicalResult
  matchAndRewrite(IREE::Util::ListCreateOp srcOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Value initialCapacity = adaptor.getInitialCapacity();
    if (initialCapacity) {
      initialCapacity = castToI32(initialCapacity, rewriter);
    } else {
      initialCapacity = IREE::VM::ConstI32Op::create(
          rewriter, srcOp.getLoc(), rewriter.getI32IntegerAttr(0));
    }
    rewriter.replaceOpWithNewOp<IREE::VM::ListAllocOp>(
        srcOp, typeConverter->convertType(srcOp.getResult().getType()),
        initialCapacity);
    return success();
  }
};

class ListConstructOpConversion
    : public OpConversionPattern<IREE::Util::ListConstructOp> {
  using Base::Base;
  LogicalResult
  matchAndRewrite(IREE::Util::ListConstructOp srcOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // Allocate list with exact capacity.
    size_t valueCount = srcOp.getValues().size();
    Value size = IREE::VM::ConstI32Op::create(
        rewriter, srcOp.getLoc(), rewriter.getI32IntegerAttr(valueCount));
    Value list = IREE::VM::ListAllocOp::create(
        rewriter, srcOp.getLoc(),
        typeConverter->convertType(srcOp.getResult().getType()), size);

    // Resize to the count (if needed).
    if (valueCount > 0) {
      IREE::VM::ListResizeOp::create(rewriter, srcOp.getLoc(), list, size);
    }

    // Add all entries.
    for (auto [i, value] : llvm::enumerate(adaptor.getValues())) {
      Value index = IREE::VM::ConstI32Op::create(
          rewriter, srcOp.getLoc(),
          rewriter.getI32IntegerAttr(static_cast<int>(i)));
      if (value.getType().isInteger(32)) {
        IREE::VM::ListSetI32Op::create(rewriter, srcOp.getLoc(), list, index,
                                       value);
      } else if (value.getType().isInteger(64)) {
        IREE::VM::ListSetI64Op::create(rewriter, srcOp.getLoc(), list, index,
                                       value);
      } else if (value.getType().isFloat(32)) {
        IREE::VM::ListSetF32Op::create(rewriter, srcOp.getLoc(), list, index,
                                       value);
      } else if (value.getType().isFloat(64)) {
        IREE::VM::ListSetF64Op::create(rewriter, srcOp.getLoc(), list, index,
                                       value);
      } else if (isa<IREE::VM::RefType>(value.getType())) {
        IREE::VM::ListSetRefOp::create(rewriter, srcOp.getLoc(), list, index,
                                       value);
      } else {
        return rewriter.notifyMatchFailure(srcOp, "invalid list element type");
      }
    }

    rewriter.replaceOp(srcOp, list);
    return success();
  }
};

class ListSizeOpConversion
    : public OpConversionPattern<IREE::Util::ListSizeOp> {
  using Base::Base;
  LogicalResult
  matchAndRewrite(IREE::Util::ListSizeOp srcOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Value size = IREE::VM::ListSizeOp::create(
        rewriter, srcOp.getLoc(), rewriter.getI32Type(), adaptor.getList());
    rewriter.replaceOp(srcOp, castToIndex(size, rewriter));
    return success();
  }
};

class ListResizeOpConversion
    : public OpConversionPattern<IREE::Util::ListResizeOp> {
  using Base::Base;
  LogicalResult
  matchAndRewrite(IREE::Util::ListResizeOp srcOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    rewriter.replaceOpWithNewOp<IREE::VM::ListResizeOp>(
        srcOp, adaptor.getList(), castToI32(adaptor.getNewSize(), rewriter));
    return success();
  }
};

class ListGetOpConversion : public OpConversionPattern<IREE::Util::ListGetOp> {
  using Base::Base;
  LogicalResult
  matchAndRewrite(IREE::Util::ListGetOp srcOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto index = castToI32(adaptor.getIndex(), rewriter);
    auto resultType = typeConverter->convertType(srcOp.getResult().getType());
    if (resultType.isInteger(32)) {
      rewriter.replaceOpWithNewOp<IREE::VM::ListGetI32Op>(
          srcOp, resultType, adaptor.getList(), index);
    } else if (resultType.isInteger(64)) {
      rewriter.replaceOpWithNewOp<IREE::VM::ListGetI64Op>(
          srcOp, resultType, adaptor.getList(), index);
    } else if (resultType.isF32()) {
      rewriter.replaceOpWithNewOp<IREE::VM::ListGetF32Op>(
          srcOp, resultType, adaptor.getList(), index);
    } else if (resultType.isF64()) {
      rewriter.replaceOpWithNewOp<IREE::VM::ListGetF64Op>(
          srcOp, resultType, adaptor.getList(), index);
    } else if (!resultType.isIntOrIndexOrFloat()) {
      rewriter.replaceOpWithNewOp<IREE::VM::ListGetRefOp>(
          srcOp, resultType, adaptor.getList(), index);
    } else {
      return srcOp.emitError() << "unsupported list element type in the VM";
    }
    return success();
  }
};

class ListSetOpConversion : public OpConversionPattern<IREE::Util::ListSetOp> {
  using Base::Base;
  LogicalResult
  matchAndRewrite(IREE::Util::ListSetOp srcOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto index = castToI32(adaptor.getIndex(), rewriter);
    auto valueType = adaptor.getValue().getType();
    if (valueType.isInteger(32)) {
      rewriter.replaceOpWithNewOp<IREE::VM::ListSetI32Op>(
          srcOp, adaptor.getList(), index, adaptor.getValue());
    } else if (valueType.isInteger(64)) {
      rewriter.replaceOpWithNewOp<IREE::VM::ListSetI64Op>(
          srcOp, adaptor.getList(), index, adaptor.getValue());
    } else if (valueType.isF32()) {
      rewriter.replaceOpWithNewOp<IREE::VM::ListSetF32Op>(
          srcOp, adaptor.getList(), index, adaptor.getValue());
    } else if (valueType.isF64()) {
      rewriter.replaceOpWithNewOp<IREE::VM::ListSetF64Op>(
          srcOp, adaptor.getList(), index, adaptor.getValue());
    } else if (!valueType.isIntOrIndexOrFloat()) {
      rewriter.replaceOpWithNewOp<IREE::VM::ListSetRefOp>(
          srcOp, adaptor.getList(), index, adaptor.getValue());
    } else {
      return srcOp.emitError() << "unsupported list element type in the VM";
    }
    return success();
  }
};

} // namespace

void populateUtilListToVMPatterns(MLIRContext *context,
                                  ConversionTarget &conversionTarget,
                                  TypeConverter &typeConverter,
                                  RewritePatternSet &patterns) {
  typeConverter.addConversion(
      [&typeConverter](IREE::Util::ListType type) -> std::optional<Type> {
        Type elementType;
        if (isa<IREE::Util::ObjectType>(type.getElementType()) ||
            isa<IREE::Util::VariantType>(type.getElementType())) {
          elementType = IREE::VM::OpaqueType::get(type.getContext());
        } else {
          elementType = typeConverter.convertType(type.getElementType());
        }
        if (!elementType) {
          return std::nullopt;
        }
        return IREE::VM::RefType::get(IREE::VM::ListType::get(elementType));
      });

  conversionTarget
      .addIllegalOp<IREE::Util::ListCreateOp, IREE::Util::ListConstructOp,
                    IREE::Util::ListSizeOp, IREE::Util::ListResizeOp,
                    IREE::Util::ListGetOp, IREE::Util::ListSetOp>();

  patterns.insert<ListCreateOpConversion, ListConstructOpConversion,
                  ListSizeOpConversion, ListResizeOpConversion,
                  ListGetOpConversion, ListSetOpConversion>(typeConverter,
                                                            context);
}

} // namespace mlir::iree_compiler
