1#import bevy_core_pipeline::fullscreen_vertex_shader::FullscreenVertexOutput
2
3struct JumpFloodUniform {
4    step_length: u32,
5    _padding: vec3<f32>
6}
7
8@group(0) @binding(0) var screen_texture: texture_2d<f32>;
9@group(0) @binding(1) var texture_sampler: sampler;
10@group(0) @binding(2) var<uniform> instance: JumpFloodUniform;
11
12@fragment
13fn fragment(in: FullscreenVertexOutput) -> @location(0) vec4<f32> {
14    let dims = vec2<f32>(textureDimensions(screen_texture));
15    let step = i32(instance.step_length);
16
17    let current = textureSample(screen_texture, texture_sampler, in.uv);
18    var closest_dist = distance(vec2<f32>(in.position.xy), vec2<f32>(current.xy));
19    var result = current;
20
21    // Check all 8 neighbouring pixels at current step distance
22    for (var dy = -1; dy <= 1; dy++) {
23        for (var dx = -1; dx <= 1; dx++) {
24            if (dx == 0 && dy == 0) {
25                continue;
26            }
27
28            let neighbour_coord = in.position.xy + vec2<f32>(f32(dx * step), f32(dy * step));
29            let neighbour = textureSample(screen_texture, texture_sampler, neighbour_coord / dims);
30            let dist = distance(vec2<f32>(in.position.xy), vec2<f32>(neighbour.xy));
31
32            if (dist < closest_dist) {
33                closest_dist = dist;
34                result = neighbour;
35            }
36        }
37    }
38    return result;
39}