1import bevy_pbr::render::pbr_lighting::D_GGX;
2import bevy_pbr::render::utils::{rand_vec2f, rand_u, rand_range_u};
3import bevy_render::maths::{PI_2, orthonormalize};
4import package::scene::bindings::{trace_ray, trace_ray_previous_frame, RAY_T_MIN, RAY_T_MAX, light_sources, directional_lights, environment_map, environment_map_sampler, environment_map_light, LightSource, LIGHT_SOURCE_KIND_DIRECTIONAL, resolve_triangle_data_full, ResolvedRayHitFull, MIRROR_ROUGHNESS_THRESHOLD};
5
6enable wgpu_ray_query;
7
8fn power_heuristic(f: f32, g: f32) -> f32 {
9    return balance_heuristic(f * f, g * g);
10}
11
12fn balance_heuristic(f: f32, g: f32) -> f32 {
13    // Need to guard against NaNs since ReSTIR reservoirs can have UCW=0
14    if f == 0.0 {
15        return 0.0;
16    }
17    // When NaN, fall back to even weights of 0.5 : 0.5
18    let weight = 1.0 / (1.0 + (g / f));
19    return select(weight, 0.5, isnan(weight));
20}
21
22/// https://gpuopen.com/download/Bounded_VNDF_Sampling_for_Smith-GGX_Reflections.pdf (Listing 1)
23/// Result is invalid when output.z <= 0.0, and must be discarded
24fn sample_ggx_vndf(wi_tangent: vec3<f32>, roughness: f32, rng: ptr<function, u32>) -> vec3<f32> {
25    // Mirror BRDF case
26    if roughness <= MIRROR_ROUGHNESS_THRESHOLD {
27        return vec3(-wi_tangent.xy, wi_tangent.z);
28    }
29
30    let i = wi_tangent;
31    let rand = rand_vec2f(rng);
32    let i_std = normalize(vec3(i.xy * roughness, i.z));
33    let phi = PI_2 * rand.x;
34    let a = roughness;
35    let s = 1.0 + length(vec2(i.xy));
36    let a2 = a * a;
37    let s2 = s * s;
38    let k = (1.0 - a2) * s2 / (s2 + a2 * i.z * i.z);
39    let b = select(i_std.z, k * i_std.z, i.z > 0.0);
40    let z = fma(1.0 - rand.y, 1.0 + b, -b);
41    let sin_theta = sqrt(saturate(1.0 - z * z));
42    let o_std = vec3(sin_theta * cos(phi), sin_theta * sin(phi), z);
43    let m_std = i_std + o_std;
44    let m = normalize(vec3(m_std.xy * roughness, m_std.z));
45    return 2.0 * dot(i, m) * m - i;
46}
47
48fn ggx_vndf_sample_invalid(ray_tangent: vec3<f32>) -> bool {
49    return !(ray_tangent.z > 0.0);
50}
51
52/// https://gpuopen.com/download/Bounded_VNDF_Sampling_for_Smith-GGX_Reflections.pdf (Listing 2)
53fn ggx_vndf_pdf(wi_tangent: vec3<f32>, wo_tangent: vec3<f32>, roughness: f32) -> f32 {
54    // Mirror BRDF case
55    if roughness <= MIRROR_ROUGHNESS_THRESHOLD {
56        let mirror_wo = vec3(-wi_tangent.xy, wi_tangent.z);
57        if distance(mirror_wo, wo_tangent) < 0.0001 {
58            return bitcast<f32>(0x7F800000u); // INF
59        } else {
60            return 0.0;
61        }
62    }
63
64    let i = wi_tangent;
65    let o = wo_tangent;
66    let m = normalize(i + o);
67    let ndf = D_GGX(roughness, saturate(m.z));
68    let ai = roughness * i.xy;
69    let len2 = dot(ai, ai);
70    let t = sqrt(len2 + i.z * i.z);
71    var pdf: f32;
72    if i.z >= 0.0 {
73        let a = roughness;
74        let s = 1.0 + length(i.xy);
75        let a2 = a * a;
76        let s2 = s * s;
77        let k = (1.0 - a2) * s2 / (s2 + a2 * i.z * i.z);
78        pdf = ndf / (2.0 * (k * i.z + t));
79    } else {
80        pdf = ndf * (t - i.z) / (2.0 * len2);
81    }
82
83    return select(pdf, 0.0, isnan(pdf));
84}
85
86fn isinf(x: f32) -> bool {
87    return (bitcast<u32>(x) & 0x7fffffffu) == 0x7f800000u;
88}
89
90fn isnan(x: f32) -> bool {
91    return (bitcast<u32>(x) & 0x7fffffffu) > 0x7f800000u;
92}
93
94const NULL_LIGHT_ID = 0xFFFFFFFFu;
95
96struct LightSample {
97    light_id: u32,
98    seed: u32,
99}
100
101struct ResolvedLightSample {
102    world_position: vec4<f32>,
103    world_normal: vec3<f32>,
104    radiance: vec3<f32>,
105    inverse_pdf: f32,
106}
107
108struct LightContribution {
109    radiance: vec3<f32>,
110    inverse_pdf: f32,
111    /// inverse_pdf may not be in solid angle measure (e.g. area measure)
112    /// but inverse_solid_angle_pdf is always in solid angle measure
113    /// for cases where it's required, e.g. MIS between light sampling techniques
114    inverse_solid_angle_pdf: f32,
115    wi: vec3<f32>,
116    brdf_rays_can_hit: bool,
117}
118
119struct GenerateRandomLightSampleResult {
120    light_sample: LightSample,
121    resolved_light_sample: ResolvedLightSample,
122}
123
124fn sample_random_light(ray_origin: vec3<f32>, origin_world_normal: vec3<f32>, rng: ptr<function, u32>) -> LightContribution {
125    let sample = generate_random_light_sample(rng);
126    var light_contribution = calculate_resolved_light_contribution(sample.resolved_light_sample, ray_origin, origin_world_normal);
127    light_contribution.radiance *= trace_visibility(ray_origin, sample.resolved_light_sample.world_position);
128    return light_contribution;
129}
130
131fn random_emissive_light_pdf(hit: ResolvedRayHitFull, ray_distance: f32, NdotV: f32) -> f32 {
132    let light_count = arrayLength(&light_sources);
133    let area_pdf = 1.0 / (f32(light_count) * f32(hit.triangle_count) * hit.triangle_area);
134    return area_pdf * (ray_distance * ray_distance) / NdotV;
135}
136
137fn sample_environment_map_light(world_direction: vec3<f32>) -> vec3<f32> {
138    if environment_map_light.intensity == 0.0 { return vec3(0.0); }
139
140    var sample_direction = environment_map_light.light_from_world * world_direction;
141    sample_direction.z = -sample_direction.z;
142
143    // Hardcoded to sample mip 0 (unfiltered radiance)
144    return textureSampleLevel(environment_map, environment_map_sampler, sample_direction, 0.0).rgb * environment_map_light.intensity;
145}
146
147fn generate_random_light_sample(rng: ptr<function, u32>) -> GenerateRandomLightSampleResult {
148    let light_count = arrayLength(&light_sources);
149    let light_id = rand_range_u(light_count, rng);
150
151    let light_source = light_sources[light_id];
152
153    var triangle_id = 0u;
154    if light_source.kind != LIGHT_SOURCE_KIND_DIRECTIONAL {
155        let triangle_count = light_source.kind >> 1u;
156        triangle_id = rand_range_u(triangle_count, rng);
157    }
158
159    let seed = rand_u(rng);
160    let light_sample = LightSample((light_id << 16u) | triangle_id, seed);
161
162    var resolved_light_sample = resolve_light_sample(light_sample, light_source);
163    resolved_light_sample.inverse_pdf *= f32(light_count);
164
165    return GenerateRandomLightSampleResult(light_sample, resolved_light_sample);
166}
167
168fn resolve_light_sample(light_sample: LightSample, light_source: LightSource) -> ResolvedLightSample {
169    if light_source.kind == LIGHT_SOURCE_KIND_DIRECTIONAL {
170        let directional_light = directional_lights[light_source.id];
171
172        var direction_to_light: vec3<f32>;
173@if(!NO_DIRECTIONAL_LIGHT_SOFT_SHADOWS) {
174        // Sample a random direction within a cone whose base is the sun approximated as a disk
175        // https://www.realtimerendering.com/raytracinggems/unofficial_RayTracingGems_v1.9.pdf#0004286901.INDD%3ASec30%3A305
176        var rng = light_sample.seed;
177        let random = rand_vec2f(&rng);
178        let cos_theta = (1.0 - random.x) + random.x * directional_light.cos_theta_max;
179        let sin_theta = sqrt(1.0 - cos_theta * cos_theta);
180        let phi = random.y * PI_2;
181        let x = cos(phi) * sin_theta;
182        let y = sin(phi) * sin_theta;
183        direction_to_light = vec3(x, y, cos_theta);
184
185        // Rotate the ray so that the cone it was sampled from is aligned with the light direction
186        direction_to_light = orthonormalize(directional_light.direction_to_light) * direction_to_light;
187} @else {
188        direction_to_light = directional_light.direction_to_light;
189}
190
191        return ResolvedLightSample(
192        vec4(direction_to_light, 0.0),
193        -direction_to_light,
194        directional_light.luminance,
195        directional_light.inverse_pdf,
196    );
197} else {
198        let triangle_count = light_source.kind >> 1u;
199        let triangle_id = light_sample.light_id & 0xFFFFu;
200        let barycentrics = triangle_barycentrics(light_sample.seed);
201        let triangle_data = resolve_triangle_data_full(light_source.id, triangle_id, barycentrics);
202
203        return ResolvedLightSample(
204            vec4(triangle_data.world_position, 1.0),
205            triangle_data.world_normal,
206            triangle_data.material.emissive.rgb,
207            f32(triangle_count) * triangle_data.triangle_area,
208        );
209    }
210}
211
212fn calculate_resolved_light_contribution(resolved_light_sample: ResolvedLightSample, ray_origin: vec3<f32>, origin_world_normal: vec3<f32>) -> LightContribution {
213    let ray = resolved_light_sample.world_position.xyz - (resolved_light_sample.world_position.w * ray_origin);
214    let light_distance = length(ray);
215    let wi = ray / light_distance;
216
217    let cos_theta_light = saturate(dot(-wi, resolved_light_sample.world_normal));
218    let light_distance_squared = light_distance * light_distance;
219    let denominator = cos_theta_light / light_distance_squared;
220
221    let radiance = resolved_light_sample.radiance * denominator;
222    let inverse_solid_angle_pdf = resolved_light_sample.inverse_pdf * denominator;
223
224    return LightContribution(radiance, resolved_light_sample.inverse_pdf, inverse_solid_angle_pdf, wi, resolved_light_sample.world_position.w == 1.0);
225}
226
227fn trace_visibility(ray_origin: vec3<f32>, point: vec4<f32>) -> f32 {
228    var ray_direction = point.xyz;
229    var ray_t_max = RAY_T_MAX;
230
231    if point.w == 1.0 {
232        let ray = ray_direction - ray_origin;
233        let dist = length(ray);
234        ray_direction = ray / dist;
235        ray_t_max = dist - RAY_T_MIN;
236    }
237
238    if ray_t_max < RAY_T_MIN { return 0.0; }
239
240    let ray_hit = trace_ray(ray_origin, ray_direction, RAY_T_MIN, ray_t_max, RAY_FLAG_TERMINATE_ON_FIRST_HIT);
241    return f32(ray_hit.kind == RAY_QUERY_INTERSECTION_NONE);
242}
243
244fn trace_visibility_previous_frame(ray_origin: vec3<f32>, point: vec4<f32>) -> f32 {
245    var ray_direction = point.xyz;
246    var ray_t_max = RAY_T_MAX;
247
248    if point.w == 1.0 {
249        let ray = ray_direction - ray_origin;
250        let dist = length(ray);
251        ray_direction = ray / dist;
252        ray_t_max = dist - RAY_T_MIN;
253    }
254
255    if ray_t_max < RAY_T_MIN { return 0.0; }
256
257    let ray_hit = trace_ray_previous_frame(ray_origin, ray_direction, RAY_T_MIN, ray_t_max, RAY_FLAG_TERMINATE_ON_FIRST_HIT);
258    return f32(ray_hit.kind == RAY_QUERY_INTERSECTION_NONE);
259}
260
261/// https://www.realtimerendering.com/raytracinggems/unofficial_RayTracingGems_v1.9.pdf#0004286901.INDD%3ASec22%3A297
262fn triangle_barycentrics(seed: u32) -> vec3<f32> {
263    var rng = seed;
264    var barycentrics = rand_vec2f(&rng);
265    if barycentrics.x + barycentrics.y > 1.0 { barycentrics = 1.0 - barycentrics; }
266    return vec3(1.0 - barycentrics.x - barycentrics.y, barycentrics);
267}