Fog Light/Density compte shader

#include <Bindings.glsl>
#include <Camera.glsl>
#include <Fog.glsl>
#include <FrameInfo.glsl>
#include <FwdLights.glsl>
#include <Lights.glsl>
#include <VTFSLightSampling.glsl>

layout(
    local_size_x = FOG_WIDTH / FOG_WORKGROUPS,
    local_size_y = FOG_HEIGHT / FOG_WORKGROUPS,
    local_size_z = FOG_DEPTH / FOG_WORKGROUPS) in;

layout(binding = UBO_FRAME_INFO) uniform FrameInfoBlock
{
    FrameInfo u_FrameInfo;
};
layout(binding = UBO_CAMERA) uniform CameraBlock
{
    Camera u_Camera;
};
layout(binding = UBO_FOG_SETTINGS) uniform FogSettingsBlock
{
    FogSettings u_FogSettings;
};

layout(binding = 0) uniform sampler3D u_Density;
layout(binding = 0, rgba16f) restrict writeonly uniform image3D img_Result;

float HenyeyGreensteinPhase(IN(vec3) a_WorldView, IN(vec3) a_WorldLightDir, float a_Aniso)
{
    float costh = dot(a_WorldView, a_WorldLightDir);
    float denom = 1.0f + a_Aniso * a_Aniso + 2.0f * a_Aniso * costh;
    return (1.0f / (4.0f * M_PI)) * (1.0f - a_Aniso * a_Aniso) / max(pow(denom, 1.5f), 0.000001);
}

float SchlickPhase(IN(vec3) a_WorldView, IN(vec3) a_WorldLightDir, float a_Aniso)
{
    float costh = dot(a_WorldView, a_WorldLightDir);
    return (1.0 - a_Aniso * a_Aniso) / (4.0 * M_PI * pow(1.0 - a_Aniso * costh, 2.0));
}

vec3 GetFogVTFSLightColor(
    IN(vec3) a_WorldCameraPosition,
    IN(vec3) a_WorldPosition,
    IN(vec3) a_NDCPosition,
    IN(float) a_Aniso)
{
    const vec3 worldView          = normalize(a_WorldPosition - a_WorldCameraPosition);
    const uvec3 vtfsClusterIndex  = VTFSClusterIndex(a_NDCPosition);
    const uint vtfsClusterIndex1D = VTFSClusterIndexTo1D(vtfsClusterIndex);
    const uint lightCount         = vtfsClusters[vtfsClusterIndex1D].count;
    vec3 totalLightColor          = vec3(0);
    for (uint i = 0; i < lightCount; i++) {
        const uint lightIndex         = vtfsClusters[vtfsClusterIndex1D].index[i];
        const int lightType           = lightBase[lightIndex].commonData.type;
        const vec3 lightPosition      = lightBase[lightIndex].commonData.position;
        const vec3 lightColor         = lightBase[lightIndex].commonData.color;
        const float lightMaxIntensity = lightBase[lightIndex].commonData.intensity;
        const float lightFalloff      = lightBase[lightIndex].commonData.falloff;
        float lightIntensity          = 0;
        vec3 L                        = vec3(0);
        if (lightType == LIGHT_TYPE_POINT || lightType == LIGHT_TYPE_SPOT) {
            const vec3 LVec   = lightPosition - a_WorldPosition;
            const float LDist = length(LVec);
            L                 = LVec / LDist;
            lightIntensity    = PointLightIntensity(LDist, lightPoint[lightIndex].range, lightMaxIntensity, lightFalloff)
                * HenyeyGreensteinPhase(worldView, L, a_Aniso);
            if (lightType == LIGHT_TYPE_SPOT) {
                const vec3 lightDir             = lightSpot[lightIndex].direction;
                const float lightInnerConeAngle = lightSpot[lightIndex].innerConeAngle;
                const float lightOuterConeAngle = lightSpot[lightIndex].outerConeAngle;
                lightIntensity *= SpotLightIntensity(L, lightDir, lightInnerConeAngle, lightOuterConeAngle);
            }
        } else {
            L              = -lightDirectional[lightIndex].direction;
            lightIntensity = lightMaxIntensity
                * HenyeyGreensteinPhase(worldView, L, a_Aniso);
        }
        totalLightColor += lightColor * lightIntensity;
    }
    return totalLightColor;
}

