1enable wgpu_ray_query;
2
3#define_import_path bevy_solari::brdf
4
5#import bevy_core_pipeline::tonemapping::tonemapping_luminance as luminance
6#import bevy_pbr::lighting::{D_GGX, V_SmithGGXCorrelated, specular_multiscatter}
7#import bevy_pbr::pbr_functions::{calculate_diffuse_color, calculate_F0}
8#import bevy_pbr::utils::{rand_f, sample_cosine_hemisphere}
9#import bevy_render::maths::{PI, orthonormalize}
10#import bevy_solari::sampling::{sample_ggx_vndf, ggx_vndf_pdf, ggx_vndf_sample_invalid}
11#import bevy_solari::scene_bindings::{ResolvedMaterial, MIRROR_ROUGHNESS_THRESHOLD, brdf_dfg_lut, brdf_dfg_lut_sampler}
12
13struct EvaluateAndSampleBrdfResult {
14    wi: vec3<f32>,
15    throughput: vec3<f32>,
16    pdf: f32,
17}
18
19fn evaluate_and_sample_brdf(
20    wo: vec3<f32>,
21    world_normal: vec3<f32>,
22    material: ResolvedMaterial,
23    rng: ptr<function, u32>,
24) -> EvaluateAndSampleBrdfResult {
25    let NdotV = dot(world_normal, wo);
26    if NdotV < 0.0001 { return EvaluateAndSampleBrdfResult(vec3(0.0), vec3(0.0), 0.0); }
27    let F0 = calculate_F0(material.base_color, material.metallic, vec3(material.reflectance));
28    let df = 1.0 - luminance(fresnel(F0, NdotV));
29
30    let diffuse_weight = mix(df, 0.0, material.metallic);
31    let specular_weight = 1.0 - diffuse_weight;
32
33    let TBN = orthonormalize(world_normal);
34    let T = TBN[0];
35    let B = TBN[1];
36    let N = TBN[2];
37
38    let wo_tangent = vec3(dot(wo, T), dot(wo, B), dot(wo, N));
39
40    var wi: vec3<f32>;
41    var wi_tangent: vec3<f32>;
42    let diffuse_selected = rand_f(rng) < diffuse_weight;
43    if diffuse_selected {
44        wi = sample_cosine_hemisphere(world_normal, rng);
45        wi_tangent = vec3(dot(wi, T), dot(wi, B), dot(wi, N));
46    } else {
47        wi_tangent = sample_ggx_vndf(wo_tangent, material.roughness, rng);
48        if ggx_vndf_sample_invalid(wi_tangent) {
49            return EvaluateAndSampleBrdfResult(vec3(0.0), vec3(0.0), 0.0);
50        }
51        wi = wi_tangent.x * T + wi_tangent.y * B + wi_tangent.z * N;
52    }
53
54    let diffuse_pdf = wi_tangent.z / PI;
55    let specular_pdf = ggx_vndf_pdf(wo_tangent, wi_tangent, material.roughness);
56    let pdf = (diffuse_weight * diffuse_pdf) + (specular_weight * specular_pdf);
57
58    var throughput = evaluate_brdf(wo, wi, world_normal, material);
59    if diffuse_selected || material.roughness > MIRROR_ROUGHNESS_THRESHOLD {
60        throughput /= pdf;
61    } else {
62        throughput /= specular_weight;
63    }
64
65    return EvaluateAndSampleBrdfResult(wi, throughput, pdf);
66}
67
68fn evaluate_brdf(
69    wo: vec3<f32>,
70    wi: vec3<f32>,
71    world_normal: vec3<f32>,
72    material: ResolvedMaterial,
73) -> vec3<f32> {
74    return evaluate_diffuse_brdf(wo, wi, world_normal, material) + evaluate_specular_brdf(wo, wi, world_normal, material);
75}
76
77fn evaluate_diffuse_brdf(wo: vec3<f32>, wi: vec3<f32>, world_normal: vec3<f32>, material: ResolvedMaterial) -> vec3<f32> {
78    let diffuse_color = calculate_diffuse_color(material.base_color, material.metallic, 0.0, 0.0) / PI;
79
80    let NdotL = dot(world_normal, wi);
81    let NdotV = dot(world_normal, wo);
82    if NdotL < 0.0001 || NdotV < 0.0001 { return vec3(0.0); }
83    let F0 = calculate_F0(material.base_color, material.metallic, vec3(material.reflectance));
84    let layering = (1.0 - fresnel(F0, NdotL)) * (1.0 - fresnel(F0, NdotV));
85
86    return diffuse_color * layering * NdotL;
87}
88
89fn evaluate_specular_brdf(wo: vec3<f32>, wi: vec3<f32>, world_normal: vec3<f32>, material: ResolvedMaterial) -> vec3<f32> {
90    let H = normalize(wi + wo);
91    let NdotL = dot(world_normal, wi);
92    let NdotH = dot(world_normal, H);
93    let LdotH = dot(wi, H);
94    let NdotV = dot(world_normal, wo);
95    if NdotL < 0.0001 || NdotH < 0.0001 || LdotH < 0.0001 || NdotV < 0.0001 { return vec3(0.0); }
96
97    let F0 = calculate_F0(material.base_color, material.metallic, vec3(material.reflectance));
98    let F = fresnel(F0, LdotH);
99
100    if material.roughness <= MIRROR_ROUGHNESS_THRESHOLD {
101        if abs(NdotH - 1.0) < 0.0001 {
102            return F;
103        } else {
104            return vec3(0.0);
105        }
106    }
107
108    let D = D_GGX(material.roughness, NdotH);
109    let Vs = V_SmithGGXCorrelated(material.roughness, NdotV, NdotL);
110    let F_ab = F_AB(material.perceptual_roughness, NdotV);
111    return specular_multiscatter(D, Vs, F, F0, F_ab, 1.0) * NdotL;
112}
113
114fn fresnel(f0: vec3<f32>, LdotH: f32) -> vec3<f32> {
115    return f0 + (1.0 - f0) * pow(1.0 - LdotH, 5.0);
116}
117
118// Scale/bias approximation
119fn F_AB(perceptual_roughness: f32, NdotV: f32) -> vec2<f32> {
120    return textureSampleLevel(brdf_dfg_lut, brdf_dfg_lut_sampler, vec2<f32>(NdotV, perceptual_roughness), 0.0).rg;
121}