Skip to main content

runmat_plot/gpu/shaders/
line3.rs

1pub const F32: &str = r#"const WORKGROUP_SIZE: u32 = {{WORKGROUP_SIZE}}u;
2
3struct VertexRaw { data: array<f32, 12u>, };
4struct Line3Params { color: vec4<f32>, count: u32, half_width_data: f32, line_style: u32, thick: u32, _pad: u32, };
5
6@group(0) @binding(0) var<storage, read> buf_x: array<f32>;
7@group(0) @binding(1) var<storage, read> buf_y: array<f32>;
8@group(0) @binding(2) var<storage, read> buf_z: array<f32>;
9@group(0) @binding(3) var<storage, read_write> out_vertices: array<VertexRaw>;
10@group(0) @binding(4) var<uniform> params: Line3Params;
11
12fn should_draw(segment: u32, style: u32) -> bool {
13  switch(style) {
14    case 0u: { return true; }
15    case 1u: { return (segment % 4u) < 2u; }
16    case 2u: { return (segment % 4u) == 0u; }
17    case 3u: { let m = segment % 6u; return (m < 2u) || (m == 3u); }
18    default: { return true; }
19  }
20}
21
22fn safe_normalize(v: vec3<f32>) -> vec3<f32> {
23  let len = length(v);
24  if (len < 0.000001) {
25    return vec3<f32>(1.0, 0.0, 0.0);
26  }
27  return v / len;
28}
29
30fn point(i: u32) -> vec3<f32> {
31  return vec3<f32>(buf_x[i], buf_y[i], buf_z[i]);
32}
33
34fn segment_dir(i0: u32, i1: u32) -> vec3<f32> {
35  let d = point(i1) - point(i0);
36  let len = length(d);
37  if (len < 0.000001) {
38    return vec3<f32>(0.0, 0.0, 0.0);
39  }
40  return d / len;
41}
42
43fn side_at(i: u32) -> vec3<f32> {
44  var prev = vec3<f32>(0.0, 0.0, 0.0);
45  var next = vec3<f32>(0.0, 0.0, 0.0);
46  var has_prev = false;
47  var has_next = false;
48
49  if (i > 0u) {
50    let d = segment_dir(i - 1u, i);
51    if (length(d) > 0.0) {
52      prev = d;
53      has_prev = true;
54    }
55  }
56
57  if (i + 1u < params.count) {
58    let d = segment_dir(i, i + 1u);
59    if (length(d) > 0.0) {
60      next = d;
61      has_next = true;
62    }
63  }
64
65  var tangent = vec3<f32>(1.0, 0.0, 0.0);
66  if (has_prev && has_next) {
67    let s = prev + next;
68    if (length(s) > 0.000001) {
69      tangent = normalize(s);
70    } else {
71      tangent = next;
72    }
73  } else if (has_prev) {
74    tangent = prev;
75  } else if (has_next) {
76    tangent = next;
77  }
78
79  let ref_axis = select(
80    vec3<f32>(1.0, 0.0, 0.0),
81    vec3<f32>(0.0, 0.0, 1.0),
82    abs(tangent.z) < 0.95,
83  );
84  var side = cross(tangent, ref_axis);
85  if (length(side) < 0.000001) {
86    side = cross(tangent, vec3<f32>(0.0, 1.0, 0.0));
87  }
88  if (length(side) < 0.000001) {
89    return vec3<f32>(0.0, 1.0, 0.0);
90  }
91  return normalize(side);
92}
93
94fn write_line_vertices(base: u32, p0: vec3<f32>, p1: vec3<f32>, color: vec4<f32>) {
95  write_vertex(base + 0u, p0, color);
96  write_vertex(base + 1u, p1, color);
97}
98
99fn write_thick_vertices(segment: u32, base: u32, p0: vec3<f32>, p1: vec3<f32>, color: vec4<f32>, half_width: f32) {
100  let dir = safe_normalize(p1 - p0);
101  let side0 = side_at(segment);
102  let side1 = side_at(segment + 1u);
103  let ext = dir * half_width;
104  let a = p0 - ext;
105  let b = p1 + ext;
106  let v0 = a + side0 * half_width;
107  let v1 = b + side1 * half_width;
108  let v2 = b - side1 * half_width;
109  let v3 = a - side0 * half_width;
110  write_vertex(base + 0u, v0, color);
111  write_vertex(base + 1u, v1, color);
112  write_vertex(base + 2u, v2, color);
113  write_vertex(base + 3u, v0, color);
114  write_vertex(base + 4u, v2, color);
115  write_vertex(base + 5u, v3, color);
116}
117
118fn write_vertex(index: u32, pos: vec3<f32>, color: vec4<f32>) {
119  var vertex: VertexRaw;
120  vertex.data[0u] = pos.x; vertex.data[1u] = pos.y; vertex.data[2u] = pos.z;
121  vertex.data[3u] = color.x; vertex.data[4u] = color.y; vertex.data[5u] = color.z; vertex.data[6u] = color.w;
122  vertex.data[7u] = 0.0; vertex.data[8u] = 0.0; vertex.data[9u] = 1.0; vertex.data[10u] = 0.0; vertex.data[11u] = 0.0;
123  out_vertices[index] = vertex;
124}
125
126@compute @workgroup_size(WORKGROUP_SIZE)
127fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
128  if (params.count < 2u) { return; }
129  let segments = params.count - 1u;
130  let idx = gid.x;
131  if (idx >= segments) { return; }
132
133  var color = params.color;
134  if (!should_draw(idx, params.line_style)) {
135    color.w = 0.0;
136  }
137
138  let p0 = point(idx);
139  let p1 = point(idx + 1u);
140  if (distance(p0, p1) < 0.000001) {
141    color.w = 0.0;
142  }
143
144  if (params.thick != 0u) {
145    let base = idx * 6u;
146    write_thick_vertices(idx, base, p0, p1, color, params.half_width_data);
147  } else {
148    let base = idx * 2u;
149    write_line_vertices(base, p0, p1, color);
150  }
151}
152"#;
153
154pub const F64: &str = r#"const WORKGROUP_SIZE: u32 = {{WORKGROUP_SIZE}}u;
155
156struct VertexRaw { data: array<f32, 12u>, };
157struct Line3Params { color: vec4<f32>, count: u32, half_width_data: f32, line_style: u32, thick: u32, _pad: u32, };
158
159@group(0) @binding(0) var<storage, read> buf_x: array<f64>;
160@group(0) @binding(1) var<storage, read> buf_y: array<f64>;
161@group(0) @binding(2) var<storage, read> buf_z: array<f64>;
162@group(0) @binding(3) var<storage, read_write> out_vertices: array<VertexRaw>;
163@group(0) @binding(4) var<uniform> params: Line3Params;
164
165fn should_draw(segment: u32, style: u32) -> bool {
166  switch(style) {
167    case 0u: { return true; }
168    case 1u: { return (segment % 4u) < 2u; }
169    case 2u: { return (segment % 4u) == 0u; }
170    case 3u: { let m = segment % 6u; return (m < 2u) || (m == 3u); }
171    default: { return true; }
172  }
173}
174
175fn safe_normalize(v: vec3<f32>) -> vec3<f32> {
176  let len = length(v);
177  if (len < 0.000001) {
178    return vec3<f32>(1.0, 0.0, 0.0);
179  }
180  return v / len;
181}
182
183fn point(i: u32) -> vec3<f32> {
184  return vec3<f32>(f32(buf_x[i]), f32(buf_y[i]), f32(buf_z[i]));
185}
186
187fn segment_dir(i0: u32, i1: u32) -> vec3<f32> {
188  let d = point(i1) - point(i0);
189  let len = length(d);
190  if (len < 0.000001) {
191    return vec3<f32>(0.0, 0.0, 0.0);
192  }
193  return d / len;
194}
195
196fn side_at(i: u32) -> vec3<f32> {
197  var prev = vec3<f32>(0.0, 0.0, 0.0);
198  var next = vec3<f32>(0.0, 0.0, 0.0);
199  var has_prev = false;
200  var has_next = false;
201
202  if (i > 0u) {
203    let d = segment_dir(i - 1u, i);
204    if (length(d) > 0.0) {
205      prev = d;
206      has_prev = true;
207    }
208  }
209
210  if (i + 1u < params.count) {
211    let d = segment_dir(i, i + 1u);
212    if (length(d) > 0.0) {
213      next = d;
214      has_next = true;
215    }
216  }
217
218  var tangent = vec3<f32>(1.0, 0.0, 0.0);
219  if (has_prev && has_next) {
220    let s = prev + next;
221    if (length(s) > 0.000001) {
222      tangent = normalize(s);
223    } else {
224      tangent = next;
225    }
226  } else if (has_prev) {
227    tangent = prev;
228  } else if (has_next) {
229    tangent = next;
230  }
231
232  let ref_axis = select(
233    vec3<f32>(1.0, 0.0, 0.0),
234    vec3<f32>(0.0, 0.0, 1.0),
235    abs(tangent.z) < 0.95,
236  );
237  var side = cross(tangent, ref_axis);
238  if (length(side) < 0.000001) {
239    side = cross(tangent, vec3<f32>(0.0, 1.0, 0.0));
240  }
241  if (length(side) < 0.000001) {
242    return vec3<f32>(0.0, 1.0, 0.0);
243  }
244  return normalize(side);
245}
246
247fn write_line_vertices(base: u32, p0: vec3<f32>, p1: vec3<f32>, color: vec4<f32>) {
248  write_vertex(base + 0u, p0, color);
249  write_vertex(base + 1u, p1, color);
250}
251
252fn write_thick_vertices(segment: u32, base: u32, p0: vec3<f32>, p1: vec3<f32>, color: vec4<f32>, half_width: f32) {
253  let dir = safe_normalize(p1 - p0);
254  let side0 = side_at(segment);
255  let side1 = side_at(segment + 1u);
256  let ext = dir * half_width;
257  let a = p0 - ext;
258  let b = p1 + ext;
259  let v0 = a + side0 * half_width;
260  let v1 = b + side1 * half_width;
261  let v2 = b - side1 * half_width;
262  let v3 = a - side0 * half_width;
263  write_vertex(base + 0u, v0, color);
264  write_vertex(base + 1u, v1, color);
265  write_vertex(base + 2u, v2, color);
266  write_vertex(base + 3u, v0, color);
267  write_vertex(base + 4u, v2, color);
268  write_vertex(base + 5u, v3, color);
269}
270
271fn write_vertex(index: u32, pos: vec3<f32>, color: vec4<f32>) {
272  var vertex: VertexRaw;
273  vertex.data[0u] = pos.x; vertex.data[1u] = pos.y; vertex.data[2u] = pos.z;
274  vertex.data[3u] = color.x; vertex.data[4u] = color.y; vertex.data[5u] = color.z; vertex.data[6u] = color.w;
275  vertex.data[7u] = 0.0; vertex.data[8u] = 0.0; vertex.data[9u] = 1.0; vertex.data[10u] = 0.0; vertex.data[11u] = 0.0;
276  out_vertices[index] = vertex;
277}
278
279@compute @workgroup_size(WORKGROUP_SIZE)
280fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
281  if (params.count < 2u) { return; }
282  let segments = params.count - 1u;
283  let idx = gid.x;
284  if (idx >= segments) { return; }
285
286  var color = params.color;
287  if (!should_draw(idx, params.line_style)) {
288    color.w = 0.0;
289  }
290
291  let p0 = point(idx);
292  let p1 = point(idx + 1u);
293  if (distance(p0, p1) < 0.000001) {
294    color.w = 0.0;
295  }
296
297  if (params.thick != 0u) {
298    let base = idx * 6u;
299    write_thick_vertices(idx, base, p0, p1, color, params.half_width_data);
300  } else {
301    let base = idx * 2u;
302    write_line_vertices(base, p0, p1, color);
303  }
304}
305"#;