#include <c10/core/impl/cow/context.h>

#include <c10/core/impl/cow/deleter.h>
#include <c10/util/Exception.h>
#include <mutex>

namespace c10::impl {

cow::Context::Context(std::unique_ptr<void, DeleterFnPtr> data)
    : data_(std::move(data)) {
  // We never wrap a Context.
  TORCH_INTERNAL_ASSERT(data_.get_deleter() != cow::delete_context);
}

auto cow::Context::increment_refcount() -> void {
  auto refcount = ++refcount_;
  TORCH_INTERNAL_ASSERT(refcount > 1);
}

auto cow::Context::decrement_refcount()
    -> std::variant<NotLastReference, LastReference> {
  auto refcount = --refcount_;
  TORCH_INTERNAL_ASSERT(refcount >= 0, refcount);
  if (refcount == 0) {
    std::unique_lock lock(mutex_);
    auto result = std::move(data_);
    lock.unlock();
    delete this;
    return {std::move(result)};
  }

  return std::shared_lock(mutex_);
}

cow::Context::~Context() {
  TORCH_INTERNAL_ASSERT(refcount_ == 0);
}

} // namespace c10::impl
