1#define_import_path bevy_pbr::utils
2
3#import bevy_pbr::rgb9e5
4#import bevy_render::maths::{PI, PI_2, orthonormalize}
5
6// Generates a random u32 in range [0, u32::MAX].
7//
8// `state` is a mutable reference to a u32 used as the seed.
9//
10// Values are generated via "white noise", with no correlation between values.
11// In shaders, you often want spatial and/or temporal correlation. Use a different RNG method for these use cases.
12//
13// https://www.pcg-random.org
14// https://www.reedbeta.com/blog/hash-functions-for-gpu-rendering
15fn rand_u(state: ptr<function, u32>) -> u32 {
16    *state = *state * 747796405u + 2891336453u;
17    let word = ((*state >> ((*state >> 28u) + 4u)) ^ *state) * 277803737u;
18    return (word >> 22u) ^ word;
19}
20
21// Generates a random f32 in range [0, 1.0].
22fn rand_f(state: ptr<function, u32>) -> f32 {
23    *state = *state * 747796405u + 2891336453u;
24    let word = ((*state >> ((*state >> 28u) + 4u)) ^ *state) * 277803737u;
25    return f32((word >> 22u) ^ word) * bitcast<f32>(0x2f800004u);
26}
27
28// Generates a random vec2<f32> where each value is in range [0, 1.0].
29fn rand_vec2f(state: ptr<function, u32>) -> vec2<f32> {
30    return vec2(rand_f(state), rand_f(state));
31}
32
33// Generates a random u32 in range [0, n).
34fn rand_range_u(n: u32, state: ptr<function, u32>) -> u32 {
35    return rand_u(state) % n;
36}
37
38// returns the (0-1, 0-1) position within the given viewport for the current buffer coords .
39// buffer coords can be obtained from `@builtin(position).xy`.
40// the view uniform struct contains the current camera viewport in `view.viewport`.
41// topleft = 0,0
42fn coords_to_viewport_uv(position: vec2<f32>, viewport: vec4<f32>) -> vec2<f32> {
43    return (position - viewport.xy) / viewport.zw;
44}
45
46// https://jcgt.org/published/0003/02/01/paper.pdf
47
48// For encoding normals or unit direction vectors as octahedral coordinates.
49fn octahedral_encode(v: vec3<f32>) -> vec2<f32> {
50    var n = v / (abs(v.x) + abs(v.y) + abs(v.z));
51    let octahedral_wrap = (1.0 - abs(n.yx)) * select(vec2(-1.0), vec2(1.0), n.xy > vec2f(0.0));
52    let n_xy = select(octahedral_wrap, n.xy, n.z >= 0.0);
53    return n_xy * 0.5 + 0.5;
54}
55
56// For decoding normals or unit direction vectors from octahedral coordinates.
57fn octahedral_decode(v: vec2<f32>) -> vec3<f32> {
58    let f = v * 2.0 - 1.0;
59    return octahedral_decode_signed(f);
60}
61
62// Like octahedral_decode, but for input in [-1, 1] instead of [0, 1].
63fn octahedral_decode_signed(v: vec2<f32>) -> vec3<f32> {
64    var n = vec3(v.xy, 1.0 - abs(v.x) - abs(v.y));
65    let t = saturate(-n.z);
66    let w = select(vec2(t), vec2(-t), n.xy >= vec2(0.0));
67    n = vec3(n.xy + w, n.z);
68    return normalize(n);
69}
70
71// https://blog.demofox.org/2022/01/01/interleaved-gradient-noise-a-different-kind-of-low-discrepancy-sequence
72fn interleaved_gradient_noise(pixel_coordinates: vec2<f32>, frame: u32) -> f32 {
73    let xy = pixel_coordinates + 5.588238 * f32(frame % 64u);
74    return fract(52.9829189 * fract(0.06711056 * xy.x + 0.00583715 * xy.y));
75}
76
77// Hammersley sequence for quasi-random points
78fn hammersley_2d(i: u32, n: u32) -> vec2f {
79    let inv_n = 1.0 / f32(n);
80    let vdc = f32(reverseBits(i)) * 2.3283064365386963e-10; // 1/2^32
81    return vec2f(f32(i) * inv_n, vdc);
82}
83
84// https://www.iryoku.com/next-generation-post-processing-in-call-of-duty-advanced-warfare (slides 120-135)
85// TODO: Use an array here instead of a bunch of constants, once arrays work properly under DX12.
86// NOTE: The names have a final underscore to avoid the following error:
87// `Composable module identifiers must not require substitution according to naga writeback rules`
88const SPIRAL_OFFSET_0_ = vec2<f32>(-0.7071,  0.7071);
89const SPIRAL_OFFSET_1_ = vec2<f32>(-0.0000, -0.8750);
90const SPIRAL_OFFSET_2_ = vec2<f32>( 0.5303,  0.5303);
91const SPIRAL_OFFSET_3_ = vec2<f32>(-0.6250, -0.0000);
92const SPIRAL_OFFSET_4_ = vec2<f32>( 0.3536, -0.3536);
93const SPIRAL_OFFSET_5_ = vec2<f32>(-0.0000,  0.3750);
94const SPIRAL_OFFSET_6_ = vec2<f32>(-0.1768, -0.1768);
95const SPIRAL_OFFSET_7_ = vec2<f32>( 0.1250,  0.0000);
96
97// https://www.realtimerendering.com/raytracinggems/unofficial_RayTracingGems_v1.9.pdf#0004286901.INDD%3ASec28%3A303
98fn sample_cosine_hemisphere(normal: vec3<f32>, rng: ptr<function, u32>) -> vec3<f32> {
99    let cos_theta = 1.0 - 2.0 * rand_f(rng);
100    let phi = PI_2 * rand_f(rng);
101    let sin_theta = sqrt(max(1.0 - cos_theta * cos_theta, 0.0));
102    let direction = normal + vec3(sin_theta * cos(phi), sin_theta * sin(phi), cos_theta);
103    let len_sq = dot(direction, direction);
104    if len_sq < 1e-8 { return normal; }
105    return direction * inverseSqrt(len_sq);
106}
107
108// https://www.pbr-book.org/3ed-2018/Monte_Carlo_Integration/2D_Sampling_with_Multidimensional_Transformations#UniformlySamplingaHemisphere
109fn sample_uniform_hemisphere(normal: vec3<f32>, rng: ptr<function, u32>) -> vec3<f32> {
110    let cos_theta = rand_f(rng);
111    let phi = PI_2 * rand_f(rng);
112    let sin_theta = sqrt(max(1.0 - cos_theta * cos_theta, 0.0));
113    let x = sin_theta * cos(phi);
114    let y = sin_theta * sin(phi);
115    let z = cos_theta;
116    return orthonormalize(normal) * vec3(x, y, z);
117}
118
119fn uniform_hemisphere_inverse_pdf() -> f32 {
120    return PI_2;
121}
122
123// https://www.realtimerendering.com/raytracinggems/unofficial_RayTracingGems_v1.9.pdf#0004286901.INDD%3ASec19%3A294
124fn sample_disk(disk_radius: f32, rng: ptr<function, u32>) -> vec2<f32> {
125    let ab = 2.0 * rand_vec2f(rng) - 1.0;
126    let a = ab.x;
127    var b = ab.y;
128    if (b == 0.0) { b = 1.0; }
129
130    var phi: f32;
131    var r: f32;
132    if (a * a > b * b) {
133        r = disk_radius * a;
134        phi = (PI / 4.0) * (b / a);
135    } else {
136        r = disk_radius * b;
137        phi = (PI / 2.0) - (PI / 4.0) * (a / b);
138    }
139
140    let x = r * cos(phi);
141    let y = r * sin(phi);
142    return vec2(x, y);
143}
144
145// Convert UV and face index to direction vector
146fn sample_cube_dir(uv: vec2f, face: u32) -> vec3f {
147    // Convert from [0,1] to [-1,1]
148    let uvc = 2.0 * uv - 1.0;
149
150    // Generate direction based on the cube face
151    var dir: vec3f;
152    switch(face) {
153        case 0u: { dir = vec3f( 1.0,  -uvc.y, -uvc.x); } // +X
154        case 1u: { dir = vec3f(-1.0,  -uvc.y,  uvc.x); } // -X
155        case 2u: { dir = vec3f( uvc.x,  1.0,   uvc.y); } // +Y
156        case 3u: { dir = vec3f( uvc.x, -1.0,  -uvc.y); } // -Y
157        case 4u: { dir = vec3f( uvc.x, -uvc.y,  1.0);  } // +Z
158        case 5u: { dir = vec3f(-uvc.x, -uvc.y, -1.0);  } // -Z
159        default: { dir = vec3f(0.0); }
160    }
161    return normalize(dir);
162}
163
164// Convert direction vector to cube face UV
165struct CubeUV {
166    uv: vec2f,
167    face: u32,
168}
169fn dir_to_cube_uv(dir: vec3f) -> CubeUV {
170    let abs_dir = abs(dir);
171    var face: u32 = 0u;
172    var uv: vec2f = vec2f(0.0);
173
174    // Find the dominant axis to determine face
175    if (abs_dir.x >= abs_dir.y && abs_dir.x >= abs_dir.z) {
176        // X axis is dominant
177        if (dir.x > 0.0) {
178            face = 0u; // +X
179            uv = vec2f(-dir.z, -dir.y) / dir.x;
180        } else {
181            face = 1u; // -X
182            uv = vec2f(dir.z, -dir.y) / abs_dir.x;
183        }
184    } else if (abs_dir.y >= abs_dir.x && abs_dir.y >= abs_dir.z) {
185        // Y axis is dominant
186        if (dir.y > 0.0) {
187            face = 2u; // +Y
188            uv = vec2f(dir.x, dir.z) / dir.y;
189        } else {
190            face = 3u; // -Y
191            uv = vec2f(dir.x, -dir.z) / abs_dir.y;
192        }
193    } else {
194        // Z axis is dominant
195        if (dir.z > 0.0) {
196            face = 4u; // +Z
197            uv = vec2f(dir.x, -dir.y) / dir.z;
198        } else {
199            face = 5u; // -Z
200            uv = vec2f(-dir.x, -dir.y) / abs_dir.z;
201        }
202    }
203
204    // Convert from [-1,1] to [0,1]
205    return CubeUV(uv * 0.5 + 0.5, face);
206}
207
208// The Porter-Duff OVER operator on RGBA, correctly computing alpha of the
209// result.
210fn porter_duff_over(bg: vec4<f32>, fg: vec4<f32>) -> vec4<f32> {
211    return vec4<f32>(mix(bg.rgb * bg.a, fg.rgb, fg.a), bg.a + fg.a * (1.0 - bg.a));
212}