XCEngine/engine/include/XCEngine/Resources/Shader/Shader.h

#pragma once

#include <XCEngine/Core/Asset/IResource.h>
#include <XCEngine/Core/Containers/Array.h>
#include <XCEngine/Core/Types.h>
#include <XCEngine/Resources/Material/MaterialRenderState.h>
#include <XCEngine/Resources/Shader/ShaderKeywordTypes.h>

namespace XCEngine {
namespace Resources {

enum class ShaderType : Core::uint8 {
    Vertex,
    Fragment,
    Geometry,
    Compute,
    Hull,
    Domain
};

enum class ShaderLanguage : Core::uint8 {
    GLSL,
    HLSL,
    SPIRV
};

enum class ShaderBackend : Core::uint8 {
    Generic = 0,
    D3D12,
    OpenGL,
    Vulkan
};

enum class ShaderPropertyType : Core::uint8 {
    Float = 0,
    Range,
    Int,
    Vector,
    Color,
    Texture2D,
    TextureCube
};

enum class ShaderResourceType : Core::uint8 {
    ConstantBuffer = 0,
    Texture2D,
    TextureCube,
    Sampler,
    StructuredBuffer,
    RawBuffer,
    RWStructuredBuffer,
    RWRawBuffer
};

struct ShaderUniform {
    Containers::String name;
    Core::uint32 location;
    Core::uint32 size;
    Core::uint32 type;
};

struct ShaderAttribute {
    Containers::String name;
    Core::uint32 location;
    Core::uint32 size;
    Core::uint32 type;
};

struct ShaderPassTagEntry {
    Containers::String name;
    Containers::String value;
};

struct ShaderPropertyDesc {
    Containers::String name;
    Containers::String displayName;
    ShaderPropertyType type = ShaderPropertyType::Float;
    Containers::String defaultValue;
    Containers::String semantic;
};

struct ShaderResourceBindingDesc {
    Containers::String name;
    ShaderResourceType type = ShaderResourceType::ConstantBuffer;
    Core::uint32 set = 0;
    Core::uint32 binding = 0;
    Containers::String semantic;
};

struct ShaderBackendCompiledBinary {
    ShaderBackend backend = ShaderBackend::Generic;
    Containers::Array<Core::uint8> payload;
};

struct ShaderStageVariant {
    ShaderType stage = ShaderType::Fragment;
    ShaderLanguage language = ShaderLanguage::GLSL;
    ShaderBackend backend = ShaderBackend::Generic;
    ShaderKeywordSet requiredKeywords;
    Containers::String entryPoint;
    Containers::String profile;
    Containers::String sourceCode;
    Containers::Array<Core::uint8> compiledBinary;
    Containers::Array<ShaderBackendCompiledBinary> backendCompiledBinaries;

    const Containers::Array<Core::uint8>* GetCompiledBinaryForBackend(
        ShaderBackend targetBackend) const;
    void SetCompiledBinaryForBackend(
        ShaderBackend targetBackend,
        const Containers::Array<Core::uint8>& binary);
    void SetCompiledBinaryForBackend(
        ShaderBackend targetBackend,
        Containers::Array<Core::uint8>&& binary);
};

struct ShaderPass {
    Containers::String name;
    bool hasFixedFunctionState = false;
    MaterialRenderState fixedFunctionState;
    Containers::Array<ShaderPassTagEntry> tags;
    Containers::Array<ShaderResourceBindingDesc> resources;
    Containers::Array<ShaderKeywordDeclaration> keywordDeclarations;
    Containers::Array<ShaderStageVariant> variants;
};

class Shader : public IResource {
public:
    Shader();
    virtual ~Shader() override;
    
    ResourceType GetType() const override { return ResourceType::Shader; }
    const Containers::String& GetName() const override { return m_name; }
    const Containers::String& GetPath() const override { return m_path; }
    ResourceGUID GetGUID() const override { return m_guid; }
    bool IsValid() const override { return m_isValid; }
    size_t GetMemorySize() const override { return m_memorySize; }
    void Release() override;
    
    void AddUniform(const ShaderUniform& uniform);
    const Containers::Array<ShaderUniform>& GetUniforms() const { return m_uniforms; }
    
    void AddAttribute(const ShaderAttribute& attribute);
    const Containers::Array<ShaderAttribute>& GetAttributes() const { return m_attributes; }

    void AddProperty(const ShaderPropertyDesc& property);
    void ClearProperties();
    const Containers::Array<ShaderPropertyDesc>& GetProperties() const { return m_properties; }
    const ShaderPropertyDesc* FindProperty(const Containers::String& propertyName) const;
    void SetFallback(const Containers::String& fallback);
    const Containers::String& GetFallback() const { return m_fallback; }

    void AddPass(const ShaderPass& pass);
    void ClearPasses();
    Core::uint32 GetPassCount() const { return static_cast<Core::uint32>(m_passes.Size()); }
    const Containers::Array<ShaderPass>& GetPasses() const { return m_passes; }

    void AddPassVariant(const Containers::String& passName, const ShaderStageVariant& variant);
    void SetPassTag(
        const Containers::String& passName,
        const Containers::String& tagName,
        const Containers::String& tagValue);
    void AddPassResourceBinding(
        const Containers::String& passName,
        const ShaderResourceBindingDesc& binding);
    void AddPassKeywordDeclaration(
        const Containers::String& passName,
        const ShaderKeywordDeclaration& declaration);
    bool HasPass(const Containers::String& passName) const;
    const ShaderPass* FindPass(const Containers::String& passName) const;
    ShaderPass* FindPass(const Containers::String& passName);
    bool PassDeclaresKeyword(
        const Containers::String& passName,
        const Containers::String& keyword) const;
    bool DeclaresKeyword(const Containers::String& keyword) const;
    const ShaderResourceBindingDesc* FindPassResourceBinding(
        const Containers::String& passName,
        const Containers::String& resourceName) const;
    const ShaderStageVariant* FindVariant(
        const Containers::String& passName,
        ShaderType stage,
        ShaderBackend backend = ShaderBackend::Generic,
        const ShaderKeywordSet& enabledKeywords = ShaderKeywordSet()) const;
    
    class IRHIShader* GetRHIResource() const { return m_rhiResource; }
    void SetRHIResource(class IRHIShader* resource);

private:
    ShaderPass& GetOrCreatePass(const Containers::String& passName);
    
    Containers::Array<ShaderUniform> m_uniforms;
    Containers::Array<ShaderAttribute> m_attributes;
    Containers::Array<ShaderPropertyDesc> m_properties;
    Containers::Array<ShaderPass> m_passes;
    Containers::String m_fallback;
    
    class IRHIShader* m_rhiResource = nullptr;
};

} // namespace Resources
} // namespace XCEngine