1import bevy_render::view::View;
2import bevy_core_pipeline::fullscreen_vertex_shader::FullscreenVertexOutput;
3
4struct InfiniteGridUniform {
5    rot_matrix: mat3x3<f32>,
6    origin: vec3<f32>,
7    normal: vec3<f32>,
8};
9
10struct InfiniteGridSettings {
11    scale: f32,
12    /// 1 / fadeout_distance
13    one_over_fadeout_distance: f32,
14    /// 1 / dot_fadeout_strength
15    one_over_dot_fadeout: f32,
16    /// 1 / major_line_interval
17    one_over_major_line_interval: f32,
18    x_axis_col: vec3<f32>,
19    z_axis_col: vec3<f32>,
20    minor_line_col: vec4<f32>,
21    major_line_col: vec4<f32>,
22};
23
24@group(0) @binding(0) var<uniform> view: View;
25
26@group(1) @binding(0) var<uniform> infinite_grid: InfiniteGridUniform;
27@group(1) @binding(1) var<uniform> grid_settings: InfiniteGridSettings;
28
29/// Same as view_transformations::position_ndc_to_world but we can't use it since
30/// it relies on bevy's view bind group
31fn position_ndc_to_world(ndc_pos: vec3<f32>) -> vec3<f32> {
32    let world_pos = view.world_from_clip * vec4(ndc_pos, 1.0);
33    return world_pos.xyz / world_pos.w;
34}
35
36struct FragmentOutput {
37    @location(0) color: vec4<f32>,
38    @builtin(frag_depth) depth: f32,
39};
40
41@fragment
42fn fragment(in: FullscreenVertexOutput) -> FragmentOutput {
43    // Reconstruct clip-space XY from UV, then unproject to world space
44    let clip_xy = in.uv * vec2(2.0, -2.0) + vec2(-1.0, 1.0);
45    let near_point = position_ndc_to_world(vec3(clip_xy, 1.0));
46    let far_point = position_ndc_to_world(vec3(clip_xy, 0.001));
47
48    // Cast a ray from the near plane towards the far plane
49    let ray_origin = near_point;
50    let ray_direction = normalize(far_point - near_point);
51    let plane_normal = infinite_grid.normal;
52    let plane_origin = infinite_grid.origin;
53
54    // Ray-plane intersection
55    // t is the signed distance to the plane
56    let point_to_point = plane_origin - ray_origin;
57    let t = dot(plane_normal, point_to_point) / dot(ray_direction, plane_normal);
58    let frag_pos_3d = ray_direction * t + ray_origin;
59
60    // Project the 3D hit point into the grid's local 2D coordinate space so
61    // that grid lines are always axis-aligned regardless of the grid's rotation.
62    let planar_offset = frag_pos_3d - plane_origin;
63    let rotation_matrix = infinite_grid.rot_matrix;
64    let plane_coords = (rotation_matrix * planar_offset).xz;
65
66    // Compute the clip-space depth of the hit point to make the lines opaque
67    let view_space_pos = view.view_from_world * vec4(frag_pos_3d, 1.);
68    let clip_space_pos = view.clip_from_view * view_space_pos;
69    let clip_depth = clip_space_pos.z / clip_space_pos.w;
70    let real_depth = -view_space_pos.z;
71
72    var out: FragmentOutput;
73    out.depth = clip_depth;
74
75    // Scale the plane coordinates to control the size of the grid
76    let scale = grid_settings.scale;
77    let coord = plane_coords * scale;
78
79    // Compute minor lines
80    let derivative = fwidth(coord);
81    let grid = abs(fract(coord - 0.5) - 0.5) / derivative;
82    let minor_line = min(grid.x, grid.y);
83
84    // Compute major lines
85    let major_scale = grid_settings.one_over_major_line_interval;
86    let derivative2 = fwidth(coord * major_scale);
87    let grid2 = abs(fract((coord * major_scale) - 0.5) - 0.5) / derivative2;
88    let major_line = min(grid2.x, grid2.y);
89
90    // Compute axis lines
91    let grid3 = abs(coord) / derivative;
92    let axis_line = min(grid3.x, grid3.y);
93
94    // Compute the alpha based on the priority of the lines, axis > major > minor
95    var alpha = vec3(1.0) - min(vec3(axis_line, major_line, minor_line), vec3(1.0));
96    alpha.y *= (1.0 - alpha.x) * grid_settings.major_line_col.a;
97    alpha.z *= (1.0 - (alpha.x + alpha.y)) * grid_settings.minor_line_col.a;
98
99    // Fade the grid linearly with distance from the camera.
100    let dist_fadeout = min(1., 1. - grid_settings.one_over_fadeout_distance * real_depth);
101    let dot_fadeout = abs(dot(infinite_grid.normal, normalize(view.world_position - frag_pos_3d)));
102    let alpha_fadeout = mix(dist_fadeout, 1., dot_fadeout)
103        * min(grid_settings.one_over_dot_fadeout * dot_fadeout, 1.);
104
105    // Normalize the alpha
106    let a_0 = alpha.x + alpha.y + alpha.z;
107    alpha /= a_0;
108    // In case a_0 is 0 we need to clamp the result to avoid NaNs
109    alpha = clamp(alpha, vec3(0.0), vec3(1.0));
110
111    // Choose the axis color based on which axis this fragment is closest to
112    let axis_color = mix(
113        grid_settings.x_axis_col,
114        grid_settings.z_axis_col,
115        step(grid3.x, grid3.y)
116    );
117
118    var grid_color = vec4(
119        axis_color * alpha.x
120            + grid_settings.major_line_col.rgb * alpha.y
121            + grid_settings.minor_line_col.rgb * alpha.z,
122        max(a_0 * alpha_fadeout, 0.0),
123    );
124    out.color = grid_color;
125
126    return out;
127}