// Copyright 2020 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 <utility>

#include "iree/compiler/Dialect/VM/IR/VMDialect.h"
#include "iree/compiler/Dialect/VM/IR/VMOps.h"
#include "iree/compiler/Dialect/VM/IR/VMTypes.h"
#include "iree/compiler/Dialect/VM/Transforms/Passes.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Diagnostics.h"
#include "mlir/Pass/Pass.h"

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

#define GEN_PASS_DEF_HOISTINLINEDRODATAPASS
#include "iree/compiler/Dialect/VM/Transforms/Passes.h.inc"

class HoistInlinedRodataPass
    : public IREE::VM::impl::HoistInlinedRodataPassBase<
          HoistInlinedRodataPass> {
  void runOnOperation() override {
    IREE::VM::ModuleOp moduleOp = getOperation();
    SymbolTable moduleSymbolTable(moduleOp);

    // Find all inline byte buffers in the module.
    SmallVector<IREE::VM::RodataInlineOp> inlineOps;
    moduleOp.walk([&](IREE::VM::RodataInlineOp inlineOp) {
      inlineOps.push_back(inlineOp);
    });

    for (auto inlineOp : inlineOps) {
      auto *parentOp = findParentContainer(inlineOp);
      OpBuilder moduleBuilder(moduleOp.getContext());
      if (parentOp) {
        moduleBuilder.setInsertionPoint(parentOp);
      } else {
        moduleBuilder.setInsertionPointToStart(&moduleOp.getBlock());
      }
      auto rodataOp = IREE::VM::RodataOp::create(
          moduleBuilder, inlineOp.getLoc(),
          inferConstantName(parentOp, inlineOp), inlineOp.getValue());
      if (auto alignmentAttr = inlineOp.getAlignmentAttr()) {
        rodataOp.setAlignmentAttr(alignmentAttr);
      }
      if (auto mimeTypeAttr = inlineOp.getMimeTypeAttr()) {
        rodataOp.setMimeTypeAttr(mimeTypeAttr);
      }
      moduleSymbolTable.insert(rodataOp);
      rodataOp.setPrivate();
      replaceInlineOpWithRodataRef(inlineOp, rodataOp);
    }
  }

private:
  Operation *findParentContainer(IREE::VM::RodataInlineOp inlineOp) {
    if (auto parentOp = inlineOp->getParentOfType<IREE::VM::InitializerOp>()) {
      return parentOp;
    } else if (auto parentOp = inlineOp->getParentOfType<IREE::VM::FuncOp>()) {
      return parentOp;
    }
    return nullptr;
  }

  std::string inferConstantName(Operation *parentOp,
                                IREE::VM::RodataInlineOp inlineOp) {
    if (auto nameAttr = inlineOp.getNameAttr()) {
      return nameAttr.str();
    }
    if (auto symbolOp = dyn_cast<SymbolOpInterface>(parentOp)) {
      return (symbolOp.getName() + "_const").str();
    }
    return "_const";
  }

  // Replaces a vm.rodata.inline op with a vm.const.ref.rodata op that
  // references the module-level |rodataOp|.
  void replaceInlineOpWithRodataRef(IREE::VM::RodataInlineOp inlineOp,
                                    IREE::VM::RodataOp rodataOp) {
    OpBuilder builder(inlineOp);
    auto refOp = IREE::VM::ConstRefRodataOp::create(builder, inlineOp.getLoc(),
                                                    rodataOp);
    inlineOp.replaceAllUsesWith(refOp.getValue());
    inlineOp.erase();
  }
};

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