vec3 GetFogShadowLightColor(
    IN(vec3) a_WorldCameraPosition,
    IN(vec3) a_WorldPosition,
    IN(float) a_Aniso,
    IN(ivec3) a_RandBase)
{
    const vec3 worldView = normalize(a_WorldPosition - a_WorldCameraPosition);
    vec3 totalLightColor = vec3(0);
    for (uint i = 0; i < u_FwdShadowsBase.count; i++) {
        const FwdShadowBase shadowBase = u_FwdShadowsBase.shadows[i];
        const int lightType            = shadowBase.light.commonData.type;
        const vec3 lightPosition       = shadowBase.light.commonData.position;
        const vec3 lightColor          = shadowBase.light.commonData.color;
        const float lightMaxIntensity  = shadowBase.light.commonData.intensity;
        const float lightFalloff       = shadowBase.light.commonData.falloff;
        float lightIntensity           = 0;
        vec3 L                         = vec3(0);
        if (lightType == LIGHT_TYPE_POINT) {
            const FwdShadowPoint shadowPoint = u_FwdShadowsPoint.shadows[i];
            const float lightRange           = shadowPoint.light.range;
            const vec3 LVec                  = lightPosition - a_WorldPosition;
            const float LDist                = length(LVec);
            L                                = normalize(LVec);
            ShadowPointData shadowData;
            shadowData.lightDir         = -L;
            shadowData.lightDist        = LDist;
            shadowData.near             = shadowPoint.projection.zNear;
            shadowData.far              = shadowPoint.projection.zFar;
            shadowData.blurRadius       = shadowPoint.blurRadius;
            shadowData.randBase         = a_RandBase;
            const float shadowIntensity = SampleShadowMap(u_FwdShadowSamplers[i], shadowData);
            lightIntensity              = PointLightIntensity(LDist, lightRange, lightMaxIntensity, lightFalloff)
                * HenyeyGreensteinPhase(worldView, L, a_Aniso)
                * shadowIntensity;
        } else if (lightType == LIGHT_TYPE_SPOT) {
            const FwdShadowSpot shadowSpot  = u_FwdShadowsSpot.shadows[i];
            const float lightRange          = shadowSpot.light.range;
            const vec3 lightDir             = shadowSpot.light.direction;
            const float lightInnerConeAngle = shadowSpot.light.innerConeAngle;
            const float lightOuterConeAngle = shadowSpot.light.outerConeAngle;
            const vec3 LVec                 = lightPosition - a_WorldPosition;
            const float LDist               = length(LVec);
            L                               = normalize(LVec);
            ShadowSpotData shadowData;
            shadowData.projection       = shadowSpot.projection.projection * shadowSpot.projection.view;
            shadowData.near             = shadowSpot.projection.zNear;
            shadowData.far              = shadowSpot.projection.zFar;
            shadowData.blurRadius       = shadowSpot.blurRadius;
            shadowData.surfacePosition  = a_WorldPosition;
            shadowData.randBase         = a_RandBase;
            const float shadowIntensity = SampleShadowMap(u_FwdShadowSamplers[i], shadowData);
            lightIntensity              = PointLightIntensity(LDist, lightRange, lightMaxIntensity, lightFalloff)
                * SpotLightIntensity(L, lightDir, lightInnerConeAngle, lightOuterConeAngle)
                * HenyeyGreensteinPhase(worldView, L, a_Aniso)
                * shadowIntensity;
        } else {
            const FwdShadowDir shadowDir = u_FwdShadowsDir.shadows[i];
            L = -shadowDir.light.direction;
            ShadowDirData shadowData;
            shadowData.surfacePosition  = a_WorldPosition;
            shadowData.projection       = shadowDir.projection.projection * shadowDir.projection.view;
            shadowData.near             = shadowDir.projection.zNear;
            shadowData.far              = shadowDir.projection.zFar;
            shadowData.blurRadius       = shadowDir.blurRadius;
            shadowData.randBase         = a_RandBase;
            const float shadowIntensity = SampleShadowMap(u_FwdShadowSamplers[i], shadowData);
            lightIntensity              = lightMaxIntensity
                * HenyeyGreensteinPhase(worldView, L, a_Aniso)
                * shadowIntensity;
        }
        totalLightColor += lightColor * lightIntensity;
    }
    return totalLightColor;
}

#ifndef ANISO
#define ANISO 0.76
#endif

void main()
{
    const mat4x4 VP       = u_Camera.projection * u_Camera.view;
    const mat4x4 invVP    = inverse(VP);
    const vec3 resultSize = vec3(FOG_WIDTH, FOG_HEIGHT, FOG_DEPTH);

    const vec3 fragCoord = gl_GlobalInvocationID + vec3(0.5f);
    const vec3 uv        = vec3(fragCoord.xy / resultSize.xy, pow(fragCoord.z / resultSize.z, 1 / FOG_DEPTH_EXP));
    const vec3 NDCPos    = uv * 2.f - 1.f;
    const vec4 projPos   = (invVP * vec4(NDCPos, 1));
    const vec3 worldPos  = projPos.xyz / projPos.w;

    const vec3 vtfsColor       = GetFogVTFSLightColor(u_Camera.position, worldPos, NDCPos, ANISO);
    const vec3 shadColor       = GetFogShadowLightColor(u_Camera.position, worldPos, ANISO, ivec3(fragCoord.xy, u_FrameInfo.frameIndex));
    const vec3 iblColor        = GetIBLColor(worldPos);
    const vec4 fogColorDensity = texture(u_Density, uv);
    const vec3 lightColor      = (vtfsColor + shadColor + iblColor) * fogColorDensity.a;

    vec4 out_Color;
    out_Color.rgb = fogColorDensity.rgb * lightColor * fogColorDensity.a;
    out_Color.a   = fogColorDensity.a;
    imageStore(img_Result, ivec3(fragCoord), out_Color);
}