// RUN: tf-opt -tf-init-text-file-to-import-saved-model-test %s | FileCheck %s

// Tests that the tf.InitializeTableFromTextFileV2 op are inlined.

func.func @init_all_tables() {
  %cst = arith.constant dense<"tokens.txt"> : tensor<!tf_type.string>
  %0 = "tf.HashTableV2"() {container = "", device = "", key_dtype = !tf_type.string, shared_name = "hash_table_/tmp/vocab.txt_-2_-1", use_node_name_sharing = false, value_dtype = i64} : () -> tensor<!tf_type.resource>

  "tf.InitializeTableFromTextFileV2"(%0, %cst) {delimiter = " ", device = "", key_index = -1 : i64, value_index = -1 : i64, vocab_size = -1 : i64} : (tensor<!tf_type.resource>, tensor<!tf_type.string>) -> ()
  func.return
  // CHECK: [[VAL:%.*]] = "tf.HashTableV2"()
  // CHECK: tf.InitializeTableFromTextFileV2"
}

// -----

// Tests that the tf.InitializeTableFromTextFileV2 op is not converted since

    "tf.InitializeTableFromTextFileV2"(%0, %arg0) {delimiter = "\09", device = "", key_index = -2 : i64, offset = 0 : i64, value_index = -1 : i64, vocab_size = 437 : i64} : (tensor<!tf_type.resource>, tensor<!tf_type.string>) -> ()
    // CHECK: "tf.InitializeTableFromTextFileV2"(%0, %arg0)
    func.return
  }
}

// -----

// RUN: tf-opt -tf-init-text-file-to-import-test %s | FileCheck %s

// Tests that the tf.InitializeTableFromTextFileV2 op are inlined.

func.func @init_all_tables() {
  %cst = arith.constant dense<"%FILE_PLACEHOLDER"> : tensor<!tf_type.string>
  %0 = "tf.HashTableV2"() {container = "", device = "", key_dtype = !tf_type.string, shared_name = "hash_table_/tmp/vocab.txt_-2_-1", use_node_name_sharing = false, value_dtype = i64} : () -> tensor<!tf_type.resource>

          %c_9 = tf_executor.island wraps "tf.InitializeTableFromTextFileV2"(%o, %arg1) <{delimiter = "\09", key_index = -2 : i64, value_index = -1 : i64, vocab_size = -1 : i64}> {_has_manual_control_dependencies = true, device = ""} : (tensor<!tf_type.resource>, tensor<!tf_type.string>) -> ()

# CHECK-SAME: tf_saved_model.initializer_type = "init_op"
# CHECK-NEXT: [[R0:%.*]] = "tf.HashTableV2"()
# CHECK-SAME: shared_name = "[[hash_table:.*]]"
# CHECK-NEXT: "tf.InitializeTableFromTextFileV2"([[R0]], [[ARG0]])


def write_vocabulary_file(vocabulary):
  """Write temporary vocab file for module construction."""
  tmpdir = tempfile.mkdtemp()
  vocabulary_file = os.path.join(tmpdir, 'tokens.txt')

    %3 = "tf.ReadVariableOp"(%0) : (tensor<*x!tf_type.resource<tensor<!tf_type.string>>>) -> tensor<!tf_type.string>
    "tf.InitializeTableFromTextFileV2"(%2, %3) {key_index = -2 : i64, value_index = -1 : i64} : (tensor<!tf_type.resource>, tensor<!tf_type.string>) -> ()
    func.return
  }
  // CHECK: func @init()
  // CHECK-NOT: tf.VarHandleOp

    %0 = "tf.HashTableV2"() {container = "", device = "", key_dtype = !tf_type.string, shared_name = "", use_node_name_sharing = false, value_dtype = i64} : () -> tensor<!tf_type.resource>
    "tf.InitializeTableFromTextFileV2"(%0, %arg0) {delimiter = "\09", device = "", key_index = -2 : i64, offset = 0 : i64, value_index = -1 : i64, vocab_size = 437 : i64} : (tensor<!tf_type.resource>, tensor<!tf_type.string>) -> ()
    func.return
  }

    %0 = "tf.HashTableV2"() {container = "", device = "", key_dtype = !tf_type.string, shared_name = "hash_table_shared_name", use_node_name_sharing = false, value_dtype = i64} : () -> tensor<!tf_type.resource>
    "tf.InitializeTableFromTextFileV2"(%0, %arg0) {delimiter = "\09", device = "", key_index = -2 : i64, value_index = -1 : i64, vocab_size = 205 : i64} : (tensor<!tf_type.resource>, tensor<!tf_type.string>) -> ()
    func.return
  }

  let constructor = "TF::CreateRewriteTPUEmbeddingOpsPass()";
}

def InitTextFileToImportPass : Pass<"tf-init-text-file-to-import", "mlir::func::FuncOp"> {
  let summary = "convert InitializeTableFromTextFileV2 ops to LookupTableImportV2Op "
           "to remove the dependency on asset files";

  let constructor = "TF::CreateInitTextFileToImportPass()";
  let options = [