forward.slang

#include "shared.h"
#include "renderer_types.h"
#include "lib/brdf.slang"
#include "lib/frustum_culling.slang"
#include "lib/shadow_sample.slang"
#include "ibl/ibl_lighting.slang"
#include "lib/cluster_lighting.slang"
#include "debug/debug_views.slang"
#include "lib/ddgi_common.slang"

[vk::push_constant]
FwdDrawPushConstants pc;

[shader("vertex")]
VtxOut vsMain(u32 vertexID: SV_VertexID, u32 drawID: SV_DrawIndex, u32 instanceID: SV_InstanceID) {
    VtxOut o;

    FrameUBO frame = *pc.frame;

    u32 item_idx = pc.visible_items[drawID].id;
    RenderItemGPU item = frame.render_items[item_idx];
    SubmeshGPU submesh = frame.submeshes[item.submesh_index];
    TransformGPU xform = frame.transforms[item.instance_id];
    Vertex v = pc.vertex[vertexID + submesh.base_vertex];
    mat3 world3x3 = (mat3)xform.world;
    mat3 normalMat = (mat3)transpose(xform.normal_world);

    vec3 N = v.normal * normalMat;
    vec3 T = v.tangent.xyz * world3x3;
    vec4 worldPos = vec4(v.pos, 1.0) * xform.world;

    o.view_pos = worldPos * frame.view;
    o.clip_pos = o.view_pos * frame.proj;
    o.world_pos = worldPos.xyz;
    o.clip_debug = o.clip_pos;
    o.normal = N;
    o.tangent = vec4(T, v.tangent.w);
    o.uv = v.uv;
    o.material_index = item.material_index;

    return o;
}

[shader("fragment")]
vec4 fsMain(VtxOut input) : SV_Target {
    FrameUBO frame = *pc.frame;

    MaterialGPU material = frame.materials[input.material_index];
    SamplerState material_sampler = frame.aniso_wrap_mips;

    vec4 albedo = material.albedo.Sample(material_sampler, input.uv) * material.base_color;
    f32 alpha = albedo.a;
    vec3 mrSample = material.metal_rough.Sample(material_sampler, input.uv).rgb;
    f32 metallic = mrSample.b * material.metallic_factor;
    f32 roughness = mrSample.g * material.roughness_factor;
    roughness = max(roughness, 0.045);
    vec3 emissive = material.emissive.Sample(material_sampler, input.uv).rgb * material.emissive_factor;

    vec3 N = normalize(input.normal);
    vec3 T = normalize(input.tangent.xyz);
    T = normalize(T - N * dot(N, T));
    vec3 B = cross(N, T) * input.tangent.w;
    mat3 TBN = mat3(T, B, N);
    vec3 normalTS = material.normal.Sample(material_sampler, input.uv).xyz * 2.0 - 1.0;
    vec3 normalWS = normalize(normalTS * TBN);

#ifdef DEBUG
    vec4 dbg_result;
    if (try_debug_view_pre(
            input, frame, albedo, alpha, normalWS, metallic, roughness, emissive, N, T, B, normalTS, dbg_result
        )) {
        return dbg_result;
    }
#endif
    vec3 worldPos = input.world_pos;
    vec3 cameraPos = frame.position_ws.xyz;
    vec3 V = normalize(cameraPos - input.world_pos);

    vec3 L = normalize(frame.light_dir);
    vec3 H = normalize(V + L);

    f32 NdotL = saturate(dot(normalWS, L));
    f32 NdotV = saturate(dot(normalWS, V));
    f32 NdotH = saturate(dot(normalWS, H));
    f32 VdotH = saturate(dot(V, H));

    // direct light
    vec3 direct_light_radiance = frame.light_color * frame.light_intensity;

    vec3 F0 = lerp(vec3(0.04, 0.04, 0.04), albedo.rgb, metallic);
    f32 D = D_GGX(NdotH, roughness);
    f32 G = G_Smith_Direct(NdotV, NdotL, roughness);
    vec3 F = F_Schlick(VdotH, F0);

    vec3 specular = (D * G * F) / max(4.0 * NdotV * NdotL, 1e-5);
    vec3 kS = F;
    vec3 kD = (1.0 - kS) * (1.0 - metallic);
    vec3 diffuse = kD * albedo.xyz * (1.0 / 3.14159265);

    // direct shadow eval
    f32 shadow_factor = sample_shadow(
        worldPos, normalWS, input.clip_pos.xy, NdotL, frame.position_ws.xyz, frame.shadows, frame.cmp_greater
    );

    // indirect ibl (diffuse + specular)
    vec3 indirect_env_ibl = evaluate_ibl(
        albedo.rgb,
        normalWS,
        V,
        NdotV,
        roughness,
        metallic,
        F0,
        frame.linear_clamp_mips,
        frame.env_irradiance,
        frame.env_prefiltered,
        frame.env_brdf,
        frame.env_prefilter_mip
    );

    // indirect ddgi
    vec3 ddgi_irradiance = sample_ddgi_irradiance(
        worldPos, normalWS, frame.ddgi_grid, frame, pc.ddgi_irradiance, pc.ddgi_distance, pc.ddgi_probe_state, false
    );

    // fade DDGI contribution out near the grid edges and let the environment diffuse fill in
    f32 volume_fade = ddgi_volume_blend_weight(worldPos, frame.ddgi_grid);
    vec3 sky_ambient = frame.env_irradiance.Sample(material_sampler, normalWS).rgb;
    ddgi_irradiance = lerp(sky_ambient, ddgi_irradiance, volume_fade);

    vec3 F_ibl = F_SchlickRoughness(NdotV, F0, roughness);
    vec3 kD_ibl = (1.0 - F_ibl) * (1.0 - metallic);
    vec3 indirect_diffuse = kD_ibl * (ddgi_irradiance) * (albedo.rgb / 3.14159265);

    vec3 direct_lighting = (diffuse + specular) * direct_light_radiance * NdotL * shadow_factor;

    ClusterRecord cr;
    vec3 point_lighting =
        evaluate_cluster_lights(input, worldPos, normalWS, V, NdotV, roughness, metallic, F0, albedo, frame, cr);

    vec3 color = direct_lighting + indirect_diffuse + emissive + point_lighting;

#ifdef DEBUG
    apply_debug_view_post(color, input, cr, frame);
#endif

    return vec4(color, alpha);
}

[shader("fragment")]
vec4 fsWireMain(VtxOut input) : SV_Target {
    return vec4(0.0, 0.85, 1.0, 1.0);
}