iwa/include/iwa/buffer.hpp

#pragma once

#if !defined(IWA_BUFFER_HPP_INCLUDED)
#define IWA_BUFFER_HPP_INCLUDED

#include <ranges>
#include <mijin/async/coroutine.hpp>
#include <mijin/container/stride_span.hpp>
#include <mijin/container/typeless_buffer.hpp>
#include <mijin/util/flag.hpp>
#include "iwa/device_memory.hpp"
#include "iwa/util/vkutil.hpp"

namespace iwa
{
MIJIN_DEFINE_FLAG(HostVisible);
MIJIN_DEFINE_FLAG(HostCoherent);

struct BufferCreationArgs
{
    vk::BufferCreateFlags flags;
    vk::DeviceSize size;
    vk::BufferUsageFlags usage;
    vk::SharingMode sharingMode = vk::SharingMode::eExclusive;
    std::vector<std::uint32_t> queueFamilyIndices;
};

class Buffer : public Object<Buffer, BaseObject, class Device>, public MixinVulkanObject<vk::Buffer>
{
public:
    ObjectPtr<DeviceMemory> mMemory;
    vk::DeviceSize mBytesSize;
public:
    Buffer(ObjectPtr<class Device> owner, const BufferCreationArgs& args);
    ~Buffer() noexcept override;

    [[nodiscard]] const ObjectPtr<DeviceMemory>& getMemory() const noexcept { return mMemory; }

    void allocateMemory(HostVisible hostVisible = HostVisible::NO, HostCoherent hostCoherent = HostCoherent::YES); // the latter is ignored if the first one is NO
    void bindMemory(ObjectPtr<DeviceMemory> memory, vk::DeviceSize offset = 0);

    mijin::Task<> c_fill(std::uint32_t data, std::size_t bytes = VK_WHOLE_SIZE, std::size_t byteOffset = 0);
    mijin::Task<> c_copyFrom(vk::Buffer srcBuffer, vk::BufferCopy region);
    mijin::Task<> c_upload(const void* data, std::size_t bytes, std::size_t byteOffset = 0);
    mijin::Task<> c_upload(const mijin::TypelessBuffer& data, std::size_t byteOffset = 0);

    template<std::ranges::contiguous_range TRange>
    mijin::Task<> c_upload(const TRange& range, std::size_t byteOffset = 0)
    {
        return c_upload(std::ranges::cdata(range), std::ranges::size(range) * sizeof(std::ranges::range_value_t<TRange>), byteOffset);
    }
};

template<typename T>
struct TypedBufferCreationArgs
{
    vk::BufferCreateFlags flags;
    vk::BufferUsageFlags usage;
    unsigned arraySize = 1;
    vk::DeviceSize stride = sizeof(T);
    vk::SharingMode sharingMode = vk::SharingMode::eExclusive;
    std::vector<std::uint32_t> queueFamilyIndices;
};

template<typename Type>
class TypedBuffer : public Object<TypedBuffer<Type>, Buffer>
{
public:
    using super_t = typename Object<TypedBuffer<Type>, Buffer>::super_t;
protected:
    unsigned mArraySize = 0;
    vk::DeviceSize mStride = 0;
public:
    TypedBuffer(ObjectPtr<class Device> owner, const TypedBufferCreationArgs<Type>& args)
        : super_t(std::move(owner), BufferCreationArgs{
        .flags = args.flags,
        .size = sizeof(Type) + (args.arraySize - 1) * args.stride,
        .usage = args.usage,
        .sharingMode = args.sharingMode,
        .queueFamilyIndices = args.queueFamilyIndices
    }), mArraySize(args.arraySize), mStride(args.stride) {}

    [[nodiscard]] unsigned getArraySize() const noexcept { return mArraySize; }
    [[nodiscard]] vk::DeviceSize getStride() const noexcept { return mStride; }
    [[nodiscard]] mijin::StrideSpan<Type> mapToSpan() const;
};

template<typename T>
struct TypedUniformBufferCreationArgs
{
    vk::BufferCreateFlags flags;
    vk::BufferUsageFlags usage;
    unsigned arraySize = 1;
    vk::SharingMode sharingMode = vk::SharingMode::eExclusive;
    std::vector<std::uint32_t> queueFamilyIndices;
};

template<typename Type>
class TypedUniformBuffer : public Object<TypedUniformBuffer<Type>, TypedBuffer<Type>>
{
public:
    using super_t = typename Object<TypedUniformBuffer<Type>, TypedBuffer<Type>>::super_t;
public:
    TypedUniformBuffer(ObjectPtr<class Device> owner, const TypedUniformBufferCreationArgs<Type>& args)
        : super_t(owner, TypedBufferCreationArgs<Type>{
            .flags = args.flags,
            .usage = args.usage | vk::BufferUsageFlagBits::eUniformBuffer,
            .arraySize = args.arraySize,
            .stride = calcVkUniformStride<Type>(*owner),
            .sharingMode = args.sharingMode,
            .queueFamilyIndices = args.queueFamilyIndices
        }) {}
};

template<typename Type>
using TypedStorageBufferCreationArgs = TypedUniformBufferCreationArgs<Type>;

template<typename Type>
class TypedStorageBuffer : public Object<TypedStorageBuffer<Type>, TypedBuffer<Type>>
{
public:
    using super_t = typename Object<TypedStorageBuffer<Type>, TypedBuffer<Type>>::super_t;
public:
    TypedStorageBuffer(ObjectPtr<class Device> owner, const TypedUniformBufferCreationArgs<Type>& args)
        : super_t(owner, TypedBufferCreationArgs<Type>{
            .flags = args.flags,
            .usage = args.usage | vk::BufferUsageFlagBits::eStorageBuffer,
            .arraySize = args.arraySize,
            .stride = calcVkStorageBufferStride<Type>(*owner),
            .sharingMode = args.sharingMode,
            .queueFamilyIndices = args.queueFamilyIndices
        }) {}
};

template<typename Type>
mijin::StrideSpan<Type> TypedBuffer<Type>::mapToSpan() const
{
    std::byte* rawMemory = static_cast<std::byte*>(this->getOwner()->getVkHandle().mapMemory(
        /* memory = */ *this->getMemory(),
        /* offset = */ 0,
        /* size   = */ VK_WHOLE_SIZE
    ));
    return mijin::StrideSpan<Type>(
        std::bit_cast<Type*>(rawMemory),
        std::bit_cast<Type*>(rawMemory + mArraySize * mStride),
        mStride
    );
}
} // namespace iwa

#endif // !defined(IWA_BUFFER_HPP_INCLUDED)