1#import bevy_render::{view::View, maths::affine3_to_square}
2
3@group(0) @binding(0) var<uniform> view: View;
4
5
6struct LineGizmoUniform {
7    world_from_local: mat3x4<f32>,
8    line_width: f32,
9    depth_bias: f32,
10    resolution: u32,
11#ifdef SIXTEEN_BYTE_ALIGNMENT
12    // WebGL2 structs must be 16 byte aligned.
13    _padding: f32,
14#endif
15}
16
17@group(1) @binding(0) var<uniform> joints_gizmo: LineGizmoUniform;
18
19struct VertexInput {
20    @location(0) position_a: vec3<f32>,
21    @location(1) position_b: vec3<f32>,
22    @location(2) position_c: vec3<f32>,
23    @location(3) color: vec4<f32>,
24    @builtin(vertex_index) index: u32,
25};
26
27struct VertexOutput {
28    @builtin(position) clip_position: vec4<f32>,
29    @location(0) color: vec4<f32>,
30};
31
32const EPSILON: f32 = 4.88e-04;
33
34@vertex
35fn vertex_bevel(vertex: VertexInput) -> VertexOutput {
36    var positions = array<vec2<f32>, 3>(
37        vec2(0, 0),
38        vec2(0, 0.5),
39        vec2(0.5, 0),
40    );
41    var position = positions[vertex.index];
42
43    let world_from_local = affine3_to_square(joints_gizmo.world_from_local);
44
45    var clip_a = view.clip_from_world * world_from_local * vec4(vertex.position_a, 1.);
46    var clip_b = view.clip_from_world * world_from_local * vec4(vertex.position_b, 1.);
47    var clip_c = view.clip_from_world * world_from_local * vec4(vertex.position_c, 1.);
48
49    // Manual near plane clipping to avoid errors when doing the perspective divide inside this shader.
50    clip_a = clip_near_plane(clip_a, clip_c);
51    clip_b = clip_near_plane(clip_b, clip_a);
52    clip_c = clip_near_plane(clip_c, clip_b);
53    clip_a = clip_near_plane(clip_a, clip_c);
54
55    let resolution = view.viewport.zw;
56    let screen_a = resolution * (0.5 * clip_a.xy / clip_a.w + 0.5);
57    let screen_b = resolution * (0.5 * clip_b.xy / clip_b.w + 0.5);
58    let screen_c = resolution * (0.5 * clip_c.xy / clip_c.w + 0.5);
59
60    var color = vertex.color;
61    var line_width = joints_gizmo.line_width;
62
63#ifdef PERSPECTIVE
64    line_width /= clip_b.w;
65#endif
66
67    // Line thinness fade from https://acegikmo.com/shapes/docs/#anti-aliasing
68    if line_width > 0.0 && line_width < 1. {
69        color.a *= line_width;
70        line_width = 1.;
71    }
72
73    let ab = normalize(screen_b - screen_a);
74    let cb = normalize(screen_b - screen_c);
75    let ab_norm = vec2(-ab.y, ab.x);
76    let cb_norm = vec2(cb.y, -cb.x);
77    let tangent = normalize(ab - cb);
78    let normal = vec2(-tangent.y, tangent.x);
79    let sigma = sign(dot(ab + cb, normal));
80
81    var p0 = line_width * sigma * ab_norm;
82    var p1 = line_width * sigma * cb_norm;
83
84    let screen = screen_b + position.x * p0 + position.y * p1;
85
86    let depth = depth(clip_b);
87
88    var clip_position = vec4(clip_b.w * ((2. * screen) / resolution - 1.), depth, clip_b.w);
89    return VertexOutput(clip_position, color);
90}
91
92@vertex
93fn vertex_miter(vertex: VertexInput) -> VertexOutput {
94    var positions = array<vec3<f32>, 6>(
95        vec3(0, 0, 0),
96        vec3(0.5, 0, 0),
97        vec3(0, 0.5, 0),
98        vec3(0, 0, 0),
99        vec3(0, 0.5, 0),
100        vec3(0, 0, 0.5),
101    );
102    var position = positions[vertex.index];
103
104    let world_from_local = affine3_to_square(joints_gizmo.world_from_local);
105
106    var clip_a = view.clip_from_world * world_from_local * vec4(vertex.position_a, 1.);
107    var clip_b = view.clip_from_world * world_from_local * vec4(vertex.position_b, 1.);
108    var clip_c = view.clip_from_world * world_from_local * vec4(vertex.position_c, 1.);
109
110    // Manual near plane clipping to avoid errors when doing the perspective divide inside this shader.
111    clip_a = clip_near_plane(clip_a, clip_c);
112    clip_b = clip_near_plane(clip_b, clip_a);
113    clip_c = clip_near_plane(clip_c, clip_b);
114    clip_a = clip_near_plane(clip_a, clip_c);
115
116    let resolution = view.viewport.zw;
117    let screen_a = resolution * (0.5 * clip_a.xy / clip_a.w + 0.5);
118    let screen_b = resolution * (0.5 * clip_b.xy / clip_b.w + 0.5);
119    let screen_c = resolution * (0.5 * clip_c.xy / clip_c.w + 0.5);
120
121    var color = vertex.color;
122    var line_width = joints_gizmo.line_width;
123
124#ifdef PERSPECTIVE
125    line_width /= clip_b.w;
126#endif
127
128    // Line thinness fade from https://acegikmo.com/shapes/docs/#anti-aliasing
129    if line_width > 0.0 && line_width < 1. {
130        color.a *= line_width;
131        line_width = 1.;
132    }
133
134    let ab = normalize(screen_b - screen_a);
135    let cb = normalize(screen_b - screen_c);
136    let ab_norm = vec2(-ab.y, ab.x);
137    let cb_norm = vec2(cb.y, -cb.x);
138    let tangent = normalize(ab - cb);
139    let normal = vec2(-tangent.y, tangent.x);
140    let sigma = sign(dot(ab + cb, normal));
141
142    var p0 = line_width * sigma * ab_norm;
143    var p1 = line_width * sigma * normal / dot(normal, ab_norm);
144    var p2 = line_width * sigma * cb_norm;
145    
146    var screen = screen_b + position.x * p0 + position.y * p1 + position.z * p2;
147
148    var depth = depth(clip_b);
149
150    var clip_position = vec4(clip_b.w * ((2. * screen) / resolution - 1.), depth, clip_b.w);
151    return VertexOutput(clip_position, color);
152}
153
154@vertex
155fn vertex_round(vertex: VertexInput) -> VertexOutput {
156    let world_from_local = affine3_to_square(joints_gizmo.world_from_local);
157
158    var clip_a = view.clip_from_world * world_from_local * vec4(vertex.position_a, 1.);
159    var clip_b = view.clip_from_world * world_from_local * vec4(vertex.position_b, 1.);
160    var clip_c = view.clip_from_world * world_from_local * vec4(vertex.position_c, 1.);
161
162    // Manual near plane clipping to avoid errors when doing the perspective divide inside this shader.
163    clip_a = clip_near_plane(clip_a, clip_c);
164    clip_b = clip_near_plane(clip_b, clip_a);
165    clip_c = clip_near_plane(clip_c, clip_b);
166    clip_a = clip_near_plane(clip_a, clip_c);
167
168    let resolution = view.viewport.zw;
169    let screen_a = resolution * (0.5 * clip_a.xy / clip_a.w + 0.5);
170    let screen_b = resolution * (0.5 * clip_b.xy / clip_b.w + 0.5);
171    let screen_c = resolution * (0.5 * clip_c.xy / clip_c.w + 0.5);
172
173    var color = vertex.color;
174    var line_width = joints_gizmo.line_width;
175
176#ifdef PERSPECTIVE
177    line_width /= clip_b.w;
178#endif
179
180    // Line thinness fade from https://acegikmo.com/shapes/docs/#anti-aliasing
181    if line_width > 0.0 && line_width < 1. {
182        color.a *= line_width;
183        line_width = 1.;
184    }
185
186    let ab = normalize(screen_b - screen_a);
187    let cb = normalize(screen_b - screen_c);
188    let ab_norm = vec2(-ab.y, ab.x);
189    let cb_norm = vec2(cb.y, -cb.x);
190
191    // We render `joints_gizmo.resolution`triangles. The vertices in each triangle are ordered as follows:
192    // - 0: The 'center' vertex at `screen_b`.
193    // - 1: The vertex closer to the ab line.
194    // - 2: The vertex closer to the cb line. 
195    var in_triangle_index = f32(vertex.index) % 3.0;
196    var tri_index = floor(f32(vertex.index) / 3.0);
197    var radius = sign(in_triangle_index) * 0.5 * line_width;
198    var theta = acos(dot(ab_norm, cb_norm));
199    let sigma = sign(dot(ab_norm, cb));
200    var angle = theta * (tri_index + in_triangle_index - 1) / f32(joints_gizmo.resolution);
201    var position_x = sigma * radius * cos(angle);
202    var position_y = radius * sin(angle);
203
204    var screen = screen_b + position_x * ab_norm + position_y * ab;
205
206    var depth = depth(clip_b);
207
208    var clip_position = vec4(clip_b.w * ((2. * screen) / resolution - 1.), depth, clip_b.w);
209    return VertexOutput(clip_position, color);
210}
211
212fn clip_near_plane(a: vec4<f32>, b: vec4<f32>) -> vec4<f32> {
213    // Move a if a is behind the near plane and b is in front. 
214    if a.z > a.w && b.z <= b.w {
215        // Interpolate a towards b until it's at the near plane.
216        let distance_a = a.z - a.w;
217        let distance_b = b.z - b.w;
218        // Add an epsilon to the interpolator to ensure that the point is
219        // not just behind the clip plane due to floating-point imprecision.
220        let t = distance_a / (distance_a - distance_b) + EPSILON;
221        return mix(a, b, t);
222    }
223    return a;
224}
225
226fn depth(clip: vec4<f32>) -> f32 {
227    var depth: f32;
228    if joints_gizmo.depth_bias >= 0. {
229        depth = clip.z * (1. - joints_gizmo.depth_bias);
230    } else {
231        // depth * (clip.w / depth)^-depth_bias. So that when -depth_bias is 1.0, this is equal to clip.w
232        // and when equal to 0.0, it is exactly equal to depth.
233        // the epsilon is here to prevent the depth from exceeding clip.w when -depth_bias = 1.0
234        // clip.w represents the near plane in homogeneous clip space in bevy, having a depth
235        // of this value means nothing can be in front of this
236        // The reason this uses an exponential function is that it makes it much easier for the
237        // user to chose a value that is convenient for them
238        depth = clip.z * exp2(-joints_gizmo.depth_bias * log2(clip.w / clip.z - EPSILON));
239    }
240    return depth;
241}
242
243struct FragmentInput {
244    @location(0) color: vec4<f32>,
245};
246
247struct FragmentOutput {
248    @location(0) color: vec4<f32>,
249};
250
251@fragment
252fn fragment(in: FragmentInput) -> FragmentOutput {
253    // return FragmentOutput(vec4(1, 1, 1, 1));
254    return FragmentOutput(in.color);
255}