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