/*******************************************************
 * Copyright (c) 2016, ArrayFire
 * All rights reserved.
 *
 * This file is distributed under 3-clause BSD license.
 * The complete license agreement can be obtained at:
 * http://arrayfire.com/licenses/BSD-3-Clause
 ********************************************************/

#pragma once

#include <thrust/device_malloc_allocator.h>
#include <thrust/device_vector.h>
#include <memory.hpp>

//Below Class definition is found at the following URL
//http://stackoverflow.com/questions/9007343/mix-custom-memory-managment-and-thrust-in-cuda

namespace cuda
{
template<typename T>
struct ThrustAllocator : thrust::device_malloc_allocator<T>
{
    // shorthand for the name of the base class
    typedef thrust::device_malloc_allocator<T> super_t;

    // get access to some of the base class's typedefs
    // note that because we inherited from device_malloc_allocator,
    // pointer is actually thrust::device_ptr<T>
    typedef typename super_t::pointer   pointer;

    typedef typename super_t::size_type size_type;

    pointer allocate(size_type elements)
    {
        return thrust::device_ptr<T>(memAlloc<T>(elements).release());// delegate to ArrayFire allocator
    }

    void deallocate(pointer p, size_type n)
    {
        memFree<T>(p.get());// delegate to ArrayFire allocator
    }
};
}
