1#define_import_path bevy_solari::sampling
2
3#import bevy_pbr::lighting::D_GGX
4#import bevy_pbr::utils::{rand_f, rand_vec2f, rand_u, rand_range_u}
5#import bevy_render::maths::{PI_2, orthonormalize}
6#import bevy_solari::scene_bindings::{trace_ray, RAY_T_MIN, RAY_T_MAX, light_sources, directional_lights, LightSource, LIGHT_SOURCE_KIND_DIRECTIONAL, resolve_triangle_data_full, ResolvedRayHitFull}
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    let sum = f + g;
14    if sum == 0.0 {
15        return 0.0;
16    }
17    return max(0.0, f / sum);
18}
19
20// https://gpuopen.com/download/Bounded_VNDF_Sampling_for_Smith-GGX_Reflections.pdf (Listing 1)
21fn sample_ggx_vndf(wi_tangent: vec3<f32>, roughness: f32, rng: ptr<function, u32>) -> vec3<f32> {
22    if roughness <= 0.001 {
23        return vec3(-wi_tangent.xy, wi_tangent.z);
24    }
25
26    let i = wi_tangent;
27    let rand = rand_vec2f(rng);
28    let i_std = normalize(vec3(i.xy * roughness, i.z));
29    let phi = PI_2 * rand.x;
30    let a = roughness;
31    let s = 1.0 + length(vec2(i.xy));
32    let a2 = a * a;
33    let s2 = s * s;
34    let k = (1.0 - a2) * s2 / (s2 + a2 * i.z * i.z);
35    let b = select(i_std.z, k * i_std.z, i.z > 0.0);
36    let z = fma(1.0 - rand.y, 1.0 + b, -b);
37    let sin_theta = sqrt(saturate(1.0 - z * z));
38    let o_std = vec3(sin_theta * cos(phi), sin_theta * sin(phi), z);
39    let m_std = i_std + o_std;
40    let m = normalize(vec3(m_std.xy * roughness, m_std.z));
41    return 2.0 * dot(i, m) * m - i;
42}
43
44// https://gpuopen.com/download/Bounded_VNDF_Sampling_for_Smith-GGX_Reflections.pdf (Listing 2)
45fn ggx_vndf_pdf(wi_tangent: vec3<f32>, wo_tangent: vec3<f32>, roughness: f32) -> f32 {
46    let i = wi_tangent;
47    let o = wo_tangent;
48    let m = normalize(i + o);
49    let ndf = D_GGX(roughness, saturate(m.z));
50    let ai = roughness * i.xy;
51    let len2 = dot(ai, ai);
52    let t = sqrt(len2 + i.z * i.z);
53    if i.z >= 0.0 {
54        let a = roughness;
55        let s = 1.0 + length(i.xy);
56        let a2 = a * a;
57        let s2 = s * s;
58        let k = (1.0 - a2) * s2 / (s2 + a2 * i.z * i.z);
59        return ndf / (2.0 * (k * i.z + t));
60    }
61    return ndf * (t - i.z) / (2.0 * len2);
62}
63
64struct LightSample {
65    light_id: u32,
66    seed: u32,
67}
68
69struct ResolvedLightSample {
70    world_position: vec4<f32>,
71    world_normal: vec3<f32>,
72    radiance: vec3<f32>,
73    inverse_pdf: f32,
74}
75
76struct LightContribution {
77    radiance: vec3<f32>,
78    inverse_pdf: f32,
79    wi: vec3<f32>,
80    brdf_rays_can_hit: bool,
81}
82
83struct LightContributionNoPdf {
84    radiance: vec3<f32>,
85    wi: vec3<f32>,
86}
87
88struct GenerateRandomLightSampleResult {
89    light_sample: LightSample,
90    resolved_light_sample: ResolvedLightSample,
91}
92
93fn sample_random_light(ray_origin: vec3<f32>, origin_world_normal: vec3<f32>, rng: ptr<function, u32>) -> LightContribution {
94    let sample = generate_random_light_sample(rng);
95    var light_contribution = calculate_resolved_light_contribution(sample.resolved_light_sample, ray_origin, origin_world_normal);
96    light_contribution.radiance *= trace_light_visibility(ray_origin, sample.resolved_light_sample.world_position);
97    return light_contribution;
98}
99
100fn random_emissive_light_pdf(hit: ResolvedRayHitFull) -> f32 {
101    let light_count = arrayLength(&light_sources);
102    return 1.0 / (f32(light_count) * f32(hit.triangle_count) * hit.triangle_area);
103}
104
105fn generate_random_light_sample(rng: ptr<function, u32>) -> GenerateRandomLightSampleResult {
106    let light_count = arrayLength(&light_sources);
107    let light_id = rand_range_u(light_count, rng);
108
109    let light_source = light_sources[light_id];
110
111    var triangle_id = 0u;
112    if light_source.kind != LIGHT_SOURCE_KIND_DIRECTIONAL {
113        let triangle_count = light_source.kind >> 1u;
114        triangle_id = rand_range_u(triangle_count, rng);
115    }
116
117    let seed = rand_u(rng);
118    let light_sample = LightSample((light_id << 16u) | triangle_id, seed);
119
120    var resolved_light_sample = resolve_light_sample(light_sample, light_source);
121    resolved_light_sample.inverse_pdf *= f32(light_count);
122
123    return GenerateRandomLightSampleResult(light_sample, resolved_light_sample);
124}
125
126fn resolve_light_sample(light_sample: LightSample, light_source: LightSource) -> ResolvedLightSample {
127    if light_source.kind == LIGHT_SOURCE_KIND_DIRECTIONAL {
128        let directional_light = directional_lights[light_source.id];
129
130#ifndef NO_DIRECTIONAL_LIGHT_SOFT_SHADOWS
131        // Sample a random direction within a cone whose base is the sun approximated as a disk
132        // https://www.realtimerendering.com/raytracinggems/unofficial_RayTracingGems_v1.9.pdf#0004286901.INDD%3ASec30%3A305
133        var rng = light_sample.seed;
134        let random = rand_vec2f(&rng);
135        let cos_theta = (1.0 - random.x) + random.x * directional_light.cos_theta_max;
136        let sin_theta = sqrt(1.0 - cos_theta * cos_theta);
137        let phi = random.y * PI_2;
138        let x = cos(phi) * sin_theta;
139        let y = sin(phi) * sin_theta;
140        var direction_to_light = vec3(x, y, cos_theta);
141
142        // Rotate the ray so that the cone it was sampled from is aligned with the light direction
143        direction_to_light = orthonormalize(directional_light.direction_to_light) * direction_to_light;
144#else
145        let direction_to_light = directional_light.direction_to_light;
146#endif
147
148        return ResolvedLightSample(
149            vec4(direction_to_light, 0.0),
150            -direction_to_light,
151            directional_light.luminance,
152            directional_light.inverse_pdf,
153        );
154    } else {
155        let triangle_count = light_source.kind >> 1u;
156        let triangle_id = light_sample.light_id & 0xFFFFu;
157        let barycentrics = triangle_barycentrics(light_sample.seed);
158        let triangle_data = resolve_triangle_data_full(light_source.id, triangle_id, barycentrics);
159
160        return ResolvedLightSample(
161            vec4(triangle_data.world_position, 1.0),
162            triangle_data.world_normal,
163            triangle_data.material.emissive.rgb,
164            f32(triangle_count) * triangle_data.triangle_area,
165        );
166    }
167}
168
169fn calculate_resolved_light_contribution(resolved_light_sample: ResolvedLightSample, ray_origin: vec3<f32>, origin_world_normal: vec3<f32>) -> LightContribution {
170    let ray = resolved_light_sample.world_position.xyz - (resolved_light_sample.world_position.w * ray_origin);
171    let light_distance = length(ray);
172    let wi = ray / light_distance;
173
174    let cos_theta_origin = saturate(dot(wi, origin_world_normal));
175    let cos_theta_light = saturate(dot(-wi, resolved_light_sample.world_normal));
176    let light_distance_squared = light_distance * light_distance;
177
178    let radiance = resolved_light_sample.radiance * cos_theta_origin * (cos_theta_light / light_distance_squared);
179
180    return LightContribution(radiance, resolved_light_sample.inverse_pdf, wi, resolved_light_sample.world_position.w == 1.0);
181}
182
183fn resolve_and_calculate_light_contribution(light_sample: LightSample, ray_origin: vec3<f32>, origin_world_normal: vec3<f32>) -> LightContributionNoPdf {
184    let resolved_light_sample = resolve_light_sample(light_sample, light_sources[light_sample.light_id >> 16u]);
185    let light_contribution = calculate_resolved_light_contribution(resolved_light_sample, ray_origin, origin_world_normal);
186    return LightContributionNoPdf(light_contribution.radiance, light_contribution.wi);
187}
188
189fn trace_light_visibility(ray_origin: vec3<f32>, light_sample_world_position: vec4<f32>) -> f32 {
190    var ray_direction = light_sample_world_position.xyz;
191    var ray_t_max = RAY_T_MAX;
192
193    if light_sample_world_position.w == 1.0 {
194        let ray = ray_direction - ray_origin;
195        let dist = length(ray);
196        ray_direction = ray / dist;
197        ray_t_max = dist - RAY_T_MIN - RAY_T_MIN;
198    }
199
200    if ray_t_max < RAY_T_MIN { return 0.0; }
201
202    let ray_hit = trace_ray(ray_origin, ray_direction, RAY_T_MIN, ray_t_max, RAY_FLAG_TERMINATE_ON_FIRST_HIT);
203    return f32(ray_hit.kind == RAY_QUERY_INTERSECTION_NONE);
204}
205
206fn trace_point_visibility(ray_origin: vec3<f32>, point: vec3<f32>) -> f32 {
207    let ray = point - ray_origin;
208    let dist = length(ray);
209    let ray_direction = ray / dist;
210
211    let ray_t_max = dist - RAY_T_MIN - RAY_T_MIN;
212    if ray_t_max < RAY_T_MIN { return 0.0; }
213
214    let ray_hit = trace_ray(ray_origin, ray_direction, RAY_T_MIN, ray_t_max, RAY_FLAG_TERMINATE_ON_FIRST_HIT);
215    return f32(ray_hit.kind == RAY_QUERY_INTERSECTION_NONE);
216}
217
218// https://www.realtimerendering.com/raytracinggems/unofficial_RayTracingGems_v1.9.pdf#0004286901.INDD%3ASec22%3A297
219fn triangle_barycentrics(seed: u32) -> vec3<f32> {
220    var rng = seed;
221    var barycentrics = rand_vec2f(&rng);
222    if barycentrics.x + barycentrics.y > 1.0 { barycentrics = 1.0 - barycentrics; }
223    return vec3(1.0 - barycentrics.x - barycentrics.y, barycentrics);
224}