/**
 * 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.
 */


#pragma once

#include <faiss/gpu/GpuResources.h>
#include <cuda.h>
#include <unordered_set>

namespace faiss { namespace gpu {

/// Allocator for Thrust that comes out of a specified memory space
class GpuResourcesThrustAllocator {
 public:
  typedef char value_type;

  inline GpuResourcesThrustAllocator(GpuResources* res,
                                     cudaStream_t stream,
                                     void* mem,
                                     size_t size)
      : res_(res),
        stream_(stream),
        start_((char*) mem),
        cur_((char*) mem),
        end_((char*) mem + size) {
  }

  inline ~GpuResourcesThrustAllocator() {
    // In the case of an exception being thrown, we may not have called
    // deallocate on all of our sub-allocations. Free them here
    for (auto p : mallocAllocs_) {
      res_->deallocMemory(getCurrentDevice(), p);
    }
  }

  inline char* allocate(std::ptrdiff_t size) {
    if (size <= (end_ - cur_)) {
      char* p = cur_;
      cur_ += size;
      FAISS_ASSERT(cur_ <= end_);

      return p;
    } else {
      // FIXME: we cannot use temporary memory for new requests because the
      // current temporary memory allocator cannot handle stream synchronization
      // at present, so just allocate through the general device
      char* p =
        (char*) res_->allocMemory(
          AllocRequest(makeDevAlloc(AllocType::Other, stream_), size));

      mallocAllocs_.insert(p);
      return p;
    }
  }

  inline void deallocate(char* p, size_t size) {
    // Allocations could be returned out-of-order; ignore those we
    // didn't cudaMalloc
    auto it = mallocAllocs_.find(p);
    if (it != mallocAllocs_.end()) {
      res_->deallocMemory(getCurrentDevice(), p);
      mallocAllocs_.erase(it);
    }
  }

 private:
  GpuResources* res_;
  cudaStream_t stream_;
  char* start_;
  char* cur_;
  char* end_;
  std::unordered_set<char*> mallocAllocs_;
};


} } // namespace
