1// Samples the rendered displacement field at gameplay probe positions.
2// LOD selection and cascade blending match the vertex path.
3
4#import bevy_aqua_core::river::{river_analytic_displacement}
5
6const LOD_COUNT: u32 = 5u;
7const CASCADE_COUNT: u32 = LOD_COUNT + 1u;
8const WAVE_COUNT: u32 = 40u;
9const MAX_QUERIES: u32 = 256u;
10const LOD_TRANSITION_START: f32 = 1.0;
11const MORPH_BLACK_POINT: f32 = 0.05;
12const MORPH_FADE_SIDES: f32 = 2.0;
13const UV_CENTER: f32 = 0.5;
14const MIN_SAMPLE_WEIGHT: f32 = 0.001;
15const MIN_NORMAL_Y: f32 = 0.0001;
16
17struct CascadeParams {
18    center: vec2<f32>,
19    scale: f32,
20    texture_res: f32,
21    inv_texture_res: f32,
22    texel_width: f32,
23    weight: f32,
24    max_wavelength: f32,
25}
26struct CascadeLayout {
27    cascades: array<CascadeParams, CASCADE_COUNT>,
28    center: vec4<f32>,
29    // XY bed-map first-texel world origin, ZW inverse world extent.
30    bed_transform: vec4<f32>,
31    // X height minimum, Y height span (negative = no bed map), Z sea level.
32    bed_range: vec4<f32>,
33}
34struct GerstnerWave {
35    direction: vec2<f32>,
36    amplitude: f32,
37    wave_number: f32,
38    angular_frequency: f32,
39    phase: f32,
40    chop_amplitude: f32,
41}
42struct AnimWavesUniform {
43    cascade_layout: CascadeLayout,
44    waves: array<GerstnerWave, WAVE_COUNT>,
45    ranges: array<vec4<u32>, LOD_COUNT>,
46    time: vec4<f32>,
47    flow: vec4<f32>,
48}
49
50// `flow`: xy current (m/s), z signed bank margin (m), w channel half-width
51// (m). Zero current selects cascade sampling.
52struct QueryRequest {
53    world_xz: vec2<f32>,
54    slot: f32,
55    kind: f32,
56    flow: vec4<f32>,
57}
58
59struct QueryResult {
60    displacement_slot: vec4<f32>,
61    normal_validity: vec4<f32>,
62    signals: vec4<f32>,
63}
64
65@group(0) @binding(0) var lod_data: texture_2d_array<f32>;
66@group(0) @binding(1) var lod_sampler: sampler;
67@group(0) @binding(2) var<uniform> params: AnimWavesUniform;
68@group(0) @binding(3) var<storage, read> requests: array<QueryRequest>;
69@group(0) @binding(4) var<storage, read_write> results: array<QueryResult>;
70
71fn lod_count() -> u32 {
72    return u32(params.cascade_layout.center.w);
73}
74
75fn lod_alpha(world_xz: vec2<f32>, cascade: CascadeParams) -> f32 {
76    let offset = abs(world_xz - params.cascade_layout.center.xy);
77    // Chebyshev distance matches the square LOD rings; Euclidean `length` is wrong here.
78    let chebyshev_distance = max(offset.x, offset.y);
79    let raw_alpha = chebyshev_distance / cascade.scale - LOD_TRANSITION_START;
80    let black_point = MORPH_BLACK_POINT;
81    let fade_width = 1.0 - MORPH_FADE_SIDES * black_point;
82    return clamp((raw_alpha - black_point) / fade_width, 0.0, 1.0);
83}
84
85fn world_to_uv(world_xz: vec2<f32>, cascade: CascadeParams) -> vec2<f32> {
86    let coverage = cascade.texel_width * cascade.texture_res;
87    return (world_xz - cascade.center) / coverage + vec2(UV_CENTER);
88}
89
90fn sample_displacement(world_xz: vec2<f32>, lod: u32, alpha: f32) -> vec3<f32> {
91    let smaller = params.cascade_layout.cascades[lod];
92    let bigger = params.cascade_layout.cascades[lod + 1u];
93    let smaller_weight = (1.0 - alpha) * smaller.weight;
94    let bigger_weight = (1.0 - smaller_weight) * bigger.weight;
95    var displacement = vec3(0.0);
96
97    if smaller_weight > MIN_SAMPLE_WEIGHT {
98        let uv = world_to_uv(world_xz, smaller);
99        displacement += smaller_weight
100            * textureSampleLevel(lod_data, lod_sampler, uv, i32(lod), 0.0).xyz;
101    }
102    if bigger_weight > MIN_SAMPLE_WEIGHT {
103        let uv = world_to_uv(world_xz, bigger);
104        displacement += bigger_weight
105            * textureSampleLevel(lod_data, lod_sampler, uv, i32(lod + 1u), 0.0).xyz;
106    }
107    return displacement;
108}
109
110fn select_lod(world_xz: vec2<f32>) -> u32 {
111    // Use square rings; beyond the coarsest ring its sentinel fades to zero.
112    for (var lod = 0u; lod < lod_count(); lod += 1u) {
113        let cascade = params.cascade_layout.cascades[lod];
114        let offset = abs(world_xz - cascade.center);
115        if max(offset.x, offset.y) < cascade.scale {
116            return lod;
117        }
118    }
119    return lod_count() - 1u;
120}
121
122fn river_surface(world_xz: vec2<f32>, request: QueryRequest) -> vec3<f32> {
123    return river_analytic_displacement(
124        world_xz,
125        request.flow,
126        params.time.x,
127    );
128}
129
130@compute @workgroup_size(64u)
131fn sample(@builtin(global_invocation_id) id: vec3<u32>) {
132    let index = id.x;
133    let count = arrayLength(&requests);
134    if index >= count || index >= MAX_QUERIES {
135        return;
136    }
137    let request = requests[index];
138
139    if request.flow.w > 0.0 {
140        let analytic = river_surface(request.world_xz, request);
141        // Central differences of the same closed form for the normal.
142        let epsilon = 0.35;
143        let height_x = river_analytic_displacement(
144            request.world_xz + vec2(epsilon, 0.0),
145            request.flow,
146            params.time.x,
147        ).x;
148        let height_z = river_analytic_displacement(
149            request.world_xz + vec2(0.0, epsilon),
150            request.flow,
151            params.time.x,
152        ).x;
153        var normal_cross = vec3(
154            -(height_x - analytic.x) / epsilon,
155            1.0,
156            -(height_z - analytic.x) / epsilon,
157        );
158        normal_cross.y = max(normal_cross.y, MIN_NORMAL_Y);
159        results[index] = QueryResult(
160            vec4(vec3(0.0, analytic.x, 0.0), request.slot),
161            vec4(normalize(normal_cross), 1.0),
162            vec4(0.0),
163        );
164        return;
165    }
166
167    if request.kind > 0.5 {
168        results[index] = QueryResult(
169            vec4(0.0, 0.0, 0.0, request.slot),
170            vec4(0.0, 1.0, 0.0, 1.0),
171            vec4(0.0),
172        );
173        return;
174    }
175
176    // Mirror cascade.wgsl's advection so probes ride the same moving water
177    // the camera sees; the shift is applied once before any wave sampling.
178    let world_xz = request.world_xz - params.flow.xy * params.time.x;
179    let detail_lod = clamp(params.cascade_layout.center.z, 0.0, f32(lod_count() - 1u));
180    let minimum_lod = u32(floor(detail_lod));
181    let detail_alpha = fract(detail_lod);
182    var lod = select_lod(world_xz);
183    var alpha = lod_alpha(world_xz, params.cascade_layout.cascades[lod]);
184    if lod < minimum_lod {
185        lod = minimum_lod;
186        alpha = detail_alpha;
187    } else if lod == minimum_lod {
188        alpha = max(alpha, detail_alpha);
189    }
190
191    let displacement = sample_displacement(world_xz, lod, alpha);
192    let texel_width = params.cascade_layout.cascades[lod].texel_width;
193    let displacement_x = sample_displacement(
194        world_xz + vec2(texel_width, 0.0),
195        lod,
196        alpha,
197    );
198    let displacement_z = sample_displacement(
199        world_xz + vec2(0.0, texel_width),
200        lod,
201        alpha,
202    );
203    let tangent_x = vec3(texel_width, 0.0, 0.0) + displacement_x - displacement;
204    let tangent_z = vec3(0.0, 0.0, texel_width) + displacement_z - displacement;
205    var normal_cross = cross(tangent_z, tangent_x);
206    normal_cross.y = max(normal_cross.y, MIN_NORMAL_Y);
207    let normal = normalize(normal_cross);
208    let horizontal_x = vec2(texel_width, 0.0)
209        + displacement_x.xz - displacement.xz;
210    let horizontal_z = vec2(0.0, texel_width)
211        + displacement_z.xz - displacement.xz;
212    let determinant = (horizontal_x.x * horizontal_z.y
213        - horizontal_x.y * horizontal_z.x) / (texel_width * texel_width);
214    let crest = clamp(0.55 - determinant, 0.0, 1.0);
215
216    results[index] = QueryResult(
217        vec4(displacement, request.slot),
218        vec4(normal, 1.0),
219        vec4(crest, 0.0, 0.0, 0.0),
220    );
221}