#include "tsl/platform/status.h"

namespace tensorflow {

VariableInfo::VariableInfo(
    int index, absl::string_view name, Var* var,
    const std::optional<ManagedStackTrace>& definition_stack_trace)
    : index_(index),
      name_(name),
      var_(var),
      definition_stack_trace_(definition_stack_trace) {}

VariableInfo::VariableInfo(VariableInfo&& other)
    : index_(other.index_),
      var_(other.var_),

// refcount on destruction.
class VariableInfo {
 public:
  explicit VariableInfo(int index, absl::string_view name, Var* var,
                        const std::optional<ManagedStackTrace>&
                            definition_stack_trace = std::nullopt);
  VariableInfo(VariableInfo&& other);

  VariableInfo& operator=(VariableInfo&& other);

  VariableInfo(const VariableInfo&) = delete;

        "defs.h",
    ],
    visibility = [":friends"],
)

# Internal targets below this point.

cc_library(
    name = "variable_info",
    srcs = ["variable_info.cc"],
    hdrs = ["variable_info.h"],
    visibility = [
        ":internal",
        # We reuse VariableInfo in TFRT's implementation of TpuExecuteOp.
        "//learning/brain/tfrt/tf_tpu:__pkg__",
        "//learning/brain/tfrt/tpu_plugin:__pkg__",

          *ptr = new Var(DT_INVALID);
          return absl::OkStatus();
        }));
    VariableInfo& variable_info = result->emplace_back(
        var_idx, handle.name(), variable, handle.definition_stack_trace());
    if (variables_updated != nullptr &&
        variables_updated->find(var_idx) == variables_updated->end()) {
      variable_info.set_read_only();
    }
  }
  return absl::OkStatus();
}

//
// If the variable is read_only(), only acquires reader locks.
Status LockVariables(absl::Span<VariableInfo*> variables)
    TF_EXCLUSIVE_LOCK_FUNCTION();
Status LockVariables(absl::Span<VariableInfo> variables)
    TF_EXCLUSIVE_LOCK_FUNCTION();

// Returns a vector of VariableInfo instances for the resource variable inputs,
// given that *all* inputs are in `inputs`. The input indices for the resource

#define TENSORFLOW_COMPILER_JIT_XLA_COMPILE_ON_DEMAND_OP_H_

#include <vector>

#include "tensorflow/compiler/jit/device_compilation_profiler.h"
#include "tensorflow/compiler/jit/variable_info.h"
#include "tensorflow/compiler/jit/variable_info_util.h"
#include "tensorflow/compiler/jit/xla_launch_util.h"
#include "tensorflow/compiler/jit/xla_platform_info.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"

#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/jit/pjrt_device_compiler_client.h"
#include "tensorflow/compiler/jit/pjrt_device_context.h"
#include "tensorflow/compiler/jit/variable_info.h"
#include "tensorflow/compiler/jit/variable_info_util.h"
#include "tensorflow/compiler/jit/xla_launch_util.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/pjrt/pjrt_client.h"

  absl::StatusOr<std::vector<VariableInfo>> variable_infos =
      GatherVariableInfo(ctx, *result, 0);
  TF_RETURN_IF_ERROR(variable_infos.status());
  TF_RETURN_IF_ERROR(LockVariables(absl::MakeSpan(*variable_infos)));
  TF_RETURN_IF_ERROR(launch_context.PopulateOutputs(
      ctx, result, execution_output.ConsumeResult(),
      /*missing_ctx_input_prefix=*/0, absl::MakeSpan(*variable_infos),
      input_output_alias, snapshot_ptrs));

  // operate on TF_VariableInfo (such as allocate temp tensor for the `var` held
  // by the underlying tensorflow::VariableInfo.
  TF_VariableInfo(int index, const std::string& name, tensorflow::Var* var) {
    var_info = tensorflow::VariableInfo{index, name, var};
  }

  tensorflow::VariableInfo var_info{0, "", nullptr};
};

TF_VariableInfo* TF_CreateVariableInfoFromContext(TF_OpKernelContext* ctx,

  Status PopulateOutputs(
      OpKernelContext* ctx,
      const XlaCompiler::CompilationResult* compilation_result,
      xla::ScopedShapedBuffer output, int missing_ctx_input_prefix,
      absl::Span<VariableInfo> variable_infos,
      const xla::HloInputOutputAliasConfig& input_output_alias,
      const std::map<int, const Tensor*>& resource_vars);

 private:
  xla::LocalClient* client_;