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