The Witcher 3 - HLSL Fog Shader (RenderDoc ready)

cbuffer cbuffer12 : register(b12) {
	float4 cb12_v0 : packoffset(c0.x);
	float4 cb12_v1 : packoffset(c1.x);
	float4 cb12_v2 : packoffset(c2.x);
	float4 cb12_v3 : packoffset(c3.x);
	float4 cb12_v4 : packoffset(c4.x);
	float4 cb12_v5 : packoffset(c5.x);
	float4 cb12_v6 : packoffset(c6.x);
	float4 cb12_v7 : packoffset(c7.x);
	float4 cb12_v8 : packoffset(c8.x);
	float4 cb12_v9 : packoffset(c9.x);
	float4 cb12_v10 : packoffset(c10.x);
	float4 cb12_v11 : packoffset(c11.x);
	float4 cb12_v12 : packoffset(c12.x);
	float4 cb12_v13 : packoffset(c13.x);
	float4 cb12_v14 : packoffset(c14.x);
	float4 cb12_v15 : packoffset(c15.x);
	float4 cb12_v16 : packoffset(c16.x);
	float4 cb12_v17 : packoffset(c17.x);
	float4 cb12_v18 : packoffset(c18.x);
	float4 cb12_v19 : packoffset(c19.x);
	float4 cb12_v20 : packoffset(c20.x);
	float4 cb12_v21 : packoffset(c21.x);
	float4 cb12_v22 : packoffset(c22.x);
	float4 cb12_v23 : packoffset(c23.x);
	float4 cb12_v24 : packoffset(c24.x);
	float4 cb12_v25 : packoffset(c25.x);
	float4 cb12_v26 : packoffset(c26.x);
	float4 cb12_v27 : packoffset(c27.x);
	float4 cb12_v28 : packoffset(c28.x);
	float4 cb12_v29 : packoffset(c29.x);
	float4 cb12_v30 : packoffset(c30.x);
	float4 cb12_v31 : packoffset(c31.x);
	float4 cb12_v32 : packoffset(c32.x);
	float4 cb12_v33 : packoffset(c33.x);
	float4 cb12_v34 : packoffset(c34.x);
	float4 cb12_v35 : packoffset(c35.x);
	float4 cb12_v36 : packoffset(c36.x);
	float4 cb12_v37 : packoffset(c37.x);
	float4 cb12_v38 : packoffset(c38.x);
	float4 cb12_v39 : packoffset(c39.x);
	float4 cb12_v40 : packoffset(c40.x);
	float4 cb12_v41 : packoffset(c41.x);
	float4 cb12_v42 : packoffset(c42.x);
	float4 cb12_v43 : packoffset(c43.x);
	float4 cb12_v44 : packoffset(c44.x);
	float4 cb12_v45 : packoffset(c45.x);
	float4 cb12_v46 : packoffset(c46.x);
	float4 cb12_v47 : packoffset(c47.x);
	float4 cb12_v48 : packoffset(c48.x);
	float4 cb12_v49 : packoffset(c49.x);
	float4 cb12_v50 : packoffset(c50.x);
	float4 cb12_v51 : packoffset(c51.x);
	float4 cb12_v52 : packoffset(c52.x);
	float4 cb12_v53 : packoffset(c53.x);
	float4 cb12_v54 : packoffset(c54.x);
	float4 cb12_v55 : packoffset(c55.x);
	float4 cb12_v56 : packoffset(c56.x);
	float4 cb12_v57 : packoffset(c57.x);
	float4 cb12_v58 : packoffset(c58.x);
	float4 cb12_v59 : packoffset(c59.x);
	float4 cb12_v60 : packoffset(c60.x);
	float4 cb12_v61 : packoffset(c61.x);
	float4 cb12_v62 : packoffset(c62.x);
	float4 cb12_v63 : packoffset(c63.x);
	float4 cb12_v64 : packoffset(c64.x);
	float4 cb12_v65 : packoffset(c65.x);
	float4 cb12_v66 : packoffset(c66.x);
	float4 cb12_v67 : packoffset(c67.x);
	float4 cb12_v68 : packoffset(c68.x);
	float4 cb12_v69 : packoffset(c69.x);
	float4 cb12_v70 : packoffset(c70.x);
	float4 cb12_v71 : packoffset(c71.x);
	float4 cb12_v72 : packoffset(c72.x);
	float4 cb12_v73 : packoffset(c73.x);
	float4 cb12_v74 : packoffset(c74.x);
	float4 cb12_v75 : packoffset(c75.x);
	float4 cb12_v76 : packoffset(c76.x);
	float4 cb12_v77 : packoffset(c77.x);
	float4 cb12_v78 : packoffset(c78.x);
	float4 cb12_v79 : packoffset(c79.x);
	float4 cb12_v80 : packoffset(c80.x);
	float4 cb12_v81 : packoffset(c81.x);
	float4 cb12_v82 : packoffset(c82.x);
	float4 cb12_v83 : packoffset(c83.x);
	float4 cb12_v84 : packoffset(c84.x);
	float4 cb12_v85 : packoffset(c85.x);
	float4 cb12_v86 : packoffset(c86.x);
	float4 cb12_v87 : packoffset(c87.x);
	float4 cb12_v88 : packoffset(c88.x);
	float4 cb12_v89 : packoffset(c89.x);
	float4 cb12_v90 : packoffset(c90.x);
	float4 cb12_v91 : packoffset(c91.x);
	float4 cb12_v92 : packoffset(c92.x);
	float4 cb12_v93 : packoffset(c93.x);
	float4 cb12_v94 : packoffset(c94.x);
	float4 cb12_v95 : packoffset(c95.x);
	float4 cb12_v96 : packoffset(c96.x);
	float4 cb12_v97 : packoffset(c97.x);
	float4 cb12_v98 : packoffset(c98.x);
	float4 cb12_v99 : packoffset(c99.x);
	float4 cb12_v100 : packoffset(c100.x);
	float4 cb12_v101 : packoffset(c101.x);
	float4 cb12_v102 : packoffset(c102.x);
	float4 cb12_v103 : packoffset(c103.x);
	float4 cb12_v104 : packoffset(c104.x);
	float4 cb12_v105 : packoffset(c105.x);
	float4 cb12_v106 : packoffset(c106.x);
	float4 cb12_v107 : packoffset(c107.x);
	float4 cb12_v108 : packoffset(c108.x);
	float4 cb12_v109 : packoffset(c109.x);
	float4 cb12_v110 : packoffset(c110.x);
	float4 cb12_v111 : packoffset(c111.x);
	float4 cb12_v112 : packoffset(c112.x);
	float4 cb12_v113 : packoffset(c113.x);
	float4 cb12_v114 : packoffset(c114.x);
	float4 cb12_v115 : packoffset(c115.x);
	float4 cb12_v116 : packoffset(c116.x);
	float4 cb12_v117 : packoffset(c117.x);
	float4 cb12_v118 : packoffset(c118.x);
	float4 cb12_v119 : packoffset(c119.x);
	float4 cb12_v120 : packoffset(c120.x);
	float4 cb12_v121 : packoffset(c121.x);
	float4 cb12_v122 : packoffset(c122.x);
	float4 cb12_v123 : packoffset(c123.x);
	float4 cb12_v124 : packoffset(c124.x);
	float4 cb12_v125 : packoffset(c125.x);
	float4 cb12_v126 : packoffset(c126.x);
	float4 cb12_v127 : packoffset(c127.x);
	float4 cb12_v128 : packoffset(c128.x);
	float4 cb12_v129 : packoffset(c129.x);
	float4 cb12_v130 : packoffset(c130.x);
	float4 cb12_v131 : packoffset(c131.x);
	float4 cb12_v132 : packoffset(c132.x);
	float4 cb12_v133 : packoffset(c133.x);
	float4 cb12_v134 : packoffset(c134.x);
	float4 cb12_v135 : packoffset(c135.x);
	float4 cb12_v136 : packoffset(c136.x);
	float4 cb12_v137 : packoffset(c137.x);
	float4 cb12_v138 : packoffset(c138.x);
	float4 cb12_v139 : packoffset(c139.x);
	float4 cb12_v140 : packoffset(c140.x);
	float4 cb12_v141 : packoffset(c141.x);
	float4 cb12_v142 : packoffset(c142.x);
	float4 cb12_v143 : packoffset(c143.x);
	float4 cb12_v144 : packoffset(c144.x);
	float4 cb12_v145 : packoffset(c145.x);
	float4 cb12_v146 : packoffset(c146.x);
	float4 cb12_v147 : packoffset(c147.x);
	float4 cb12_v148 : packoffset(c148.x);
	float4 cb12_v149 : packoffset(c149.x);
	float4 cb12_v150 : packoffset(c150.x);
	float4 cb12_v151 : packoffset(c151.x);
	float4 cb12_v152 : packoffset(c152.x);
	float4 cb12_v153 : packoffset(c153.x);
	float4 cb12_v154 : packoffset(c154.x);
	float4 cb12_v155 : packoffset(c155.x);
	float4 cb12_v156 : packoffset(c156.x);
	float4 cb12_v157 : packoffset(c157.x);
	float4 cb12_v158 : packoffset(c158.x);
	float4 cb12_v159 : packoffset(c159.x);
	float4 cb12_v160 : packoffset(c160.x);
	float4 cb12_v161 : packoffset(c161.x);
	float4 cb12_v162 : packoffset(c162.x);
	float4 cb12_v163 : packoffset(c163.x);
	float4 cb12_v164 : packoffset(c164.x);
	float4 cb12_v165 : packoffset(c165.x);
	float4 cb12_v166 : packoffset(c166.x);
	float4 cb12_v167 : packoffset(c167.x);
	float4 cb12_v168 : packoffset(c168.x);
	float4 cb12_v169 : packoffset(c169.x);
	float4 cb12_v170 : packoffset(c170.x);
	float4 cb12_v171 : packoffset(c171.x);
	float4 cb12_v172 : packoffset(c172.x);
	float4 cb12_v173 : packoffset(c173.x);
	float4 cb12_v174 : packoffset(c174.x);
	float4 cb12_v175 : packoffset(c175.x);
	float4 cb12_v176 : packoffset(c176.x);
	float4 cb12_v177 : packoffset(c177.x);
	float4 cb12_v178 : packoffset(c178.x);
	float4 cb12_v179 : packoffset(c179.x);
	float4 cb12_v180 : packoffset(c180.x);
	float4 cb12_v181 : packoffset(c181.x);
	float4 cb12_v182 : packoffset(c182.x);
	float4 cb12_v183 : packoffset(c183.x);
	float4 cb12_v184 : packoffset(c184.x);
	float4 cb12_v185 : packoffset(c185.x);
	float4 cb12_v186 : packoffset(c186.x);
	float4 cb12_v187 : packoffset(c187.x);
	float4 cb12_v188 : packoffset(c188.x);
	float4 cb12_v189 : packoffset(c189.x);
	float4 cb12_v190 : packoffset(c190.x);
	float4 cb12_v191 : packoffset(c191.x);
	float4 cb12_v192 : packoffset(c192.x);
	float4 cb12_v193 : packoffset(c193.x);
	float4 cb12_v194 : packoffset(c194.x);
	float4 cb12_v195 : packoffset(c195.x);
	float4 cb12_v196 : packoffset(c196.x);
	float4 cb12_v197 : packoffset(c197.x);
	float4 cb12_v198 : packoffset(c198.x);
	float4 cb12_v199 : packoffset(c199.x);
	float4 cb12_v200 : packoffset(c200.x);
	float4 cb12_v201 : packoffset(c201.x);
	float4 cb12_v202 : packoffset(c202.x);
	float4 cb12_v203 : packoffset(c203.x);
	float4 cb12_v204 : packoffset(c204.x);
	float4 cb12_v205 : packoffset(c205.x);
	float4 cb12_v206 : packoffset(c206.x);
	float4 cb12_v207 : packoffset(c207.x);
	float4 cb12_v208 : packoffset(c208.x);
	float4 cb12_v209 : packoffset(c209.x);
	row_major float4x4 screenToWorld : packoffset(c210.x);
};


