AccelerationStructure.h

#pragma once
 
#include "Common.h"
 
#include "Buffer.h"
#include "Device.h"
 
namespace Mandrill
{
    struct AccelerationStructureBuildInfo {
        VkAccelerationStructureBuildGeometryInfoKHR geometry;
        VkAccelerationStructureBuildSizesInfoKHR size;
        const VkAccelerationStructureBuildRangeInfoKHR* range;
    };
 
    struct BLAS {
        VkAccelerationStructureKHR accelerationStructure;
        VkAccelerationStructureGeometryKHR geometry;
        VkAccelerationStructureBuildRangeInfoKHR buildRange;
 
        AccelerationStructureBuildInfo buildInfo;
    };
 
    // Forward declare Descriptor and Scene
    class Descriptor;
    class Scene;
 
    class AccelerationStructure : public std::enable_shared_from_this<AccelerationStructure>
    {
    public:
        MANDRILL_NON_COPYABLE(AccelerationStructure)
 
        
        MANDRILL_API AccelerationStructure(ptr<Device> pDevice, std::weak_ptr<Scene> wpScene,
                                           VkBuildAccelerationStructureFlagsKHR flags);
 
        MANDRILL_API ~AccelerationStructure();
 
        MANDRILL_API void update(VkBuildAccelerationStructureFlagsKHR flags);
 
        MANDRILL_API VkAccelerationStructureKHR getAccelerationStructure() const
        {
            return mTLAS;
        }
 
        MANDRILL_API VkWriteDescriptorSet getWriteDescriptor(uint32_t binding) const
        {
            static VkWriteDescriptorSetAccelerationStructureKHR as = {
                .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET_ACCELERATION_STRUCTURE_KHR,
                .accelerationStructureCount = 1,
                .pAccelerationStructures = &mTLAS,
            };
 
            VkWriteDescriptorSet write = {
                .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
                .pNext = &as,
                .dstBinding = binding,
                .dstArrayElement = 0,
                .descriptorCount = 1,
                .descriptorType = VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR,
            };
 
            return write;
        }
 
    private:
        MANDRILL_API void createBLASes(VkBuildAccelerationStructureFlagsKHR flags);
 
        MANDRILL_API void createTLAS(VkBuildAccelerationStructureFlagsKHR flags, bool update = false);
 
        ptr<Device> mpDevice;
        std::weak_ptr<Scene> mwpScene; // Use weak pointer to scene so scene can destruct freely
 
        VkAccelerationStructureKHR mTLAS;
        VkAccelerationStructureGeometryKHR mGeometry;
        VkAccelerationStructureBuildRangeInfoKHR mBuildRange;
 
        AccelerationStructureBuildInfo mBuildInfo;
 
        std::vector<BLAS> mBLASes;
 
        ptr<Buffer> mpBLASBuffer;
        ptr<Buffer> mpTLASBuffer;
        ptr<Buffer> mpScratch;
        ptr<Buffer> mpInstances;
    };
} // namespace Mandrill