1import bevy_core_pipeline::tonemapping::tonemapping_luminance as luminance;
2import bevy_pbr::render::pbr_lighting::{D_GGX, V_SmithGGXCorrelated, specular_multiscatter};
3import bevy_pbr::render::pbr_functions::calculate_F0_dielectric;
4import bevy_pbr::render::utils::{rand_f, sample_cosine_hemisphere};
5import bevy_render::maths::{PI, orthonormalize};
6import package::scene::sampling::{sample_ggx_vndf, ggx_vndf_pdf, ggx_vndf_sample_invalid};
7import package::scene::bindings::{ResolvedMaterial, MIRROR_ROUGHNESS_THRESHOLD, brdf_dfg_lut, brdf_dfg_lut_sampler};
8
9enable wgpu_ray_query;
10
11struct EvaluateAndSampleBrdfResult {
12    wi: vec3<f32>,
13    throughput: vec3<f32>,
14    pdf: f32,
15    diffuse_selected: bool,
16}
17
18struct LobeReflectances {
19    specular: vec3<f32>,
20    diffuse: vec3<f32>,
21}
22
23/// Hemispherical reflectance of each lobe
24fn lobe_reflectances(F0_metal: vec3<f32>, F0_dielectric: vec3<f32>, material: ResolvedMaterial, F_ab: vec2<f32>) -> LobeReflectances {
25    let multiscattering_factor = 1.0 / (F_ab.x + F_ab.y) - 1.0;
26    let rho_specular_metallic = (F0_metal * F_ab.x + F_ab.y) * (1.0 + F0_metal * multiscattering_factor);
27    let rho_specular_dielectric = (F0_dielectric * F_ab.x + F_ab.y) * (1.0 + F0_dielectric * multiscattering_factor);
28    return LobeReflectances(
29        mix(rho_specular_dielectric, rho_specular_metallic, material.metallic),
30        (1.0 - material.metallic) * (1.0 - rho_specular_dielectric) * material.base_color,
31    );
32}
33
34fn specular_lobe_sampling_probability(rho: LobeReflectances, perceptual_roughness: f32) -> f32 {
35    let specular_luminance = luminance(rho.specular);
36    let total_luminance = specular_luminance + luminance(rho.diffuse);
37    let energy_probability = specular_luminance / max(total_luminance, 0.0001);
38
39    // Give smooth dielectric reflections more specular samples to prevent undersampling
40    let roughness_biased_probability = mix(0.5, energy_probability, perceptual_roughness);
41    return max(energy_probability, roughness_biased_probability);
42}
43
44fn evaluate_and_sample_brdf(
45    wo: vec3<f32>,
46    world_normal: vec3<f32>,
47    material: ResolvedMaterial,
48    F_ab: vec2<f32>,
49    rng: ptr<function, u32>,
50) -> EvaluateAndSampleBrdfResult {
51    let NdotV = dot(world_normal, wo);
52    if NdotV < 0.0001 { return EvaluateAndSampleBrdfResult(vec3(0.0), vec3(0.0), 0.0, false); }
53    let F0_metal = material.base_color;
54    let F0_dielectric = calculate_F0_dielectric(vec3(material.reflectance));
55    let rho = lobe_reflectances(F0_metal, F0_dielectric, material, F_ab);
56    let specular_weight = specular_lobe_sampling_probability(rho, material.perceptual_roughness);
57    let diffuse_weight = 1.0 - specular_weight;
58
59    let TBN = orthonormalize(world_normal);
60    let T = TBN[0];
61    let B = TBN[1];
62    let N = TBN[2];
63
64    let wo_tangent = vec3(dot(wo, T), dot(wo, B), dot(wo, N));
65
66    var wi: vec3<f32>;
67    var wi_tangent: vec3<f32>;
68    let diffuse_selected = rand_f(rng) < diffuse_weight;
69    if diffuse_selected {
70        wi = sample_cosine_hemisphere(world_normal, rng);
71        wi_tangent = vec3(dot(wi, T), dot(wi, B), dot(wi, N));
72    } else {
73        wi_tangent = sample_ggx_vndf(wo_tangent, material.roughness, rng);
74        if ggx_vndf_sample_invalid(wi_tangent) {
75            return EvaluateAndSampleBrdfResult(vec3(0.0), vec3(0.0), 0.0, false);
76        }
77        wi = wi_tangent.x * T + wi_tangent.y * B + wi_tangent.z * N;
78
79        // Mirror specular is a delta function
80        if material.roughness <= MIRROR_ROUGHNESS_THRESHOLD {
81            return EvaluateAndSampleBrdfResult(
82                wi,
83                evaluate_specular_brdf(wo, wi, world_normal, material, F_ab) / specular_weight,
84                bitcast<f32>(0x7F800000u), // INF
85                false,
86            );
87        }
88    }
89
90    let diffuse_pdf = wi_tangent.z / PI;
91    let specular_pdf = ggx_vndf_pdf(wo_tangent, wi_tangent, material.roughness);
92    let pdf = (diffuse_weight * diffuse_pdf) + (specular_weight * specular_pdf);
93    let throughput = evaluate_brdf(wo, wi, world_normal, material, F_ab) / pdf;
94    return EvaluateAndSampleBrdfResult(wi, throughput, pdf, diffuse_selected);
95}
96
97fn evaluate_brdf(
98    wo: vec3<f32>,
99    wi: vec3<f32>,
100    world_normal: vec3<f32>,
101    material: ResolvedMaterial,
102    F_ab: vec2<f32>,
103) -> vec3<f32> {
104    return evaluate_diffuse_brdf(wo, wi, world_normal, material, F_ab) + evaluate_specular_brdf(wo, wi, world_normal, material, F_ab);
105}
106
107fn evaluate_diffuse_brdf(wo: vec3<f32>, wi: vec3<f32>, world_normal: vec3<f32>, material: ResolvedMaterial, F_ab: vec2<f32>) -> vec3<f32> {
108    let NdotL = dot(world_normal, wi);
109    let NdotV = dot(world_normal, wo);
110    if NdotL < 0.0001 || NdotV < 0.0001 { return vec3(0.0); }
111    let F0_metal = material.base_color;
112    let F0_dielectric = calculate_F0_dielectric(vec3(material.reflectance));
113    let rho = lobe_reflectances(F0_metal, F0_dielectric, material, F_ab);
114    return rho.diffuse / PI * NdotL;
115}
116
117fn evaluate_specular_brdf(wo: vec3<f32>, wi: vec3<f32>, world_normal: vec3<f32>, material: ResolvedMaterial, F_ab: vec2<f32>) -> vec3<f32> {
118    let H = normalize(wi + wo);
119    let NdotL = dot(world_normal, wi);
120    let NdotH = min(dot(world_normal, H), 1.0);
121    let LdotH = dot(wi, H);
122    let NdotV = dot(world_normal, wo);
123    if NdotL < 0.0001 || NdotH < 0.0001 || LdotH < 0.0001 || NdotV < 0.0001 { return vec3(0.0); }
124
125    let F0_metal = material.base_color;
126    let F0_dielectric = calculate_F0_dielectric(vec3(material.reflectance));
127
128    if material.roughness <= MIRROR_ROUGHNESS_THRESHOLD {
129        let mirror_wi = 2.0 * NdotV * world_normal - wo;
130        if distance(mirror_wi, wi) < 0.0001 {
131            let F_metal = fresnel(F0_metal, LdotH);
132            let F_dielectric = fresnel(F0_dielectric, LdotH);
133            return mix(F_dielectric, F_metal, material.metallic);
134        } else {
135            return vec3(0.0);
136        }
137    }
138
139    let D = D_GGX(material.roughness, NdotH);
140    let Vs = V_SmithGGXCorrelated(material.roughness, NdotV, NdotL);
141    let F_metal = fresnel(F0_metal, LdotH);
142    let F_dielectric = fresnel(F0_dielectric, LdotH);
143    return mix(specular_multiscatter(D, Vs, F_dielectric, F0_dielectric, F_ab, 1.0),
144               specular_multiscatter(D, Vs, F_metal, F0_metal, F_ab, 1.0),
145               material.metallic) * NdotL;
146}
147
148fn brdf_pdf(wo: vec3<f32>, wi: vec3<f32>, world_normal: vec3<f32>, material: ResolvedMaterial, F_ab: vec2<f32>) -> f32 {
149    let NdotV = max(dot(world_normal, wo), 0.0001);
150    let F0_metal = material.base_color;
151    let F0_dielectric = calculate_F0_dielectric(vec3(material.reflectance));
152    let rho = lobe_reflectances(F0_metal, F0_dielectric, material, F_ab);
153    let specular_weight = specular_lobe_sampling_probability(rho, material.perceptual_roughness);
154    let diffuse_weight = 1.0 - specular_weight;
155
156    let TBN = orthonormalize(world_normal);
157    let T = TBN[0];
158    let B = TBN[1];
159    let N = TBN[2];
160
161    let wo_tangent = vec3(dot(wo, T), dot(wo, B), dot(wo, N));
162    let wi_tangent = vec3(dot(wi, T), dot(wi, B), dot(wi, N));
163
164    let diffuse_pdf = wi_tangent.z / PI;
165    let specular_pdf = ggx_vndf_pdf(wo_tangent, wi_tangent, material.roughness);
166    return (diffuse_weight * diffuse_pdf) + (specular_weight * specular_pdf);
167}
168
169fn fresnel(f0: vec3<f32>, LdotH: f32) -> vec3<f32> {
170    return f0 + (1.0 - f0) * pow(saturate(1.0 - LdotH), 5.0);
171}
172
173/// Scale/bias approximation
174fn F_AB(perceptual_roughness: f32, NdotV: f32) -> vec2<f32> {
175    return textureSampleLevel(brdf_dfg_lut, brdf_dfg_lut_sampler, vec2<f32>(NdotV, perceptual_roughness), 0.0).rg;
176}