cbuffer cbuffer3 : register(b3) {
	float4 cb3_v0 : packoffset(c0.x);
	float4 cb3_v1 : packoffset(c1.x);
};

#define OVERRIDE_FOG

Texture2D texDepth : register (t0);
Texture2D texColor : register (t1);
Texture2D texAO    : register (t2);

SamplerState samplerPointClamp  : register (s0);
SamplerState samplerPointWrap   : register (s1);
SamplerState samplerLinearClamp : register (s2);
SamplerState samplerLinearWrap  : register (s3);
SamplerState samplerAnisoClamp  : register (s4);
SamplerState samplerAnisoWrap   : register (s5);

struct VS_OUTPUT
{
 	float4 PositionH : SV_Position;
	float2 Texcoords : Texcoord0;
};

struct FogResult
{
    float4 paramsFog;
	float4 paramsAerial;
};

float3 getWorldPos( float2 screenPos, float depth )
{
   float4 worldPos = float4(screenPos, depth, 1.0);
   worldPos = mul( worldPos, screenToWorld );

   return worldPos.xyz / worldPos.w;
}

float linstep( float min, float max, float v )
{
    return saturate( (v - min) / (max - min) );
}


float4 SelectDensityParams(const bool isSky)
{
    float4 FogDensityParamsScene = cb12_v43;
    float4 FogDensityParamsSky = cb12_v44;

    if (isSky)
        return FogDensityParamsSky;
    else
        return FogDensityParamsScene;
}

