}

// Visits an op and it's associated operands. IslandOps are handled differently
// where it's regions op operands are also visited as values may be implicitly
// captured within. NextIterationSourceOp will also visit it's associated
// NextIterationSinkOp.
void VisitOp(GraphOp graph, Operation* op,
             llvm::SmallPtrSetImpl<Operation*>* reachable_ops,

//===----------------------------------------------------------------------===//
// tf_executor.NextIteration.Source
//===----------------------------------------------------------------------===//

LogicalResult NextIterationSourceOp::verify() {
  NextIterationSourceOp source = *this;
  Value token = source.getToken();
  if (!token.hasOneUse())
    return source.emitOpError() << "expects a single user for produced token";

  for (Value value : operands)
    if (!mlir::isa<mlir::BlockArgument>(value)) filtered.push_back(value);
}

// Split the tfg.NextIteration into tf_executor::NextIterationSourceOp and
// tf_executor::NextIterationSinkOp to break the cycle introduced by itself.
static void SplitNextIteration(Block &block) {
  // TODO(b/207144333): Supports callback for unregistered ops

  auto& sink_island_control = sink_island_controls[0];
  island_op.getControl().replaceAllUsesWith(sink_island_control);
  // All existing outputs need to add sink_island_control as control input.
  // GraphOp, YieldOp and NextIterationSourceOp don't have control inputs so
  // exclude them below.
  for (Value out : island_op.getOutputs()) {
    for (auto& use : out.getUses()) {
      Operation* owner = use.getOwner();

  };

  // Check if tf_executor.NextIteration.Source/tf_executor.NextIteration.Sink
  // pair has matching devices or no devices.
  if (auto source = llvm::dyn_cast<tf_executor::NextIterationSourceOp>(op)) {
    return ops_have_same_device(source, source.GetSink());
  } else if (auto sink = llvm::dyn_cast<tf_executor::NextIterationSinkOp>(op)) {
    return ops_have_same_device(sink.GetSource(), sink);
  }

    // Replaces the input node with NextIteration sink if it is a NextIteration
    // source.
    if (auto next_iter_source =
            llvm::dyn_cast<mlir::tf_executor::NextIterationSourceOp>(
                input_inst))
      input_inst = next_iter_source.GetSink();

    auto node_it = nodes_.find(input_inst);
    TF_RET_CHECK(node_it != nodes_.end())

    // Replaces the input node with NextIteration sink if it is a NextIteration
    // source.
    if (auto next_iter_source =
            llvm::dyn_cast<mlir::tf_executor::NextIterationSourceOp>(
                input_inst))
      input_inst = next_iter_source.GetSink();

    auto node_it = nodes_.find(input_inst);
    TF_RET_CHECK(node_it != nodes_.end())

                              operands.end());
      $_state.attributes.append(attributes.begin(), attributes.end());
    }]>
   ];

  let extraClassDeclaration = [{
    NextIterationSourceOp GetSource();
  }];

  let assemblyFormat = " `[` $token `]` $input (`,` $controlInputs^)? `:` type($input) attr-dict";

  let hasCustomAssemblyFormat = 0;
}

                   mlir::tf_executor::EnterOp, mlir::tf_executor::ExitOp,
                   mlir::tf_executor::NextIterationSinkOp,
                   mlir::tf_executor::NextIterationSourceOp>(op);
}

mlir::LogicalResult IsValidGraph(mlir::ModuleOp module) {
  auto result = module.walk([&](Operation* op) {
    return IsControlFlowV1Op(op) ? mlir::WalkResult::interrupt()

                                           op->getResults());
  }
  if (auto iter_sink = dyn_cast<tf_executor::NextIterationSinkOp>(op)) {
    auto iter_source = cast<tf_executor::NextIterationSourceOp>(
        iter_sink.getToken().getDefiningOp());
    return RefineTypeForPassThroughOperands(
        op, iter_sink.getOperands().drop_front().take_front(),
        iter_source.getResults());
  }