// RUN: tf-opt %s -cse  | FileCheck %s
// Checks that CSE runs without generating invalid IR and doesn't CSE ops like
// NextIteration.Source and NextIteration.Sink.
// CHECK-LABEL: func @main
// CHECK: tf_executor.NextIteration.Source
// CHECK: tf_executor.NextIteration.Source
// CHECK: tf_executor.NextIteration.Sink
// CHECK: tf_executor.NextIteration.Sink
func.func @main() -> (tensor<1xi32>, tensor<1xi32>) {
  %0, %1 = tf_executor.graph {

}

def TfExecutor_NextIterationSourceOp : TfExecutor_Op<"NextIteration.Source",
    [HasParent<"GraphOp">]> {
  let summary = [{
    The "tf_executor.NextIteration.Source" is paired with a
    "tf_executor.NextIteration.sink" to represent NextIteration op in
    Tensorflow.
  }];

  let description = [{
    Tensorflow NextIteration operation forwards its input to the next iteration

# Verify that a NextIteration node feeding two different merge nodes is properly
# Imported.

# CHECK-LABEL: func @main()
# CHECK:         %[[NEXTITERATION:[a-z_0-9]+]], {{.*}} = tf_executor.NextIteration.Source
# CHECK:         tf_executor.Merge {{.*}}, %[[NEXTITERATION]]
# CHECK:         tf_executor.Merge {{.*}}, %[[NEXTITERATION]]

node {
  name: "Const"
  op: "Const"
  attr {

  // CHECK:  op: "RefMerge"
  // CHECK:  op: "RefSwitch"
  // CHECK:  op: "RefExit"
  // CHECK:  op: "RefNextIteration"
  tf_executor.graph {
    %0:3 = tf_executor.NextIteration.Source : tensor<*x!tf_type.int32ref> {device = "", T = "tfdtype$DT_INT32"} loc("while/NextIteration")

    tf_executor.NextIteration.Sink[%1#1] %1#0 : i32
    tf_executor.fetch %1#0 : i32
  }
  func.return %0 : i32
}

// Check that NextIteration.sink node is deleted when the source does not have
// any user other than the sink.
// CHECK-LABEL: func @nextiteration_deleted(
func.func @nextiteration_deleted(%arg0 : i32) -> i32 {
// CHECK-NOT: tf_executor.NextIteration.Source
// CHECK-NOT: tf_executor.NextIteration.Sink

  }
  func.return
}

// Tests tf_executor.NextIteration.Source/tf_executor.NextIteration.Sink has a
// device when an intermediate op in its loop has a device.

// CHECK-LABEL: func @testNextIterationNoDevice
func.func @testNextIterationNoDevice() {
  tf_executor.graph {
    // CHECK:      tf_executor.NextIteration.Source
    // CHECK-SAME: device = "/job:localhost/replica:0/task:0/device:TPU:0"

# Verify that importing a Graph with a backedge leads to two NextIteration nodes
# to break the cycle.

# CHECK-LABEL: func @main()
# CHECK:    %[[NEXTITERATION:[a-z0-9]+]], %[[NEXTITERATION_token:[a-z0-9]+]], {{.*}} = tf_executor.NextIteration.Source
# CHECK:    tf_executor.Merge {{.*}} %[[NEXTITERATION]]

# CHECK:    tf_executor.NextIteration.Sink[%[[NEXTITERATION_token]]]

node {
  name: "Const"
  op: "Const"

  // CHECK-NEXT: input: "while/Enter"
  // CHECK-NEXT: input: "while/NextIteration"
  // CHECK:      name: "while/NextIteration"
  // CHECK-NEXT: op: "NextIteration"
  // CHECK-NEXT: input: "while/Add"
  tf_executor.graph {
    %0:3 = tf_executor.NextIteration.Source : tensor<*xi32> {device = "", T = "tfdtype$DT_INT32"} loc("while/NextIteration")

    %NextIteration, %ctl_1 = NextIteration(%Add) name("while/NextIteration") {T = i32} : (tensor<*xi32>) -> (tensor<*xi32>)
    %Merge:2, %ctl_2 = Merge(%Enter, %NextIteration) name("while/Merge") {N = 2 : i64, T = i32} : (tensor<*xi32>, tensor<*xi32>) -> (tensor<*xi32>, tensor<*xi32>)

    %NextIteration, %ctl_1 = NextIteration(%Const) [%arg.ctl] name("while/NextIteration") {T = i32} : (tensor<i32>) -> (tensor<*xi32>)
    %Merge:2, %ctl_2 = Merge(%Enter, %NextIteration) [%arg.ctl] name("while/Merge") {N = 2 : i64, T = i32} : (tensor<*xi32>, tensor<*xi32>) -> (tensor<*xi32>, tensor<i32>)