float AdjustAmbientOcclusion(in float inputAO, in float worldToCameraDistance)
{
    // *** Inputs *** //
    const float aoDistanceStart = cb3_v0.x;
    const float aoDistanceEnd = cb3_v0.y;
    const float aoStrengthStart = cb3_v0.z;
    const float aoStrengthEnd = cb3_v0.w;

    // * Adjust AO
    float aoDistanceIntensity = linstep( aoDistanceStart, aoDistanceEnd, worldToCameraDistance );
    float aoStrength = lerp(aoStrengthStart, aoStrengthEnd, aoDistanceIntensity);
    float adjustedAO = lerp(1.0, inputAO, aoStrength);

    return adjustedAO;
}

FogResult CalculateFog(float3 fragPosWorldSpace, float3 customCameraPos, float fogStartDist, float ao, const bool isSky)
{
    // *** Inputs *** //
    float3 FogSunDir = cb12_v38.xyz;
    float3 FogColorFront = cb12_v39.xyz;
    float3 FogColorMiddle = cb12_v40.xyz;
    float3 FogColorBack = cb12_v41.xyz;

    float4 FogBaseParams = cb12_v42;
    float4 FogDensityParamsScene = cb12_v43;
    float4 FogDensityParamsSky = cb12_v44;

    float3 AerialColorFront = cb12_v45.xyz;
    float3 AerialColorMiddle = cb12_v46.xyz;
    float3 AerialColorBack = cb12_v47.xyz;
    float4 AerialParams = cb12_v48;

    float3 GlobalLightDirection = FogSunDir;

    float4 density_params = SelectDensityParams(isSky);


     // *** Shader data *** //
    float density = density_params.x;
    float dist_clamp = FogBaseParams.z;
    float final_exp_fog = FogBaseParams.w;
    float final_exp_aerial = AerialParams.x;

    float3 frag_vec = fragPosWorldSpace.xyz - customCameraPos.xyz;
    float frag_dist = length(frag_vec);

    float3 frag_dir = frag_vec / frag_dist;

    customCameraPos += frag_dir * fogStartDist;
    frag_dist = max(0, frag_dist - fogStartDist);

    frag_dist = (isSky) ? dist_clamp : min(dist_clamp, frag_dist);
    frag_vec = frag_dir * frag_dist;

    const float inv_num_samples = 1.0 / 16;
    float3 frag_step = frag_vec * inv_num_samples;
    float density_sample_scale = frag_dist * density * inv_num_samples;
    float dot_fragDirSunDir = dot(GlobalLightDirection.xyz, frag_dir);

    float density_factor = 1.0;
    {
        float density_shift = FogBaseParams.x;

        float fc_t = dot_fragDirSunDir;
        fc_t = (fc_t + density_shift) / (1.0 + density_shift);
        fc_t = saturate(fc_t);
        density_factor = lerp(density_params.z, density_params.y, fc_t);
    }

    float3 curr_col_fog;
    float3 curr_col_aerial;
    {
        float _dot = dot_fragDirSunDir;

        float _dd = _dot;
        {
            const float _distOffset = -150;
            const float _distRange = 500;
            const float _mul = 1.0 / _distRange;
            const float _bias = _distOffset * _mul;

            _dd = abs(_dd);
            _dd *= _dd;
            _dd *= saturate( frag_dist * _mul + _bias );
        }

        curr_col_fog = lerp( FogColorMiddle.xyz, (_dot>0.0f ? FogColorFront.xyz : FogColorBack.xyz), _dd );
        curr_col_aerial = lerp( AerialColorMiddle.xyz, (_dot>0.0f ? AerialColorFront.xyz : AerialColorBack.xyz), _dd );
    }

    float fog_amount = 1;
    float fog_amount_scale = 0;
    [branch]
    if ( frag_dist >= AerialParams.y )
    {
        float curr_pos_z_base = (customCameraPos.z + FogBaseParams.y) * density_factor;
        float curr_pos_z_step = frag_step.z * density_factor;

        [unroll] for ( int i=16; i>0; --i )
        {
            fog_amount *= 1 - saturate( density_sample_scale / (1 + max( 0.0, curr_pos_z_base + (i) * curr_pos_z_step ) ) );
        }

        fog_amount = 1 - fog_amount;
        fog_amount_scale = saturate( (frag_dist - AerialParams.y) * AerialParams.z );
    }
    //ret.paramsFog.xyz = fog_amount;
    //return ret;

    FogResult ret;

    ret.paramsFog = float4 ( curr_col_fog, fog_amount_scale * pow( abs(fog_amount), final_exp_fog ) );
    ret.paramsAerial = float4 ( curr_col_aerial, fog_amount_scale * pow( abs(fog_amount), final_exp_aerial ) );


#ifdef OVERRIDE_FOG

    // Override
    float fog_influence = ret.paramsFog.w; // r0.w

	float override1ColorScale = cb12_v189.x;
	float override1ColorBias = cb12_v189.y;
	float3 override1Color = cb12_v188.rgb;

	float override1InfluenceScale = cb12_v189.z;
	float override1InfluenceBias = cb12_v189.w;
	float override1Influence = cb12_v188.w;


    float override1ColorAmount = saturate(fog_influence * override1ColorScale + override1ColorBias);
    float override1InfluenceAmount = saturate(fog_influence * override1InfluenceScale + override1InfluenceBias);

    float4 paramsFogOverride;
    paramsFogOverride.rgb = lerp(curr_col_fog, override1Color, override1ColorAmount ); // ***r5.xyz


    float param1 = lerp(1.0, override1Influence, override1InfluenceAmount); // r0.x
    paramsFogOverride.w = saturate(fog_influence * param1); // ** r5.w


    const float extraFogOverride = cb12_v192.x;

    [branch]
    if (extraFogOverride > 0.0)
    {
        float override2ColorScale = cb12_v191.x;
        float override2ColorBias = cb12_v191.y;
        float3 override2Color = cb12_v190.rgb;

        float override2InfluenceScale = cb12_v191.z;
        float override2InfluenceBias = cb12_v191.w;
        float override2Influence = cb12_v190.w;

        float override2ColorAmount = saturate(fog_influence * override2ColorScale + override2ColorBias);
        float override2InfluenceAmount = saturate(fog_influence * override2InfluenceScale + override2InfluenceBias);

        float4 paramsFogOverride2;
        paramsFogOverride2.rgb = lerp(curr_col_fog, override2Color, override2ColorAmount); // r3.xyz

        float  ov_param1 = lerp(1.0, override2Influence, override2InfluenceAmount); // r0.z
        paramsFogOverride2.w = saturate(fog_influence * ov_param1); // r3.w

        paramsFogOverride = lerp(paramsFogOverride, paramsFogOverride2, extraFogOverride);

    }
    ret.paramsFog = paramsFogOverride;

#endif

    return ret;
}


