Program.hpp

#pragma once
 
#include "Common.hpp"
 
#include "graphics/GPUTypes.hpp"
#include "resources/Buffer.hpp"
#include "resources/Texture.hpp"
 
#include <unordered_map>
#include <array>
 
// Forward declarations
VK_DEFINE_HANDLE(VkBuffer)
VK_DEFINE_HANDLE(VkImageView)
VK_DEFINE_HANDLE(VkSampler)
VK_DEFINE_HANDLE(VkShaderModule)
VK_DEFINE_HANDLE(VkPipelineLayout)
VK_DEFINE_HANDLE(VkDescriptorSetLayout)
 
#define UNIFORMS_SET 0
#define SAMPLERS_SET 1
#define IMAGES_SET 2
#define BUFFERS_SET 3
 
class Program {
public:
 
    static uint ALL_MIPS;
 
    enum class Type {
        GRAPHICS, 
        COMPUTE 
    };
 
    struct UniformDef {
 
        enum class Type {
            BOOL, BVEC2, BVEC3, BVEC4,
            INT, IVEC2, IVEC3, IVEC4,
            UINT, UVEC2, UVEC3, UVEC4,
            FLOAT, VEC2, VEC3, VEC4,
            MAT2, MAT3, MAT4,
            OTHER
        };
 
        std::string name; 
        Type type; 
 
        struct Location {
            uint binding; 
            uint offset; 
        };
 
        std::vector<Location> locations; 
 
    };
 
    struct ImageDef {
        std::string name; 
        TextureShape shape; 
        uint binding; 
        uint set; 
        uint count; 
        bool storage; 
    };
 
    struct BufferDef {
        std::vector<UniformDef> members; 
        std::string name; 
        uint binding; 
        uint size; 
        uint set; 
        uint count; 
        bool storage; 
    };
 
    using Uniforms = std::unordered_map<std::string, UniformDef>; 
 
    Program(const std::string & name, const std::string & vertexContent, const std::string & fragmentContent, const std::string & tessControlContent = "", const std::string & tessEvalContent = "");
 
    Program(const std::string & name, const std::string & computeContent);
 
    void reload(const std::string & vertexContent, const std::string & fragmentContent, const std::string & tessControlContent = "", const std::string & tessEvalContent = "");
 
    void reload(const std::string & computeContent);
 
    bool reloaded() const;
    
    bool reloaded(bool absorb);
 
    void use() const;
 
    void clean();
 
    void buffer(const UniformBufferBase& buffer, uint slot);
 
    void buffer(const Buffer& buffer, uint slot);
 
    void bufferArray(const std::vector<const Buffer *> & buffers, uint slot);
 
    void texture(const Texture* texture, uint slot, uint mip = Program::ALL_MIPS);
 
    void texture(const Texture& texture, uint slot, uint mip = Program::ALL_MIPS);
 
    void textureArray(const std::vector<const Texture *> & textures, uint slot, uint mip = Program::ALL_MIPS);
 
    void textures(const std::vector<const Texture *> & textures, size_t slot = 0);
 
    void defaultTexture(uint slot);
 
    void transitionResourcesTo(Program::Type type);
 
    void update();
 
    void uniform(const std::string & name, bool t);
 
    void uniform(const std::string & name, int t);
 
    void uniform(const std::string & name, uint t);
 
    void uniform(const std::string & name, float t);
 
    void uniform(const std::string & name, const glm::vec2 & t);
 
    void uniform(const std::string & name, const glm::vec3 & t);
 
    void uniform(const std::string & name, const glm::vec4 & t);
 
    void uniform(const std::string & name, const glm::ivec2 & t);
 
    void uniform(const std::string & name, const glm::ivec3 & t);
 
    void uniform(const std::string & name, const glm::ivec4 & t);
 
    void uniform(const std::string & name, const glm::mat3 & t);
 
