using mlir::OpBuilder;
using mlir::Operation;
using mlir::OperationState;
using mlir::UnknownLoc;
using mlir::chlo::ChloDialect;
using mlir::TF::TensorFlowDialect;
using tensorflow::tf2xla::internal::IsInBridgeAcceptableDialects;

class SharedUtilsTest : public ::testing::Test {};

TEST_F(SharedUtilsTest, IsInFunctionalDialectPasses) {
  MLIRContext context;
  context.loadDialect<TensorFlowDialect>();

namespace tensorflow {
namespace tf2xla {
namespace internal {

// Returns true if the op has a valid namespace during clustering & tf dialect
// to executor components of the Bridge.
bool IsInBridgeAcceptableDialects(mlir::Operation* op);

}  // namespace internal
}  // namespace tf2xla
}  // namespace tensorflow

#include "mlir/IR/Operation.h"  // from @llvm-project
#include "mlir/IR/Visitors.h"  // from @llvm-project

namespace tensorflow {
namespace tf2xla {
namespace internal {

bool IsInBridgeAcceptableDialects(mlir::Operation* op) {
  const std::set<std::string> kBuiltinNamespaces = {"func", "return",
                                                    "builtin"};

};

void VerifyClusteringPass::runOnOperation() {
  Operation* func_op = getOperation();

  auto walk_result = func_op->walk([&](Operation* op) {
    if (!tensorflow::tf2xla::internal::IsInBridgeAcceptableDialects(op)) {
      std::string error = "op is in dialect " +
                          op->getDialect()->getNamespace().str() +
                          " not in tf functional dialect";
      op->emitError() << error;

}

void VerifyInputDialectToExecutorPass::runOnOperation() {
  Operation* func_op = getOperation();

  auto walk_result = func_op->walk([&](Operation* op) {
    if (!tensorflow::tf2xla::internal::IsInBridgeAcceptableDialects(op)) {
      std::string error = "op is in dialect " +
                          op->getDialect()->getNamespace().str() +
                          " which is not an accepted dialect";
      op->emitError() << error;