// Test that no ops are added to @main when @tf_quant__save function does
// not exist.

module attributes {tf_saved_model.semantics} {
  func.func @main(%arg: tensor<!tf_type.string> {tf_saved_model.index_path = ["__tf_file_prefix"]}) -> ()
      attributes {tf.entry_function = {inputs = "tf_file_prefix:0", outputs = ""}, tf_saved_model.exported_names = ["main"]} {
    tf_executor.graph {
      tf_executor.fetch
    }
    return
  }
}

# CHECK: func {{@[a-zA-Z_0-9]+}}(
# CHECK-SAME: %[[ARG_1:.*]]: tensor<3x1xf32> {tf_saved_model.index_path = ["x"]}
# CHECK-SAME: %[[ARG_2:.*]]: tensor<!tf_type.resource<tensor<1x3xf32>>> {tf_saved_model.bound_input = @y}
# CHECK-SAME: -> (tensor<3x3xf32> {tf_saved_model.index_path = ["r"]})
# CHECK-SAME: attributes {{.*}} tf_saved_model.exported_names = ["key"]

# CHECK-SAME:             -> (tensor<3x3xf32> {tf_saved_model.index_path = ["r"]})
# CHECK-SAME: attributes {{.*}} tf_saved_model.exported_names = ["key"]

# CHECK-NEXT: [[R0:%.*]] = "tf.ReadVariableOp"([[ARG1]]) {{{.*}}} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
# CHECK-NEXT: [[R1:%.*]] = "tf.MatMul"([[ARG0]], [[R0]]) <{{{.*}}}> {{{.*}}} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32>

// Test for case with no session initialize op.
module attributes {tf_saved_model.semantics} {
  // CHECK-LABEL: serving_default
  func.func @serving_default(%arg0: tensor<1x10xf32> {tf_saved_model.index_path = ["x"]}) ->

# CHECK-SAME:             -> (tensor<3x3xf32> {tf_saved_model.index_path = ["r"]})
# CHECK-SAME: attributes {{.*}} tf_saved_model.exported_names = ["key"]

# CHECK:      func {{@[a-zA-Z_0-9]+}}(
# CHECK-SAME:   [[ARG2:%.*]]: tensor<3x1xf32> {tf_saved_model.index_path = ["x"]},
# CHECK-SAME:   [[ARG3:%.*]]: tensor<!tf_type.resource<tensor<1x3xf32>>> {tf_saved_model.bound_input = @[[VAR]]})

// CHECK: func @test_hoist_hash_table
func.func @hoist_hash_table(%arg: tensor<?x!tf_type.string> {tf_saved_model.index_path = ["input"]}, %default: tensor<i64> {tf_saved_model.index_path = ["default"]}) -> (tensor<i64> {tf_saved_model.index_path = ["r"]}, tensor<*xi64> {tf_saved_model.index_path = ["r1"]})
  attributes {tf_saved_model.exported_names = ["test_hoist_hash_table"]} {
  // CHECK-NOT: tf.HashTableV2
  // CHECK-NOT: tf.LookupTableSizeV2

// CHECK: func @test_hoist_hash_table
func.func @hoist_hash_table(%arg: tensor<?x!tf_type.string> {tf_saved_model.index_path = ["input"]}, %default: tensor<i64> {tf_saved_model.index_path = ["default"]}) -> (tensor<i64> {tf_saved_model.index_path = ["r"]}, tensor<*xi64> {tf_saved_model.index_path = ["r1"]})
  attributes {tf_saved_model.exported_names = ["test_hoist_hash_table"]} {
  // CHECK-NOT: tf.HashTableV2
  // CHECK-NOT: tf.LookupTableSizeV2

# CHECK-SAME: min_consumer
# CHECK-SAME: producer

# CHECK: tf_saved_model.session_initializer
# CHECK-SAME: initializers = [@[[restore:.*]]]

# CHECK: "tf_saved_model.asset"()
# CHECK-SAME: <{filename = [[filename:.*]], sym_name = "[[sym_name:.*]]"}> : () -> ()

# CHECK:      func @[[restore]](
# CHECK-SAME:   [[variable_path:%.*]]: tensor<!tf_type.string> {tf_saved_model.bound_input = @[[sym_name]]}

module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1506 : i32}, tf_saved_model.semantics} {
  "tf_saved_model.session_initializer"() {initializers = [@init_all_tables]} : () -> ()
  func.func @init_all_tables() attributes {tf_saved_model.exported_names = ["__tf_saved_model_session_initializer_init_all_tables"], tf_saved_model.initializer_type = "init_op"} {

#include "mlir/Support/LogicalResult.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"

namespace mlir {
namespace tf_saved_model {

// Trace upwards until we hit a (block) argument a value originates from.
static LogicalResult traceUpwardsToArgument(Value v, llvm::DenseSet<Value> seen,