# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

# RUN: %p/hash_table_v1 | FileCheck %s

# pylint: disable=missing-docstring,line-too-long
import tensorflow.compat.v1 as tf
from tensorflow.compiler.mlir.tensorflow.tests.tf_saved_model import common_v1

// HashTable is initialized if it appears in the initializer since it can't
// check the actual value.
bool IsResourceInitialized(ModuleOp module_op, Operation* hash_table) {
  StringRef shared_name = GetSharedName(hash_table);
  if (shared_name.empty()) return false;

  for (func::FuncOp init_func_op :
       tf_saved_model::GetInitializerFunctions(module_op)) {

    // CHECK: [[HASHTABLE:%.+]] = "tf.HashTableV2"()
    // CHECK: "tf.InitializeTableFromTextFileV2"([[HASHTABLE]], [[CST]])
    // CHECK: [[HASHTABLE_1:%.+]] = "tf.HashTableV2"()
    // CHECK: "tf.InitializeTableFromTextFileV2"([[HASHTABLE_1]], [[CST_1]])
    func.return
  }
}

// -----

// Test running the pass on a module that does not have
// tf_saved_model.semantics.

  // CHECK-NEXT:  [[HASH_TABLE:%.*]] = "tfl.hashtable"() <{key_dtype = !tf_type.string, table_id = -1323619995 : i32, value_dtype = i64}> : () -> tensor<1x!tf_type.resource>
  // CHECK-NEXT:   "tfl.hashtable_import"([[HASH_TABLE]], [[CST]], [[CST_1]]) : (tensor<1x!tf_type.resource>, tensor<3x!tf_type.string>, tensor<3xi64>) -> ()

py_strict_binary(
    name = "keras",
    srcs = ["keras.py"],
    deps = [
        ":common",
        "//tensorflow:tensorflow_py",
    ],
)

py_strict_binary(
    name = "hash_table_v1",
    srcs = ["hash_table_v1.py"],
    deps = [
        ":common_v1",
        "//tensorflow:tensorflow_py",
    ],
)

py_strict_binary(
    name = "include_variables_in_init_v1",

    return
  }
  // CHECK: @init_func_restore_op
  // CHECK-SAME: tf_saved_model.initializer_type = "restore_op"
  // CHECK: %[[CST:.*]] = "tf.Const"()
  // CHECK-NEXT: %[[HASH_TABLE:.*]] = "tf.HashTableV2"()
  // CHECK-NEXT: "tf.AssignVariableOp"(%[[HASH_TABLE]], %[[CST]])
}

// -----


// Nothing happens when there are no `tf_saved_model.session_initializer`.

// CHECK-LABEL: module

# CHECK-SAME: [[ARG1:%.*]]: tensor<!tf_type.string> {tf_saved_model.bound_input = @[[asset1]]}
# 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()