//
// Created by postaron on 25/03/24.
//

#ifndef RAIISAFECUDA_MALLOC_UNMANAGED_HPP
#define RAIISAFECUDA_MALLOC_UNMANAGED_HPP

#include <concepts>
#include <memory>
#include <variant>
#include <cuda_runtime_api.h>

namespace safe_cuda::unmanaged {
    using deviceMallocDestroyType = decltype(&cudaFree);
    using hostMallocDestroyType = decltype(&cudaFreeHost);

    template<typename D>
    concept CudaArrayDestroyer = std::is_same_v<D, deviceMallocDestroyType> || std::is_same_v<D, hostMallocDestroyType>;

    template<typename T, CudaArrayDestroyer D> requires std::integral<T> || std::floating_point<T>
    using safePtrType = std::unique_ptr<T, D>;

    /**
     * \brief It allocates unmanaged memory with cuda runtime API.
     *
     * It can allocate unmanaged memory on device and on Host for pinned memory.
     * \tparam T bare and built-in type.
     * \tparam D destroyer type. It determines if it allocates on device or pinned host.
     * \param byteDataSize
     * \return
     */
    template<typename T, CudaArrayDestroyer D>
    std::variant<safePtrType<T, D>, cudaError_t> cuda_malloc(const std::size_t byteDataSize) noexcept {
        T *ptr_tmp = nullptr;
        if constexpr (std::is_same_v<D, deviceMallocDestroyType>) {
            const cudaError_t error = cudaMalloc(&ptr_tmp, byteDataSize);
            if (error != cudaSuccess) {
                return error;
            }
        } else if constexpr (std::is_same_v<D, hostMallocDestroyType>) {
            const cudaError_t error = cudaMallocHost(&ptr_tmp, byteDataSize);
            if (error != cudaSuccess) {
                return error;
            }
        }
        return safePtrType<T, D>{ ptr_tmp, cudaFree };
    }
}

#endif //RAIISAFECUDA_MALLOC_UNMANAGED_HPP