#include "mlir/IR/BuiltinOps.h"
#include "mlir/Pass/Pass.h"

#include "src/Compiler/CompilerUtils.hpp"

using namespace mlir;

namespace onnx_mlir {

namespace {

struct MessagePass : PassWrapper<MessagePass, OperationPass<ModuleOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MessagePass)
  StringRef getArgument() const override { return "message-pass"; }
  StringRef getDescription() const override { return "Print compilation status"; }

  MessagePass(std::string message)
  : message(message) {}
  MessagePass(const MessagePass& pass) {}

  void runOnOperation() final { showCompilePhase(message); }

private:
  std::string message;
};

} // namespace

std::unique_ptr<Pass> createMessagePass(std::string message) { return std::make_unique<MessagePass>(message); }

} // namespace onnx_mlir
