AssetManager.h

Go to the documentation of this file.

#pragma once
#include "DescriptorIndexAllocator.h"
#include "Engine/Core/Settings.h"
#include "Engine/Ecs/FrameProcessing.h"
#include "Engine/Logging/FunctionNotImplemented.h"
#include "Engine/Logging/TracyMacros.hpp"
#include "Engine/Renderer/VulkanBuffer.h"
#include "Engine/Renderer/VulkanStagedBuffer.h"
#include "Engine/Texture/TextureAsset.h"
#include "Engine/Texture/TextureHandleRegistry.h"
 
#include <Engine/Math/BidirectionalMap.h>
#include <Engine/Mesh/Mesh.h>
#include <Engine/Texture/Texture.h>
#include <Engine/Texture/TextureLoadData.h>
#include <filesystem>
#include <future>
#include <memory>
#include <optional>
#include <string>
#include <tiny_gltf.h>
#include <unordered_map>
 
namespace EngineCore
{
    class MeshAsset;
    class MeshAssetManager;
    class ModelAssetManager;
    class RenderingDataManager;
    class TextureAsset;
    struct UnifiedMeshlet;
    struct MeshletBounds;
    class RenderProcess;
    class Renderer;
}  // namespace EngineCore
 
namespace EngineCore
{
    class Engine;
 
    struct PathHasher
    {
        std::size_t operator()( const std::filesystem::path & p ) const noexcept
        {
            return std::hash<std::string>{}( p.string() );  // Hash the string representation
        }
    };
 
    struct PackedVertex
    {
        glm::vec3 position;        // 12 bytes (Keep full precision)
        uint32_t  packedNormal;    //  4 bytes (A2B10G10R10_SNORM)
        uint32_t  packedColor;     //  4 bytes (R8G8B8A8_UNORM)
        uint32_t  packedTexCoord;  //  4 bytes (R16G16_SFLOAT)
        // Total: 24 bytes
    };
 
    // Pack vec3 normal [-1, 1] into A2B10G10R10_SNORM uint32_t
    inline uint32_t packNormalA2B10G10R10_SNORM( const glm::vec3 & n )
    {
        // Clamp input just in case
        float x = std::max( -1.0f, std::min( 1.0f, n.x ) );
        float y = std::max( -1.0f, std::min( 1.0f, n.y ) );
        float z = std::max( -1.0f, std::min( 1.0f, n.z ) );
 
        // Scale to signed 10-bit range [-511, 511]
        int32_t ix = static_cast<int32_t>( std::round( x * 511.0f ) );
        int32_t iy = static_cast<int32_t>( std::round( y * 511.0f ) );
        int32_t iz = static_cast<int32_t>( std::round( z * 511.0f ) );
 
        // Pack into uint32_t (masking handles two's complement correctly)
        // Layout: | 2 unused | 10 bits B | 10 bits G | 10 bits R |
        return ( static_cast<uint32_t>( iz & 0x3FF ) << 20 ) | ( static_cast<uint32_t>( iy & 0x3FF ) << 10 ) |
               ( static_cast<uint32_t>( ix & 0x3FF ) );
        // Alpha bits (highest 2) remain 0
    }
 
    // Pack vec3 color [0, 1] into R8G8B8A8_UNORM uint32_t (Alpha = 1.0)
    inline uint32_t packColorR8G8B8A8_UNORM( const glm::vec3 & c )
    {
        // Clamp input
        float r = std::max( 0.0f, std::min( 1.0f, c.x ) );
        float g = std::max( 0.0f, std::min( 1.0f, c.y ) );
        float b = std::max( 0.0f, std::min( 1.0f, c.z ) );
 
        // Scale to [0, 255]
        uint32_t ur = static_cast<uint32_t>( std::round( r * 255.0f ) );
        uint32_t ug = static_cast<uint32_t>( std::round( g * 255.0f ) );
        uint32_t ub = static_cast<uint32_t>( std::round( b * 255.0f ) );
        uint32_t ua = 255;  // Fully opaque alpha
 
        // Pack (assuming typical RGBA little-endian layout)
        return ( ua << 24 ) | ( ub << 16 ) | ( ug << 8 ) | ur;
    }
 
    // Pack vec2 texCoord into R16G16_SFLOAT uint32_t
    inline uint32_t packTexCoordR16G16_SFLOAT( const glm::vec2 & tc )
    {
        // Use glm's packing function (requires <glm/gtc/packing.hpp>)
        // It converts two floats to two 16-bit half-floats and packs them.
        return glm::packHalf2x16( tc );
    }
 
    struct TextureStorage
    {
        uint32_t                 index   = 0u;
        std::shared_ptr<Texture> texture = nullptr;
 
        template <typename T>
        bool operator>( const T & other )
        {
            return index > other.index;
        }
 
        template <typename T>
        bool operator<( const T & other )
        {
            return index < other.index;
        }
    };
 
    class AssetManager
    {
        friend ApplicationContext;
        friend RenderingDataManager;
 
