// Test that VarHandleOp is not marked as "potentially written" if it is not
// assigned inside the function called by "tf.BatchFunction".

module {
// CHECK-LABEL: @serving_default
  func.func @serving_default() -> (tensor<*xi32>) {
    // expected-remark@below {{device: "", container: "", shared_name: "var_0", is_potentially_written: false}}

    %0 = "tf.VarHandleOp"() {container = "", shared_name = "v"} : () -> tensor<!tf_type.resource<tensor<3xf32>>>
    %1 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<3xf32>>>) -> tensor<3xf32>
    %2 = "tf.VarHandleOp"() {container = "", shared_name = "v"} : () -> tensor<!tf_type.resource<tensor<?xf32>>>

  }

  StringRef getDescription() const final {
    return "Removes `tf.AssignVariableOp(tf.VarHandleOp, tf.Const)` patterns "
           "from the initializer function of type 'restore_op'.";
  }

  void runOnOperation() override;
};

// Finds and removes the `tf.AssignVariableOp(tf.VarHandleOp, tf.Const)`
// pattern. `tf.VarHandleOp` and `tf.Const` are removed unless they are used by
// other ops.

  attributes {tf.entry_function = {control_outputs = "", inputs = "", outputs = "dense_2/Add:0"}, tf_saved_model.exported_names = ["serving_default"]} {
    // CHECK: "tf.VarHandleOp"
    %0 = "tf.VarHandleOp"() {_class = ["loc:@dense/kernel"], allowed_devices = [], container = "", device = "invalid", shared_name = "invalid_var"} : () -> tensor<!tf_type.resource<tensor<100x50xf32>>>

    %1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
    "tf.AssignVariableOp"(%1, %0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
    return
  }
}

// -----
// single variable: VarHandleOp is before RestoreV2

// CHECK-LABEL:   func.func @varhandle_before_restore() {

    // CHECK: tf_executor.island wraps "tf.VarHandleOp"() <{container = "a", shared_name = "x"}> {_mlir_name = "x", _output_shapes = [#tf_type.shape<>], allowed_devices = [], device = "/device:CPU:0", dtype = i64, shape = #tf_type.shape<>} : () -> tensor<!tf_type.resource<tensor<i64>>>

    %handle_0 = "tf.VarHandleOp"() {container="c", shared_name="a"} : () -> tensor<!tf_type.resource<tensor<1x10xf32>>>
    %handle_1 = "tf.VarHandleOp"() {container="c", shared_name="b"} : () -> tensor<!tf_type.resource<tensor<1x10xi64>>>

  } else if (auto next = llvm::dyn_cast<TF::IteratorGetNextOp>(op)) {
    resource = next.getIterator();
  }

  if (llvm::dyn_cast_or_null<TF::VarHandleOp>(resource.getDefiningOp())) {
    return;  // We're already directly after a VarHandleOp.
  }

  const TF::ResourceDataflowState* state =
      solver.lookupState<TF::ResourceDataflowState>(resource);
  if (!state) {

// Variable is used by both CPU and TPU
//
// CHECK-LABEL: func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32>
// CHECK-NEXT:    [[HANDLE:%.*]] = "tf.VarHandleOp"()
// CHECK-NEXT:    [[ARRAYKEY:%.*]], [[FURTURE:%.*]] = "tf_mlrt.tf_ifrt_load_variable"([[HANDLE]])
// CHECK-SAME:       <{used_by_host = true}> : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> (tensor<!tf_type.string>, !mlrt.future)

# RUN: tf-mlir-translate -graphdef-to-mlir -tf-enable-shape-inference-on-import=false %s -tf-output-arrays=func_call -o - | FileCheck %s

node {
  name: "x"
  op: "VarHandleOp"
  device: "/CPU:0"
  attr {
    key: "container"
    value {
      s: "a"
    }
  }
  attr {
    key: "dtype"
    value {
      type: DT_INT64
    }
  }
  attr {
    key: "shape"
    value {
      shape {
      }
    }
  }