//
// This file is part of MorphoGraphX - http://www.MorphoGraphX.org
// Copyright (C) 2012-2016 Richard S. Smith and collaborators.
//
// If you use MorphoGraphX in your work, please cite:
//   http://dx.doi.org/10.7554/eLife.05864
//
// MorphoGraphX is free software, and is licensed under under the terms of the 
// GNU General (GPL) Public License version 2.0, http://www.gnu.org/licenses.
// 
#include <CachedAlloc.hpp>

// Modified for use with MorphoGraphX
//
// Example by Nathan Bell and Jared Hoberock
// (modified by Mihail Ivakhnenko)
//
// This example demonstrates how to intercept calls to get_temporary_buffer
// and return_temporary_buffer to control how Thrust allocates temporary storage
// during algorithms such as thrust::reduce. The idea will be to create a simple
// cache of allocations to search when temporary storage is requested. If a hit
// is found in the cache, we quickly return the cached allocation instead of
// resorting to the more expensive thrust::cuda::malloc.
//
// Note: this implementation cached_allocator is not thread-safe. If multiple
// (host) threads use the same cached_allocator then they should gain exclusive
// access to the allocator before accessing its methods.

#ifdef THRUST_BACKEND_CUDA
namespace mgx
{
  // One instance shared by all
  CachedAllocator cachedAlloc;
  
  char* CachedAllocator::allocate(size_t n)
  {
    char* result = 0;
 
    // search the cache for a free block
    FreeBlockList::iterator iter = freeBlocks.begin();
    while(iter != freeBlocks.end()) {
      if(iter->first == n) {
        result = iter->second;
        freeBlocks.erase(iter);
        break;
      }
      ++iter;
    }
    // if no allocation of the right size exists, create one
    if(!result)
      result = thrust::cuda::malloc<char>(n).get();

    // insert the allocated pointer into the allocated_blocks map
    allocatedBlocks.insert(std::make_pair(result, n));
 
    return result;
  }

  void CachedAllocator::deallocate(char* ptr, size_t sz)
  {
    // erase the allocated block from the allocated blocks map
    AllocatedBlockMap::iterator iter = allocatedBlocks.find(ptr);
    size_t n = iter->second;
    allocatedBlocks.erase(iter);

    // insert the block into the free blocks map
    freeBlocks.push_front(std::make_pair(n, ptr));
    
    // Limit number of free buffers
    while(freeBlocks.size() > maxBuffers) {
      thrust::cuda::free(thrust::cuda::pointer<char>(freeBlocks.back().second));
      freeBlocks.pop_back();
    }
  }

  void CachedAllocator::clear()
  {
    // deallocate all outstanding blocks in both lists
    for(FreeBlockList::iterator i = freeBlocks.begin(); i != freeBlocks.end(); i++)
      // transform the pointer to cuda::pointer before calling cuda::free
      thrust::cuda::free(thrust::cuda::pointer<char>(i->second));
 
    for (AllocatedBlockMap::iterator i = allocatedBlocks.begin(); i != allocatedBlocks.end(); i++) {
      // transform the pointer to cuda::pointer before calling cuda::free
      thrust::cuda::free(thrust::cuda::pointer<char>(i->first));
    }
    freeBlocks.clear();
    allocatedBlocks.clear();
  }
}
#endif