    public:
        AssetManager();
        AssetManager( ApplicationContext * context );
        ~AssetManager();

        void attachToRenderer( Renderer * renderer );

        void setRenderingDataManager( RenderingDataManager * renderingDataManager );

        void loadEcsModel( const std::filesystem::path & path );

        MeshAsset * getMeshAsset( const Asset::Path & asset );

        void loadEcsTexture( const std::filesystem::path & path );

        [[nodiscard]] uint32_t getTextureDescriptorIndex( const std::filesystem::path & path );

        [[nodiscard]] std::vector<const Mesh *> getAllMeshes() const;
        [[nodiscard]] uint32_t                  getTotalPrimitiveCount() const;

        bool doesMeshAlreadyExist( const std::filesystem::path & meshPath ) const;
 
        uint32_t getImageCount() const;
 
        void logAllLoadedAssets() const;
        void reloadMeshOffsets();
 
        void cleanup();
 
 
 
        std::vector<VkDescriptorImageInfo> getTextureDescriptorInfos() const;
 
        uint32_t                      getShaderCount() const;
        std::vector<MaterialShader *> getShaders() const;
 
        PipelineNames materialNameFromString( const std::string & string );
 
        MaterialShader * getShaderByName( const std::string & shaderName );
        MaterialShader * getShaderByName( const PipelineNames & shaderName );
 
        std::unordered_map<Mesh *, uint32_t> getMeshOffsets() const
        {
            return meshOffsets;
        }
 
        std::vector<Texture *> getTexturesToUpload();
        uint32_t               getMeshOffset( Mesh * mesh ) const;
 
        MaterialShader * registerShader( MaterialShader shader, const PipelineNames & materialName );
        void             unregisterShader( MaterialShader * shader );
 
        VulkanBuffer & getMaterialStorageBufferByName( PipelineNames name );
 
        Ecs::TextureAssetPipeline & getTextureAssetPipeline();
        Ecs::ModelAssetPipeline &   getMeshAssetPipeline();
        void                        unloadAllData() const;
 
        [[nodiscard]] MeshAsset * getMeshAsset( const Asset::Path & asset ) const;

        [[nodiscard]] MaterialAssetManager * getMaterialAssetManager() const
        {
            return materialAssetManager;
        }

        [[nodiscard]] ModelAssetManager * getModelAssetManager() const
        {
            return modelAssetManager;
        }

        [[nodiscard]] MeshAssetManager * getMeshAssetManager() const
        {
            return meshAssetManager;
        }

        [[nodiscard]] TextureHandleRegistry * getTextureHandleRegistry()
        {
            return &textureHandleRegistry_;
        }
 
    private:
        NamedThreadPool * threadPool          = new NamedThreadPool( 4, "Asset Loader" );
        NamedThreadPool * threadedCalculation = new NamedThreadPool( 10, "Bounding Sphere" );
 
        TextureAssetManager *     textureAssetManager;
        Ecs::TextureAssetPipeline texturePipeline;
        MeshAssetManager *        meshAssetManager;
        MaterialAssetManager *    materialAssetManager;
        ModelAssetManager *       modelAssetManager;
 
        // IMPORTANT: textureHandleRegistry_ must be declared BEFORE modelAssetPipeline
        // because modelAssetPipeline's constructor receives 'this' and may access the registry.
        // C++ initializes members in declaration order, not initializer list order.
        TextureHandleRegistry textureHandleRegistry_;  
 
        Ecs::ModelAssetPipeline   modelAssetPipeline;
 
        std::vector<CpuMeshData> cpuMeshData;
 
        std::vector<UnifiedMeshlet>    meshlets;
        std::vector<MeshletBounds>     meshletBounds;
        std::vector<ObjectCullingData> objectCullingData;
        std::vector<PrimitiveMeshletData> objectMeshletData;
 
        std::vector<const MeshPrimitive *>
            getMeshPrimitivesFromMeshes( const std::vector<const Mesh *> & meshes ) const;
        std::vector<MeshPrimitive *> getMeshPrimitivesFromMeshes( const std::vector<Mesh *> & meshes );
        uint32_t                     getMeshletCount( const std::vector<const Mesh *> & meshes ) const;
 
    public:
        NamedThreadPool * getCalculationPool() const
        {
            return threadedCalculation;
        }
 
        NamedThreadPool * getAssetLoaderPool() const
        {
            return threadPool;
        }
 
    private:
        GltfLoader gltfLoader = GltfLoader( threadPool );
 
        template <typename T>
        void initMaterialStorageBuffer( PipelineNames name, uint32_t materialCount );

        template <typename T>
        uint32_t addMaterialData( T data );

        template <typename MaterialDataType>
        class MaterialData
        {
        public:
            explicit MaterialData( uint32_t materialCount );
 
            [[nodiscard]] std::vector<MaterialDataType>                   getData() const;
            [[nodiscard]] const DescriptorIndexAllocator & getIndexAllocator() const;
 
