/**
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */


#include <faiss/gpu/utils/DeviceMemory.h>
#include <faiss/gpu/utils/DeviceUtils.h>
#include <faiss/impl/FaissAssert.h>

namespace faiss { namespace gpu {

DeviceMemoryReservation::DeviceMemoryReservation()
    : state_(nullptr),
      device_(0),
      data_(nullptr),
      size_(0),
      stream_(0) {
}

DeviceMemoryReservation::DeviceMemoryReservation(DeviceMemory* state,
                                             int device,
                                             void* p,
                                             size_t size,
                                             cudaStream_t stream)
    : state_(state),
      device_(device),
      data_(p),
      size_(size),
      stream_(stream) {
}

DeviceMemoryReservation::DeviceMemoryReservation(
  DeviceMemoryReservation&& m) noexcept {

  state_ = m.state_;
  device_ = m.device_;
  data_ = m.data_;
  size_ = m.size_;
  stream_ = m.stream_;

  m.data_ = nullptr;
}

DeviceMemoryReservation::~DeviceMemoryReservation() {
  if (data_) {
    FAISS_ASSERT(state_);
    state_->returnAllocation(*this);
  }

  data_ = nullptr;
}

DeviceMemoryReservation&
DeviceMemoryReservation::operator=(DeviceMemoryReservation&& m) {
  if (data_) {
    FAISS_ASSERT(state_);
    state_->returnAllocation(*this);
  }

  state_ = m.state_;
  device_ = m.device_;
  data_ = m.data_;
  size_ = m.size_;
  stream_ = m.stream_;

  m.data_ = nullptr;

  return *this;
}

DeviceMemory::~DeviceMemory() {
}

} } // namespace
