1enable wgpu_ray_query;
2
3#define_import_path bevy_solari::world_cache
4
5#import bevy_pbr::utils::{rand_f, rand_vec2f}
6#import bevy_render::maths::orthonormalize
7#import bevy_solari::realtime_bindings::{
8    world_cache_life,
9    world_cache_checksums,
10    world_cache_radiance,
11    world_cache_geometry_data,
12    world_cache_luminance_deltas,
13    world_cache_a,
14    world_cache_b,
15    world_cache_active_cell_indices,
16    world_cache_active_cells_count,
17    WorldCacheGeometryData,
18}
19
20/// How responsive the world cache is to changes in lighting (higher is less responsive but more stable, lower is more responsive but less stable)
21const WORLD_CACHE_MAX_TEMPORAL_SAMPLES: f32 = 32.0;
22/// How many direct light samples each cell takes when updating each frame
23const WORLD_CACHE_DIRECT_LIGHT_SAMPLE_COUNT: u32 = 32u;
24/// Maximum amount of distance to trace GI rays between two cache cells
25const WORLD_CACHE_MAX_GI_RAY_DISTANCE: f32 = 50.0;
26/// Soft upper limit on the amount of cache cells to update each frame
27const WORLD_CACHE_CELL_UPDATES_SOFT_CAP: u32 = 40000u;
28
29/// Maximum amount of frames a cell can live for without being queried
30const WORLD_CACHE_CELL_LIFETIME: u32 = 10u;
31/// Maximum amount of attempts to find a cache entry after a hash collision
32const WORLD_CACHE_MAX_SEARCH_STEPS: u32 = 3u;
33
34/// Size of a cache cell at the lowest LOD in meters
35const WORLD_CACHE_POSITION_BASE_CELL_SIZE: f32 = 0.15;
36/// How fast the world cache transitions between LODs as a function of distance to the camera
37const WORLD_CACHE_POSITION_LOD_SCALE: f32 = 15.0;
38
39/// Marker value for an empty cell
40const WORLD_CACHE_EMPTY_CELL: u32 = 0u;
41
42#ifndef WORLD_CACHE_NON_ATOMIC_LIFE_BUFFER
43fn query_world_cache(world_position_in: vec3<f32>, world_normal: vec3<f32>, view_position: vec3<f32>, ray_t: f32, cell_lifetime: u32, rng: ptr<function, u32>) -> vec3<f32> {
44    var world_position = world_position_in;
45    var cell_size = get_cell_size(world_position, view_position, rng);
46
47#ifdef WORLD_CACHE_FIRST_BOUNCE_LIGHT_LEAK_PREVENTION
48    if ray_t < cell_size {
49        // Prevent light leaks
50        cell_size = WORLD_CACHE_POSITION_BASE_CELL_SIZE;
51    }
52#endif
53
54#ifndef NO_JITTER_WORLD_CACHE
55    // Jitter query point, which essentially blurs the cache a bit so it's not so grid-like
56    // https://tomclabault.github.io/blog/2025/regir, jitter_world_position_tangent_plane
57    let TBN = orthonormalize(world_normal);
58    let offset = (rand_vec2f(rng) * 2.0 - 1.0) * cell_size * 0.5;
59    world_position += offset.x * TBN[0] + offset.y * TBN[1];
60    cell_size = get_cell_size(world_position, view_position, rng);
61#endif
62
63    let world_position_quantized = bitcast<vec3<u32>>(quantize_position(world_position, cell_size));
64    let world_normal_quantized = bitcast<vec3<u32>>(quantize_normal(world_normal));
65    var key = compute_key(world_position_quantized, world_normal_quantized);
66    let checksum = compute_checksum(world_position_quantized, world_normal_quantized);
67
68    for (var i = 0u; i < WORLD_CACHE_MAX_SEARCH_STEPS; i++) {
69        let existing_checksum = atomicCompareExchangeWeak(&world_cache_checksums[key], WORLD_CACHE_EMPTY_CELL, checksum).old_value;
70
71        // Cell already exists or is empty - reset lifetime
72        if existing_checksum == checksum || existing_checksum == WORLD_CACHE_EMPTY_CELL {
73#ifndef WORLD_CACHE_QUERY_ATOMIC_MAX_LIFETIME
74            atomicStore(&world_cache_life[key], cell_lifetime);
75#else
76            atomicMax(&world_cache_life[key], cell_lifetime);
77#endif
78        }
79
80        if existing_checksum == checksum {
81            // Cache entry already exists - get radiance
82            return world_cache_radiance[key].rgb;
83        } else if existing_checksum == WORLD_CACHE_EMPTY_CELL {
84            // Cell is empty - initialize it
85            world_cache_geometry_data[key].world_position = world_position;
86            world_cache_geometry_data[key].world_normal = world_normal;
87            return vec3(0.0);
88        } else {
89            // Collision - linear probe to next entry
90            key += 1u;
91        }
92    }
93
94    return vec3(0.0);
95}
96#endif
97
98fn get_cell_size(world_position: vec3<f32>, view_position: vec3<f32>, rng: ptr<function, u32>) -> f32 {
99    let camera_distance = distance(view_position, world_position) / WORLD_CACHE_POSITION_LOD_SCALE;
100    let lod_f = log2(1.0 + camera_distance);
101    let lod_fract = fract(lod_f);
102    let lod = floor(lod_f) + select(0.0, 1.0, rand_f(rng) < lod_fract * lod_fract * lod_fract);
103    return WORLD_CACHE_POSITION_BASE_CELL_SIZE * exp2(lod);
104}
105
106fn quantize_position(world_position: vec3<f32>, quantization_factor: f32) -> vec3<f32> {
107    return floor(world_position / quantization_factor + 0.0001);
108}
109
110fn quantize_normal(world_normal: vec3<f32>) -> vec3<f32> {
111    return floor(world_normal * 2.0 + 0.0001);
112}
113
114// TODO: Clustering
115fn compute_key(world_position: vec3<u32>, world_normal: vec3<u32>) -> u32 {
116    var key = pcg_hash(world_position.x);
117    key = pcg_hash(key + world_position.y);
118    key = pcg_hash(key + world_position.z);
119    key = pcg_hash(key + world_normal.x);
120    key = pcg_hash(key + world_normal.y);
121    key = pcg_hash(key + world_normal.z);
122    return wrap_key(key);
123}
124
125fn compute_checksum(world_position: vec3<u32>, world_normal: vec3<u32>) -> u32 {
126    var key = iqint_hash(world_position.x);
127    key = iqint_hash(key + world_position.y);
128    key = iqint_hash(key + world_position.z);
129    key = iqint_hash(key + world_normal.x);
130    key = iqint_hash(key + world_normal.y);
131    key = iqint_hash(key + world_normal.z);
132    return max(key, 1u); // 0u is reserved for WORLD_CACHE_EMPTY_CELL
133}
134
135fn pcg_hash(input: u32) -> u32 {
136    let state = input * 747796405u + 2891336453u;
137    let word = ((state >> ((state >> 28u) + 4u)) ^ state) * 277803737u;
138    return (word >> 22u) ^ word;
139}
140
141fn iqint_hash(input: u32) -> u32 {
142    let n = (input << 13u) ^ input;
143    return n * (n * n * 15731u + 789221u) + 1376312589u;
144}
145
146fn wrap_key(key: u32) -> u32 {
147    return key & (#{WORLD_CACHE_SIZE} - 1u);
148}