float3 ApplyFog(FogResult fog, float3 color)
{
    const float3 LuminanceFactors = float3(0.333f, 0.555f, 0.222f);

	float3 aerialColor = dot(LuminanceFactors, color) * fog.paramsAerial.xyz;
	color = lerp(color, aerialColor, fog.paramsAerial.w);
	color = lerp(color, fog.paramsFog.xyz, fog.paramsFog.w);
    return color.xyz;

    //return fog.paramsFog.xyz;
}


float4 EditedShaderPS( in VS_OUTPUT Input ) : SV_Target0
{
    // Get coordinates for load
    uint3 ssPosition = uint3( (uint2) Input.PositionH.xy, 0);

    // Load depth from texture
    float hardwareDepth = texDepth.Load(ssPosition).x;

    // Calculate inverted depth
    float depthScale = cb12_v22.x;
    float depthBias  = cb12_v22.y;
    float invDepth = hardwareDepth * depthScale + depthBias;

    // Output color
    float3 outColor = 0.0.xxx;

    [branch]
    if (invDepth < 1.0)
    {
        float3 worldPos = getWorldPos( ssPosition.xy, hardwareDepth );
        float3 CameraPosition = cb12_v0.xyz;

        // Get HDR color
        float3 colorHDR = texColor.Load( ssPosition ).rgb;

        // Get AO and adjust it
        float ao = texAO.Load( ssPosition ).x;
        ao = max(cb3_v1.x, ao);
        ao = AdjustAmbientOcclusion(ao, length(worldPos-CameraPosition));

        // Calculate fog
        float fogStart = cb12_v22.z; // near plane

        FogResult fog = CalculateFog( worldPos, CameraPosition, fogStart, ao, false );

        // Apply AO to influence
        fog.paramsFog.w *= ao;
        fog.paramsAerial.w *= ao;

        // Mix fog with scene color
        outColor = ApplyFog(fog, colorHDR);
    }

    return float4( outColor, 1.0);
}