1import constants::WORLD_CACHE_SIZE;
2import bevy_pbr::render::utils::{rand_f, rand_vec2f};
3import bevy_render::maths::orthonormalize;
4import package::realtime::bindings::{
5    world_cache,
6    constants,
7};
8
9enable wgpu_ray_query;
10
11/// Maximum amount of frames a cell can live for without being queried
12const WORLD_CACHE_CELL_LIFETIME: u32 = 10u;
13/// Maximum amount of attempts to find a cache entry after a hash collision
14const WORLD_CACHE_MAX_SEARCH_STEPS: u32 = 3u;
15
16/// Marker value for an empty cell
17const WORLD_CACHE_EMPTY_CELL: u32 = 0u;
18
19@if(!WORLD_CACHE_NON_ATOMIC_LIFE_BUFFER)
20fn 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> {
21    var world_position = world_position_in;
22    var cell_size = get_cell_size(world_position, view_position, ray_t, rng);
23
24@if(!NO_JITTER_WORLD_CACHE) {
25    // Jitter query point, which essentially blurs the cache a bit so it's not so grid-like
26    // https://tomclabault.github.io/blog/2025/regir, jitter_world_position_tangent_plane
27    let TBN = orthonormalize(world_normal);
28    let offset = (rand_vec2f(rng) * 2.0 - 1.0) * cell_size * 0.5;
29    world_position += offset.x * TBN[0] + offset.y * TBN[1];
30    cell_size = get_cell_size(world_position, view_position, ray_t, rng);
31}
32
33    let world_position_quantized = bitcast<vec3<u32>>(quantize_position(world_position, cell_size));
34    let world_normal_quantized = bitcast<vec3<u32>>(quantize_normal(world_normal));
35    var key = compute_key(world_position_quantized, world_normal_quantized);
36    let checksum = compute_checksum(world_position_quantized, world_normal_quantized);
37
38    for (var i = 0u; i < WORLD_CACHE_MAX_SEARCH_STEPS; i++) {
39        let cas = atomicCompareExchangeWeak(&world_cache.checksums[key], WORLD_CACHE_EMPTY_CELL, checksum);
40        let existing_checksum = cas.old_value;
41
42        // Cell already exists or is empty - reset lifetime
43        if existing_checksum == checksum || existing_checksum == WORLD_CACHE_EMPTY_CELL {
44            @if(!WORLD_CACHE_QUERY_ATOMIC_MAX_LIFETIME)
45            atomicStore(&world_cache.life[key], cell_lifetime);
46            @else
47            atomicMax(&world_cache.life[key], cell_lifetime);
48        }
49
50        if existing_checksum == checksum {
51            // Cache entry already exists - get radiance
52            return world_cache.radiance[key].rgb;
53        } else if existing_checksum == WORLD_CACHE_EMPTY_CELL && cas.exchanged {
54            // Cell is empty - initialize it
55            world_cache.geometry_data[key].world_position = world_position;
56            world_cache.geometry_data[key].world_normal = world_normal;
57            return vec3(0.0);
58        } else {
59            // Collision - linear probe to next entry
60            key += 1u;
61        }
62    }
63
64    return vec3(0.0);
65}
66
67fn get_cell_size(world_position: vec3<f32>, view_position: vec3<f32>, ray_t: f32, rng: ptr<function, u32>) -> f32 {
68    let camera_distance = distance(view_position, world_position) / constants.world_cache_position_lod_scale;
69    let lod_f = log2(1.0 + camera_distance);
70    let lod_fract = fract(lod_f);
71    let lod = floor(lod_f) + select(0.0, 1.0, rand_f(rng) < lod_fract * lod_fract * lod_fract);
72    return constants.world_cache_position_base_cell_size * exp2(lod);
73}
74
75fn quantize_position(world_position: vec3<f32>, quantization_factor: f32) -> vec3<f32> {
76    return floor(world_position / quantization_factor + 0.0001);
77}
78
79fn quantize_normal(world_normal: vec3<f32>) -> vec3<f32> {
80    return floor(world_normal + 0.0001);
81}
82
83fn compute_key(world_position: vec3<u32>, world_normal: vec3<u32>) -> u32 {
84    var key = pcg_hash(world_position.x);
85    key = pcg_hash(key + world_position.y);
86    key = pcg_hash(key + world_position.z);
87    key = pcg_hash(key + world_normal.x);
88    key = pcg_hash(key + world_normal.y);
89    key = pcg_hash(key + world_normal.z);
90    return wrap_key(key);
91}
92
93fn compute_checksum(world_position: vec3<u32>, world_normal: vec3<u32>) -> u32 {
94    var key = iqint_hash(world_position.x);
95    key = iqint_hash(key + world_position.y);
96    key = iqint_hash(key + world_position.z);
97    key = iqint_hash(key + world_normal.x);
98    key = iqint_hash(key + world_normal.y);
99    key = iqint_hash(key + world_normal.z);
100    return max(key, 1u); // 0u is reserved for WORLD_CACHE_EMPTY_CELL
101}
102
103fn pcg_hash(input: u32) -> u32 {
104    let state = input * 747796405u + 2891336453u;
105    let word = ((state >> ((state >> 28u) + 4u)) ^ state) * 277803737u;
106    return (word >> 22u) ^ word;
107}
108
109fn iqint_hash(input: u32) -> u32 {
110    let n = (input << 13u) ^ input;
111    return n * (n * n * 15731u + 789221u) + 1376312589u;
112}
113
114fn wrap_key(key: u32) -> u32 {
115    return key & (u32(WORLD_CACHE_SIZE) - 1u);
116}