Prism/Editor/assets/shaders/SSR.glsl

#type vertex
#version 430

layout(location = 0) in vec3 a_Position;
layout(location = 1) in vec2 a_TexCoord;

out vec2 v_TexCoord;

void main()
{
    vec4 position = vec4(a_Position.xy, 0.0, 1.0);
    v_TexCoord = a_TexCoord;
    gl_Position = position;
}

#type fragment
#version 430

layout(location = 0) out vec4 o_Color;

in vec2 v_TexCoord;

uniform sampler2D u_ColorTexture;
uniform sampler2D u_DepthTexture;
uniform sampler2D u_MaterialInfoTexture;

uniform mat4 u_Projection;
uniform mat4 u_InvProjection;

uniform vec2 u_ScreenSize;
uniform float u_CameraNear;
uniform float u_CameraFar;

uniform int u_Steps = 64;
uniform float u_Thickness = 0.5;
uniform float u_MaxDistance = 30.0;
uniform float u_Intensity = 1.0;

float LinearizeDepth(float depth)
{
    float z = depth * 2.0 - 1.0;
    return (2.0 * u_CameraNear * u_CameraFar) / (u_CameraFar + u_CameraNear - z * (u_CameraFar - u_CameraNear));
}

vec3 ReconstructViewPos(vec2 tc, float depth)
{
    float z = depth * 2.0 - 1.0;
    vec4 clipPos = vec4(tc * 2.0 - 1.0, z, 1.0);
    vec4 viewPos4 = u_InvProjection * clipPos;
    return viewPos4.xyz / viewPos4.w;
}

vec3 ReconstructNormalFromDepth(vec2 tc, float depth)
{
    vec2 texelSize = 1.0 / u_ScreenSize;

    vec2 tcL = tc - vec2(texelSize.x, 0.0);
    vec2 tcR = tc + vec2(texelSize.x, 0.0);
    vec2 tcB = tc - vec2(0.0, texelSize.y);
    vec2 tcT = tc + vec2(0.0, texelSize.y);

    float depthL = texture(u_DepthTexture, tcL).r;
    float depthR = texture(u_DepthTexture, tcR).r;
    float depthB = texture(u_DepthTexture, tcB).r;
    float depthT = texture(u_DepthTexture, tcT).r;

    vec3 pL = ReconstructViewPos(tcL, depthL);
    vec3 pR = ReconstructViewPos(tcR, depthR);
    vec3 pB = ReconstructViewPos(tcB, depthB);
    vec3 pT = ReconstructViewPos(tcT, depthT);

    vec3 dx = pR - pL;
    vec3 dy = pT - pB;

    return normalize(cross(dx, dy));
}

void main()
{
    float depth = texture(u_DepthTexture, v_TexCoord).r;

    if (depth >= 1.0)
    {
        o_Color = vec4(0.0);
        return;
    }

    vec4 matInfo = texture(u_MaterialInfoTexture, v_TexCoord);
    float metalness = matInfo.r;
    float roughness = matInfo.g;

    if (roughness > 0.9)
    {
        o_Color = vec4(0.0);
        return;
    }

    vec3 viewPos = ReconstructViewPos(v_TexCoord, depth);
    vec3 normal = ReconstructNormalFromDepth(v_TexCoord, depth);

    vec3 viewDir = normalize(-viewPos);
    float NdotV = max(dot(normal, viewDir), 0.0);

    vec3 reflectDir = reflect(-viewDir, normal);

    if (dot(reflectDir, normal) < 0.0)
    {
        o_Color = vec4(0.0);
        return;
    }

    vec3 rayOrigin = viewPos + normal * 0.02;
    vec3 rayEnd = rayOrigin + reflectDir * u_MaxDistance;

    vec4 p0 = u_Projection * vec4(rayOrigin, 1.0);
    vec4 p1 = u_Projection * vec4(rayEnd, 1.0);

    float k0 = 1.0 / p0.w;
    float k1 = 1.0 / p1.w;

    p0.xyz *= k0;
    p1.xyz *= k1;

    vec3 v0 = rayOrigin * k0;
    vec3 v1 = rayEnd * k1;

    vec2 s0 = p0.xy * 0.5 + 0.5;
    vec2 s1 = p1.xy * 0.5 + 0.5;

    vec2 sDelta = s1 - s0;
    float screenDist = length(sDelta * u_ScreenSize);

    if (screenDist < 1.0)
    {
        o_Color = vec4(0.0);
        return;
    }

    float divisions = min(screenDist, float(u_Steps));
    vec3 dV = (v1 - v0) / divisions;
    float dK = (k1 - k0) / divisions;
    vec2 dS = sDelta / divisions;

    vec3 curV = v0;
    float curK = k0;
    vec2 curS = s0;

    vec2 hitUV = vec2(0.0);
    bool found = false;

    for (int i = 0; i < int(divisions); i++)
    {
        curV += dV;
        curK += dK;
        curS += dS;

        if (curS.x < 0.0 || curS.x > 1.0 || curS.y < 0.0 || curS.y > 1.0)
            break;

        float sampleDepth = texture(u_DepthTexture, curS).r;
        float surfDepth = LinearizeDepth(sampleDepth);
        float rayDepth = -curV.z / curK;

        float depthDiff = rayDepth - surfDepth;

        if (depthDiff > 0.0 && depthDiff < u_Thickness)
        {
            hitUV = curS;
            found = true;
            break;
        }
    }

    if (!found)
    {
        o_Color = vec4(0.0);
        return;
    }

    vec3 hitNormal = ReconstructNormalFromDepth(hitUV, texture(u_DepthTexture, hitUV).r);
    if (dot(hitNormal, -reflectDir) < 0.0)
    {
        o_Color = vec4(0.0);
        return;
    }

    vec3 reflectColor = texture(u_ColorTexture, hitUV).rgb;

    float dist = distance(viewPos, ReconstructViewPos(hitUV, texture(u_DepthTexture, hitUV).r));
    float fadeFactor = 1.0 - smoothstep(u_MaxDistance * 0.5, u_MaxDistance, dist);

    float edgeFade = 1.0;
    edgeFade *= smoothstep(0.0, 0.15, hitUV.x);
    edgeFade *= smoothstep(0.0, 0.15, hitUV.y);
    edgeFade *= smoothstep(0.0, 0.15, 1.0 - hitUV.x);
    edgeFade *= smoothstep(0.0, 0.15, 1.0 - hitUV.y);

    vec3 F0 = mix(vec3(0.04), reflectColor, metalness);
    vec3 Fresnel = F0 + (1.0 - F0) * pow(1.0 - NdotV, 5.0);

    float roughnessFade = 1.0 - roughness * roughness;

    vec3 ssrContrib = reflectColor * Fresnel * fadeFactor * edgeFade * roughnessFade * u_Intensity;
    float ssrAlpha = length(Fresnel) * fadeFactor * edgeFade * roughnessFade;

    o_Color = vec4(ssrContrib, ssrAlpha);
}