    void uniform(const std::string & name, const glm::mat4 & t);
 
    void getUniform(const std::string & name, bool & t) const;
 
    void getUniform(const std::string & name, int & t) const;
 
    void getUniform(const std::string & name, uint & t) const;
 
    void getUniform(const std::string & name, float & t) const;
 
    void getUniform(const std::string & name, glm::vec2 & t) const;
 
    void getUniform(const std::string & name, glm::vec3 & t) const;
 
    void getUniform(const std::string & name, glm::vec4 & t) const;
 
    void getUniform(const std::string & name, glm::ivec2 & t) const;
 
    void getUniform(const std::string & name, glm::ivec3 & t) const;
 
    void getUniform(const std::string & name, glm::ivec4 & t) const;
 
    void getUniform(const std::string & name, glm::mat3 & t) const;
 
    void getUniform(const std::string & name, glm::mat4 & t) const;
 
    const Uniforms & uniforms() const {
        return _uniforms;
    }
 
    const std::string & name() const {
        return _name;
    }
 
    Program::Type type() const {
        return _type;
    }
 
    const glm::uvec3 & size() const;
 
    Program & operator=(const Program &) = delete;
    
    Program(const Program &) = delete;
    
    Program & operator=(Program &&) = delete;
    
    Program(Program &&) = default;
 
    struct State {
        std::vector<VkDescriptorSetLayout> setLayouts; 
        VkPipelineLayout layout = VK_NULL_HANDLE; 
    };
 
    const State& getState() const {
        return _state;
    }
 
    struct Stage {
        std::vector<ImageDef> images; 
        std::vector<BufferDef> buffers; 
        VkShaderModule module = VK_NULL_HANDLE; 
        glm::uvec3 size = glm::uvec3(0); 
        void reset();
    };
 
    Stage& stage(ShaderType type){
        return _stages[uint(type)];
    }
 
private:
 
    void reflect();
 
 
    void updateUniformMetric() const;
 
    inline const char* retrieveUniform(const UniformDef::Location& location) const {
        return &(_dynamicBuffers.at(location.binding).buffer->data[location.offset]);
    }
 
    inline char* retrieveUniformNonConst(const UniformDef::Location& location) {
        DynamicBufferState& buffState = _dynamicBuffers.at(location.binding);
        buffState.dirty = true;
        _dirtySets[0] = true;
        return &(buffState.buffer->data[location.offset]);
    }
 
    struct DynamicBufferState {
        std::shared_ptr<UniformBuffer<char>> buffer; 
        uint descriptorIndex = 0; 
        bool dirty = true; 
    };
 
    struct TextureState {
        std::string name; 
        TextureShape shape = TextureShape::D2; 
        std::vector<const Texture*> textures = { nullptr }; 
        std::vector<VkImageView> views = { VK_NULL_HANDLE }; 
        uint count = 1; 
        uint mip = 0xFFFF; 
        bool storage = false; 
    };
 
    struct StaticBufferState {
        std::string name; 
        std::vector<VkBuffer> buffers = { VK_NULL_HANDLE}; 
        std::vector<uint> offsets = {0}; 
        uint size = 0; 
        uint count = 1; 
        uint lastSet = 0; 
        bool storage = false; 
    };
 
 
    std::string _name; 
    std::array<Stage, int(ShaderType::COUNT)> _stages; 
    State _state; 
 
    std::unordered_map<std::string, UniformDef> _uniforms; 
 
    std::unordered_map<int, DynamicBufferState> _dynamicBuffers; 
    std::unordered_map<int, TextureState> _textures; 
    std::unordered_map<int, StaticBufferState> _staticBuffers; 
 
    std::array<bool, 4> _dirtySets; 
    std::array<DescriptorSet, 4> _currentSets; 
    std::vector<uint32_t> _currentOffsets; 
 
    bool _reloaded = false; 
    const Type _type; 
 
    friend class GPU; 
};