Scene.h

#pragma once
 
#include "Common.h"
 
#include "AccelerationStructure.h"
#include "Camera.h"
#include "Descriptor.h"
#include "Device.h"
#include "Layout.h"
#include "Sampler.h"
#include "Swapchain.h"
#include "Texture.h"
 
namespace Mandrill
{
    struct Vertex {
        alignas(16) glm::vec3 position; // Position of vertex in 3D space
        alignas(16) glm::vec3 normal;   // Normal vector of vertex
        alignas(16) glm::vec2 texcoord; // Texture coordinates
        alignas(16) glm::vec3 tangent;  // Tangent vector for normal mapping
        alignas(16) glm::vec3 binormal; // Binormal vector for normal mapping
 
        bool operator==(const Vertex& other) const
        {
            return position == other.position && normal == other.normal && texcoord == other.texcoord &&
                   tangent == other.tangent && binormal == other.binormal;
        }
    };
 
    struct Mesh {
        std::vector<Vertex> vertices;
        std::vector<uint32_t> indices;
        uint32_t materialIndex{};
 
        // Device offset are set when uploading to device
        VkDeviceSize deviceVerticesOffset{};
        VkDeviceSize deviceIndicesOffset{};
    };
 
    struct alignas(16) MaterialParams {
        glm::vec3 diffuse;
        float shininess;
        glm::vec3 specular;
        float indexOfRefraction;
        glm::vec3 ambient;
        float opacity;
        glm::vec3 emission;
        uint32_t hasTexture;
    };
 
    struct alignas(16) MaterialDevice {
        MaterialParams params;
        uint32_t diffuseTextureIndex;
        uint32_t specularTextureIndex;
        uint32_t ambientTextureIndex;
        uint32_t emissionTextureIndex;
        uint32_t normalTextureIndex;
    };
 
    struct Material {
        MaterialParams params{};
        MaterialParams* paramsDevice{}; // This is set during Scene::compile()
        VkDeviceSize paramsOffset{};    // This is set during Scene::compile()
 
        std::string diffuseTexturePath;
        std::string specularTexturePath;
        std::string ambientTexturePath;
        std::string emissionTexturePath;
        std::string normalTexturePath;
 
        ptr<Descriptor> pDescriptor;
    };
 
    struct InstanceData {
        uint32_t verticesOffset; // Offset into global vertex buffer
        uint32_t indicesOffset;  // Offset into global index buffer
    };
 
    class Scene; // Forward declare scene so Node can befriend it
    class Pipeline;
 
    class Node
    {
    public:
        MANDRILL_API Node();
 
        MANDRILL_API ~Node();
 
        MANDRILL_API void render(VkCommandBuffer cmd, const ptr<Camera> pCamera, const ptr<const Scene> pScene) const;
 
        MANDRILL_API void addMesh(uint32_t meshIndex)
        {
            mMeshIndices.push_back(meshIndex);
        }
 
        MANDRILL_API void setPipeline(ptr<Pipeline> pPipeline)
        {
            mpPipeline = pPipeline;
        }
 
        MANDRILL_API glm::mat4 getTransform() const
        {
            return mTransform;
        }
 
        MANDRILL_API void setTransform(glm::mat4 transform)
        {
            mTransform = transform;
        }
 
        MANDRILL_API void setVisible(bool visible)
        {
            mVisible = visible;
        }
 
        MANDRILL_API bool getVisible() const
        {
            return mVisible;
        }
 
        MANDRILL_API std::vector<uint32_t>& getMeshIndices()
        {
            return mMeshIndices;
        }
 
    private:
        friend Scene;
 
        ptr<Pipeline> mpPipeline;
 
        std::vector<uint32_t> mMeshIndices;
 
        glm::mat4 mTransform;
        glm::mat4* mpTransformDevice;
        ptr<Descriptor> pDescriptor;
        std::vector<VkDescriptorSet> mDescriptors;
 
        bool mVisible;
 
        std::vector<Node*> mChildren;
    };
 
    class Scene : public std::enable_shared_from_this<Scene>
    {
    public:
        MANDRILL_NON_COPYABLE(Scene)
 
        
        MANDRILL_API Scene(ptr<Device> pDevice, ptr<Swapchain> pSwapchain, bool supportRayTracing = false);
 
        MANDRILL_API ~Scene();
 
        MANDRILL_API void render(VkCommandBuffer cmd, const ptr<Camera> pCamera) const;
 
