* See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.gradle.internal.isolation;

import org.gradle.internal.hash.Hashable;
import org.gradle.internal.snapshot.ValueSnapshot;

import javax.annotation.Nullable;

/**
 * Isolatable objects can return an isolated instance of the given type T from which this object was created.

module attributes {tf_saved_model.semantics} {

// Test hoisting hash table op.

// CHECK-LABEL: func @_tfrt_resource_init
// CHECK: [[handle:%.*]] = "tf.HashTableV2"()
// CHECK-SAME: shared_name = "x"
// CHECK: "tf._TfrtSetResource"([[handle]]) <{index = [[handle_id:.*]] : i64}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"}
// CHECK: [[x:%.*]] = "tf.LookupTableSizeV2"([[handle]])

  let arguments = (ins
    TFL_TensorOf<[Complex<F<32>>, Complex<F<64>>]>:$input
  );

  let results = (outs
    TFL_TensorOf<[F32, F64]>:$output
  );
}

def TFL_HashtableOp: TFL_Op<"hashtable", []> {
  let summary = "Creates a non-initialized hash table.";
  let description = [{
This op creates a hash table, specifying the type of its keys and values.

           device == rhs.device && op == rhs.op;
  }

  // Make ResourceHandle hashable.
  friend ::llvm::hash_code hash_value(const ResourceHandle& resource_handle);

  StringRef container;
  StringRef name;
  StringRef device;
  Operation* op = nullptr;
};

// Make ResourceHandle hashable.
inline ::llvm::hash_code hash_value(const ResourceHandle& resource_handle) {

    %cst = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi64>
    %cst_0 = arith.constant dense<["a", "b", "c", "d"]> : tensor<4x!tf_type.string>
    %0 = "tf.HashTableV2"() {container = "", device = "", key_dtype = i64, shared_name = "hash_table_dba2ccaa-f1b1-46d6-b276-98008f69da71", use_node_name_sharing = false, value_dtype = !tf_type.string} : () -> tensor<!tf_type.resource>

      %out_0, %ctl_1 = tf_executor.island wraps "tf.Const"() {device = "", value = dense<[1]> : tensor<1xi64>} : () -> tensor<1xi64>
      %out_1, %ctl_2 = tf_executor.island wraps "tf.HashTableV2"() {container = "", device = "", key_dtype = !tf_type.string, shared_name = "1", use_node_name_sharing = false, value_dtype = i64} : () -> tensor<!tf_type.resource>

    // expected-remark@+1 {{Cost: 1}}
    %0 = "tf.HashTableV2"() {container = "", device = "/job:localhost/replica:0/task:0/device:CPU:0", key_dtype = !tf_type.string, shared_name = "hash_table", use_node_name_sharing = false, value_dtype = i64} : () -> tensor<!tf_type.resource>
    // expected-remark@+1 {{Cost: 1024}}

}

// CHECK-LABEL: @merge_hashtable_ops
// CHECK: %[[OUT_0:.*]], %[[CTL_0:.*]] = tf_executor.island wraps "tf.HashTableV2"()
// CHECK: %[[OUT_1:.*]], %[[CTL_1:.*]] = tf_executor.island wraps "tf.LookupTableSizeV2"(%[[OUT_0]])

// Check that LookupTableImportV2 is using the same HashTableV2 with LookupTableSizeV2.

    "tf.AssignVariableOp"(%0, %1) : (tensor<*x!tf_type.resource<tensor<!tf_type.string>>>, tensor<!tf_type.string>) -> ()
    %2 = "tf.HashTableV2"() {key_dtype = !tf_type.string, shared_name = "table", value_dtype = i64} : () -> tensor<!tf_type.resource>
    %3 = "tf.ReadVariableOp"(%0) : (tensor<*x!tf_type.resource<tensor<!tf_type.string>>>) -> tensor<!tf_type.string>

              StrEq("arg_0:0"));
}

TEST_F(ConvertMlirModuleToExportedModelTest,
       InitNodeNameSetToLocOfControlOutput) {
  // Define a module that initializes a tf.HashTableV2 whose control output node
  // for the initialization is named "init_op_init_all_tables".
  mlir::OwningOpRef<mlir::ModuleOp> module_op = ParseModuleOpString(R"mlir(
    module attributes {tf_saved_model.semantics} {