            uint32_t addMaterialData( MaterialDataType materialData );
 
            [[nodiscard]] size_t getMaterialCount() const;
 
            void removeFromIndex( uint32_t materialIndex );
 
            [[nodiscard]] size_t getCapacity() const;
 
        private:
            std::vector<MaterialDataType>           data{};
            DescriptorIndexAllocator allocator{ 0 };
        };
 
        std::unordered_map<PipelineNames, VulkanBuffer> materialStorageBuffers;
 
        DescriptorIndexAllocator                            shaderAllocator = DescriptorIndexAllocator( 100 );
        std::array<std::optional<MaterialShader>, 100>      shaderData;
        std::vector<MaterialShader *>                       shaders;
        std::unordered_map<PipelineNames, MaterialShader *> shadersByName;
 
        std::unordered_map<MaterialShader *, uint32_t> shaderOffsets;
 
        ApplicationContext * context  = nullptr;
        Renderer *           renderer = nullptr;
 
        std::unordered_map<std::filesystem::path, Texture *, PathHasher> textures = {};
        BidirectionalMap<uint32_t, Texture *>                            textureIndexMap;
        BidirectionalMap<uint32_t, Mesh *>                               meshIndexMap;
 
        std::array<std::optional<Texture>, MAX_TEXTURE_COUNT> textureData;
        DescriptorIndexAllocator                              textureAllocator;

        bool doesTextureAlreadyExist( const std::filesystem::path & path );

        Texture * registerTexture( const std::filesystem::path & path, Texture texture );
        void      unregisterTexture( Texture * texture );
 
        TRACY_LOCKABLE( std::mutex, textureMutex, "Texture mutex" );
        TRACY_LOCKABLE( std::mutex, meshesMutex, "Meshes mutex" );
 
        Mesh * registerMesh( const std::filesystem::path & path, GltfLoader::GltfMeshData & gltfMeshData );
        void   unregisterMesh( Mesh * mesh );
 
        std::array<std::optional<Mesh>, MAX_MESH_COUNT>               meshData{};
        DescriptorIndexAllocator                                      meshAllocator;
        std::unordered_map<std::filesystem::path, Mesh *, PathHasher> meshes = {};
 
        bool                                 bIsMeshDirty = true;
        std::unordered_map<Mesh *, uint32_t> meshOffsets;
 
        std::vector<TextureStorage> texturesToUpload{};
        std::vector<Texture *>      texturesToCopyImageData;
    };
 
    template <typename T>
    void AssetManager::initMaterialStorageBuffer( const PipelineNames name, const uint32_t materialCount )
    {
        new ( &materialStorageBuffers[name] ) VulkanBuffer( context );
 
        const size_t bufferSize = sizeof( T ) * materialCount;
 
        materialStorageBuffers[name].overrideSize( bufferSize );
        materialStorageBuffers[name].create(
            bufferSize,
            VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
            VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT
        );
        materialStorageBuffers[name].setDebugName( MaterialShader::getPipelineNameFromString( name ) + " Buffer" );
 
        PLOGI << "Created material storage buffer with size " << bufferSize << " bytes and a count of "
              << materialCount;
    }
 
    template <typename T>
    uint32_t AssetManager::addMaterialData( T data )
    {
        FUNCTION_DEPRECATED_RETURN
        /*constexpr PipelineNames materialName = MaterialNameTrait<T>::value;
 
        if constexpr ( materialName == NORMALS_SHADER )
        {
            return normalMaterialData.addMaterialData( data );
        }
        if constexpr ( materialName == DIFFUSE_FLAT_COLOR )
        {
            return diffuseFlatColorMaterialData.addMaterialData( data );
        }
        // TODO: Add rest of all material data
        throw std::runtime_error( "Material not yet implemented!" );*/
    }
 
    template <typename T>
    AssetManager::MaterialData<T>::MaterialData( uint32_t materialCount )
      : data( materialCount ), allocator( materialCount )
    {
    }
 
    template <typename T>
    std::vector<T> AssetManager::MaterialData<T>::getData() const
    {
        return data;
    }
 
    template <typename T>
    const DescriptorIndexAllocator & AssetManager::MaterialData<T>::getIndexAllocator() const
    {
        return allocator;
    }
 
    template <typename T>
    uint32_t AssetManager::MaterialData<T>::addMaterialData( T materialData )
    {
        uint32_t freeIndex = allocator.allocate();
        data[freeIndex]    = materialData;
        return freeIndex;
    }
 
    template <typename T>
    size_t AssetManager::MaterialData<T>::getMaterialCount() const
    {
        return allocator.getUsedIndices().size();
    }
 
    template <typename T>
    void AssetManager::MaterialData<T>::removeFromIndex( uint32_t materialIndex )
    {
        allocator.free( materialIndex );
        data[materialIndex] = {};
    }
 
    template <typename T>
    size_t AssetManager::MaterialData<T>::getCapacity() const
    {
        return allocator.getCapacity();
    }
}  // namespace EngineCore