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