        MANDRILL_API ptr<Node> addNode();
 
        MANDRILL_API uint32_t addMaterial(Material material);
 
        MANDRILL_API uint32_t addMesh(const std::vector<Vertex> vertices, const std::vector<uint32_t> indices,
                                      uint32_t materialIndex);
 
        MANDRILL_API std::vector<uint32_t> addMeshFromFile(const std::filesystem::path& path,
                                                           const std::filesystem::path& materialPath = "");
 
        MANDRILL_API void compile();
 
        MANDRILL_API void syncToDevice();
 
        MANDRILL_API void updateAccelerationStructure(
            VkBuildAccelerationStructureFlagsKHR flags = VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_TRACE_BIT_KHR);
 
        MANDRILL_API void bindRayTracingDescriptors(VkCommandBuffer cmd, ptr<Camera> pCamera, VkPipelineLayout layout);
 
        MANDRILL_API void setSampler(const ptr<Sampler> pSampler);
 
        MANDRILL_API ptr<Layout> getLayout();
 
        MANDRILL_API std::vector<Node>& getNodes()
        {
            return mNodes;
        }
 
        MANDRILL_API uint32_t getVertexCount() const
        {
            return mVertexCount;
        }
 
        MANDRILL_API uint32_t getIndexCount() const
        {
            return mIndexCount;
        }
 
        MANDRILL_API uint32_t getMeshCount()
        {
            return count(mMeshes);
        }
 
        MANDRILL_API uint32_t getMaterialCount() const
        {
            return count(mMaterials);
        }
 
        MANDRILL_API uint32_t getTextureCount() const
        {
            return count(mTextures);
        }
 
        MANDRILL_API uint32_t getMeshVertexCount(uint32_t meshIndex) const
        {
            return count(mMeshes[meshIndex].vertices);
        }
 
        MANDRILL_API uint32_t getMeshIndexCount(uint32_t meshIndex) const
        {
            return count(mMeshes[meshIndex].indices);
        }
 
        MANDRILL_API VkDeviceAddress getMeshVertexAddress(uint32_t meshIndex) const
        {
            return mpVertexBuffer->getDeviceAddress() + mMeshes[meshIndex].deviceVerticesOffset;
        }
 
        MANDRILL_API VkDeviceAddress getMeshIndexAddress(uint32_t meshIndex) const
        {
            return mpIndexBuffer->getDeviceAddress() + mMeshes[meshIndex].deviceIndicesOffset;
        }
 
        MANDRILL_API uint32_t getMeshMaterialIndex(uint32_t meshIndex) const
        {
            return mMeshes[meshIndex].materialIndex;
        }
 
        MANDRILL_API ptr<AccelerationStructure> getAccelerationStructure() const
        {
            return mpAccelerationStructure;
        }
 
        MANDRILL_API void setEnvironmentMap(ptr<Texture> pTexture)
        {
            mpEnvironmentMap = pTexture;
        }
 
    private:
        friend Node;
 
        std::vector<uint32_t> loadFromOBJ(const std::filesystem::path& path,
                                          const std::filesystem::path& materialPath = "");
        std::vector<uint32_t> loadFromGLTF(const std::filesystem::path& path,
                                           const std::filesystem::path& materialPath = "");
        void addTexture(std::string texturePath);
        void createDescriptors();
 
        ptr<Device> mpDevice;
        ptr<Swapchain> mpSwapchain;
 
        std::vector<Mesh> mMeshes;
        std::vector<Node> mNodes;
        std::vector<Material> mMaterials;
        std::unordered_map<std::string, ptr<Texture>> mTextures;
        ptr<Texture> mpEnvironmentMap;
        ptr<Descriptor> mpEnvironmentMapDescriptor;
 
        ptr<Buffer> mpVertexBuffer;
        ptr<Buffer> mpIndexBuffer;
        ptr<Buffer> mpTransforms;
        ptr<Buffer> mpMaterialParams;
 
        ptr<Texture> mpMissingTexture;
 
        bool mSupportRayTracing;
        ptr<AccelerationStructure> mpAccelerationStructure;
        ptr<Descriptor> mpRayTracingDescriptor;
        ptr<Buffer> mpMaterialBuffer; // Almost same as mpMaterialParams but for ray tracing
        ptr<Buffer> mpInstanceDataBuffer;
 
        uint32_t mVertexCount;
        uint32_t mIndexCount;
    };
}; // namespace Mandrill