#include "mlir/Transforms/RegionUtils.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_side_effects.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"

namespace mlir {
namespace tf_executor {
namespace {

#define GEN_PASS_DEF_EXECUTORGRAPHPRUNINGPASS

  %graph:2 = tf_executor.graph {
    %0:2 = tf_executor.island wraps "tf.Const"() {device = "", dtype = "tfdtype$DT_FLOAT", value = dense<0.000000e+00> : tensor<f32>} : () -> tensor<f32> loc("const")
    %1:2 = tf_executor.island wraps "tf.Identity"(%0#0) {T = "tfdtype$DT_FLOAT", device = ""} : (tensor<f32>) -> tensor<f32> loc("identity")

  func.func @main() attributes {tf.entry_function = {inputs = "", outputs = ""}} {
    tf_executor.graph {
      %0:2 = tf_executor.island wraps "tf.Const"() {device = "TPU:0", name = "const", dtype = "tfdtype$DT_INT32", value = dense<[1, 2]> : tensor<2xi32>} : () -> tensor<2xi32>
      tf_executor.fetch
    }
    func.return
  }
}

// CHECK-NOT: node

// CHECK: library
// CHECK-NEXT: function

// RUN: tf-mlir-translate -mlir-to-graphdef %s -o - | FileCheck %s

func.func @main() {
  tf_executor.graph {
    %0:2 = tf_executor.island wraps "tf.RecvTPUEmbeddingActivations"() {config = "test_config_recv_embedding"} : () -> tensor<512x256xf32> loc("RecvTPUEmbedding")

tf_executor::IslandOp CreateIsland(TypeRange result_types,
                                   const tf_executor::ControlType& control_type,
                                   const Location& loc, Operation& sub_op,
                                   tf_executor::IslandOp original_island) {
  OpBuilder builder(original_island);
  auto island = builder.create<tf_executor::IslandOp>(

#include "tensorflow/core/common_runtime/eager/eager_executor.h"

struct TFE_Executor {
  explicit TFE_Executor(bool async, bool enable_streaming_enqueue,
                        int in_flight_nodes_limit)
      : owned_executor(new tensorflow::EagerExecutor(
            async, enable_streaming_enqueue, in_flight_nodes_limit)) {}

  explicit TFE_Executor(tensorflow::EagerExecutor* executor)

  let summary = "Lifts tf_executor.island inner ops from a tf_executor.graph";

  let description = [{
    This pass converts tf_executor.graphs consisting of only tf_executor.islands and
    a tf_executor.fetch into a sea of nodes consisting of TensorFlow Dialect ops by
    lifting such ops out of a tf_executor.graph's tf_executor.islands. If V1 control

  %graph = tf_executor.graph {
    %result:2 = tf_executor.island wraps "tf.TPUPartitionedCall"(%arg0, %arg1) {config = "", config_proto = "", executor_type = "", f = @simple_callee} : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
    tf_executor.fetch %result#0 : tensor<2xi32>
  }
  func.return %graph : tensor<2xi32>
}

// CHECK:       library {

  %graph = tf_executor.graph {
    // This node would be renamed to bar1 [note: if imported from TF graphdef this would not be possible]
    %2:2 = tf_executor.island wraps "tf.Identity"(%arg1) {device = "", dtype = "tfdtype$DT_INT32"} : (tensor<10xi32>) -> tensor<10xi32> loc ("bar")
    // The following node would be renamed to bar2

  attributes  {tf.entry_function = {inputs = "args_0,args_1", outputs = "rets_0,rets_1"}} {
    %0:2 = tf_executor.graph {
      %1:3 = tf_executor.island wraps "tf.IdentityN"(%arg0, %arg1) {T = ["tfdtype$DT_FLOAT", "tfdtype$DT_INT32"], device = "", name = "identity"} : (tensor<*xf32>, tensor<2x4x6x8xi32>) -> (tensor<*xf32>, tensor<2x4x6x8xi32>)
      tf_executor.fetch %1#0, %1#1 : tensor<*xf32>, tensor<2x4x6x8xi32>
    }