IDeviceMemory.h

#pragma once
#include <UnCompute/Backend/IBuffer.h>
#include <UnCompute/Backend/MemoryKindFlags.h>
#include <UnCompute/Base/Flags.h>
#include <UnCompute/Containers/ArraySlice.h>
#include <limits>
 
namespace UN
{
    struct DeviceMemoryDesc
    {
        const char* Name      = nullptr;               
        UInt64 Size           = 0;                     
        MemoryKindFlags Flags = MemoryKindFlags::None; 
        ArraySlice<IDeviceObject*> Objects;            
 
        inline DeviceMemoryDesc() = default;
 
        inline DeviceMemoryDesc(const char* name, MemoryKindFlags flags, UInt64 size, ArraySlice<IDeviceObject*> objects)
            : Name(name)
            , Size(size)
            , Flags(flags)
            , Objects(objects)
        {
        }
    };
 
    class IDeviceMemory : public IDeviceObject
    {
    public:
        using DescriptorType = DeviceMemoryDesc;
 
        inline static constexpr UInt64 WholeSize = std::numeric_limits<UInt64>::max();
 
        [[nodiscard]] virtual const DescriptorType& GetDesc() const = 0;
 
        virtual ResultCode Init(const DescriptorType& desc) = 0;
 
        virtual ResultCode Map(UInt64 byteOffset, UInt64 byteSize, void** ppData) = 0;
 
        virtual void Unmap() = 0;
 
        [[nodiscard]] virtual bool IsCompatible(IDeviceObject* pObject, UInt64 sizeLimit) = 0;
 
        [[nodiscard]] virtual bool IsCompatible(IDeviceObject* pObject) = 0;
    };
 
    class DeviceMemorySlice final
    {
        IDeviceMemory* m_pMemory = nullptr;
        UInt64 m_ByteOffset      = 0;
        UInt64 m_ByteSize        = 0;
 
    public:
        inline DeviceMemorySlice() = default;
 
        inline explicit DeviceMemorySlice(IDeviceMemory* pMemory, UInt64 byteOffset = 0,
                                          UInt64 byteSize = IDeviceMemory::WholeSize)
        {
            m_pMemory    = pMemory;
            m_ByteOffset = byteOffset;
            m_ByteSize   = std::min(byteSize, pMemory->GetDesc().Size - byteOffset);
            UN_Assert(byteSize == IDeviceMemory::WholeSize || byteSize <= pMemory->GetDesc().Size - byteOffset,
                      "Invalid memory slice size");
        }
 
        [[nodiscard]] inline IDeviceMemory* GetDeviceMemory() const
        {
            return m_pMemory;
        }
 
        [[nodiscard]] inline UInt64 GetByteOffset() const
        {
            return m_ByteOffset;
        }
 
        [[nodiscard]] inline UInt64 GetByteSize() const
        {
            return m_ByteSize;
        }
 
        [[nodiscard]] inline ResultCode Map(void** ppData) const
        {
            return m_pMemory->Map(m_ByteOffset, m_ByteSize, ppData);
        }
 
        [[nodiscard]] inline void* Map() const
        {
            void* pResult;
            if (auto result = m_pMemory->Map(m_ByteOffset, m_ByteSize, &pResult); Failed(result))
            {
                UN_Error(false, "Couldn't map memory, IDeviceMemory::Map returned {}", result);
                return nullptr;
            }
 
            return pResult;
        }
 
        [[nodiscard]] inline ResultCode Map(UInt64 byteOffset, UInt64 byteSize, void** ppData) const
        {
            UN_Assert(byteSize == IDeviceMemory::WholeSize || byteSize <= m_pMemory->GetDesc().Size - byteOffset,
                      "Invalid memory map size");
 
            return m_pMemory->Map(
                m_ByteOffset + byteOffset, std::min(byteSize, m_pMemory->GetDesc().Size - byteOffset - m_ByteOffset), ppData);
        }
 
        [[nodiscard]] inline void* Map(UInt64 byteOffset, UInt64 byteSize = IDeviceMemory::WholeSize) const
        {
            UN_Assert(byteSize == IDeviceMemory::WholeSize || byteSize <= m_pMemory->GetDesc().Size - byteOffset,
                      "Invalid memory map size");
 
            void* pResult;
            if (Succeeded(m_pMemory->Map(m_ByteOffset + byteOffset,
                                         std::min(byteSize, m_pMemory->GetDesc().Size - byteOffset - m_ByteOffset),
                                         &pResult)))
            {
                return pResult;
            }
 
            UN_Error(false, "Couldn't map memory");
            return nullptr;
        }
 
        inline void Unmap() const
        {
            m_pMemory->Unmap();
        }
 
        [[nodiscard]] inline bool IsCompatible(IDeviceObject* pObject) const
        {
            return m_pMemory->IsCompatible(pObject, m_ByteSize);
        }
    };
 
    template<class T>
    class MemoryMapHelper
    {
        DeviceMemorySlice m_MemorySlice;
        T* m_Map = nullptr;
 
        inline MemoryMapHelper() = default;
 
        inline explicit MemoryMapHelper(const DeviceMemorySlice& memorySlice, T* map)
            : m_MemorySlice(memorySlice)
            , m_Map(map)
        {
        }
 
    public:
        inline MemoryMapHelper(MemoryMapHelper&& other) noexcept
            : m_MemorySlice(other.m_MemorySlice)
            , m_Map(other.m_Map)
        {
            other.m_MemorySlice = {};
            other.m_Map         = nullptr;
        }
 
        inline ~MemoryMapHelper()
        {
            if (m_Map)
            {
                m_MemorySlice.Unmap();
            }
        }
 
        inline T& operator[](USize index)
        {
            UN_Assert(IsValid(), "MemoryMapHelper was in invalid state");
            UN_Assert(index * sizeof(T) < m_MemorySlice.GetByteSize(), "Index out of range in MemoryMapHelper");
            return m_Map[index];
        }
 
        [[nodiscard]] inline UInt64 Length() const
        {
            UN_Assert(IsValid(), "MemoryMapHelper was in invalid state");
            return m_MemorySlice.GetByteSize() / sizeof(T);
        }
 
        [[nodiscard]] inline UInt64 ByteSize() const
        {
            UN_Assert(IsValid(), "MemoryMapHelper was in invalid state");
            return m_MemorySlice.GetByteSize();
        }
 
        [[nodiscard]] inline bool IsValid() const
        {
            return m_Map;
        }
 
        inline explicit operator bool() const
        {
            return IsValid();
        }
 
        [[nodiscard]] inline T* Data() const
        {
            UN_Assert(IsValid(), "MemoryMapHelper was in invalid state");
            return m_Map;
        }
 
        inline static MemoryMapHelper Map(const DeviceMemorySlice& memorySlice)
        {
            void* map;
            auto result = memorySlice.Map(&map);
            if (Succeeded(result))
            {
                return MemoryMapHelper(memorySlice, static_cast<T*>(map));
            }
 
            return {};
        }
 
        inline static MemoryMapHelper Map(IDeviceMemory* pMemory, UInt64 byteOffset = 0,
                                          UInt64 byteSize = IDeviceMemory::WholeSize)
        {
            return Map(DeviceMemorySlice(pMemory, byteOffset, byteSize));
        }
 
        inline T* begin()
        {
            UN_Assert(IsValid(), "MemoryMapHelper was in invalid state");
            return m_Map;
        }
 
        inline T* end()
        {
            UN_Assert(IsValid(), "MemoryMapHelper was in invalid state");
            return m_Map + Length();
        }
    };
} // namespace UN