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