1use crate::context::shared_wgpu_context;
4use crate::core::{
5 vertex_utils, AlphaMode, BoundingBox, DrawCall, GpuVertexBuffer, Material, PipelineType,
6 RenderData, Vertex,
7};
8use crate::gpu::stairs::StairsGpuInputs;
9use crate::gpu::util::readback_scalar_buffer_f64;
10use crate::plots::line::LineMarkerAppearance;
11use glam::{Vec3, Vec4};
12
13#[derive(Debug, Clone)]
14pub struct StairsPlot {
15 pub x: Vec<f64>,
16 pub y: Vec<f64>,
17 pub color: Vec4,
18 pub line_width: f32,
19 pub label: Option<String>,
20 pub visible: bool,
21 vertices: Option<Vec<Vertex>>,
22 bounds: Option<BoundingBox>,
23 dirty: bool,
24 gpu_vertices: Option<GpuVertexBuffer>,
25 gpu_vertex_count: Option<usize>,
26 gpu_bounds: Option<BoundingBox>,
27 gpu_inputs: Option<StairsGpuInputs>,
28 marker: Option<LineMarkerAppearance>,
29 marker_vertices: Option<Vec<Vertex>>,
30 marker_gpu_vertices: Option<GpuVertexBuffer>,
31 marker_dirty: bool,
32}
33
34impl StairsPlot {
35 pub async fn export_scene_xy_data(&self) -> Result<(Vec<f64>, Vec<f64>), String> {
36 if !self.x.is_empty() && self.x.len() == self.y.len() {
37 return Ok((self.x.clone(), self.y.clone()));
38 }
39 if !self.x.is_empty() || !self.y.is_empty() {
40 return Err(format!(
41 "stairs plot has partial CPU source data: x has {} values, y has {} values",
42 self.x.len(),
43 self.y.len()
44 ));
45 }
46
47 if let Some(inputs) = &self.gpu_inputs {
48 let context = shared_wgpu_context().ok_or_else(|| {
49 "stairs plot has GPU source data but no shared WGPU context is installed"
50 .to_string()
51 })?;
52 let len = inputs.len as usize;
53 let x = readback_scalar_buffer_f64(
54 &context.device,
55 &context.queue,
56 &inputs.x_buffer,
57 len,
58 inputs.scalar,
59 )
60 .await?;
61 let y = readback_scalar_buffer_f64(
62 &context.device,
63 &context.queue,
64 &inputs.y_buffer,
65 len,
66 inputs.scalar,
67 )
68 .await?;
69 return Ok((x, y));
70 }
71
72 if self.gpu_vertices.is_some() {
73 return Err(
74 "stairs plot has GPU render vertices but no exportable source data".to_string(),
75 );
76 }
77
78 Ok((Vec::new(), Vec::new()))
79 }
80
81 pub fn new(x: Vec<f64>, y: Vec<f64>) -> Result<Self, String> {
82 if x.len() != y.len() || x.is_empty() {
83 return Err("stairs: X and Y must be same non-zero length".to_string());
84 }
85 Ok(Self {
86 x,
87 y,
88 color: Vec4::new(0.0, 0.5, 1.0, 1.0),
89 line_width: 1.0,
90 label: None,
91 visible: true,
92 vertices: None,
93 bounds: None,
94 dirty: true,
95 gpu_vertices: None,
96 gpu_vertex_count: None,
97 gpu_bounds: None,
98 gpu_inputs: None,
99 marker: None,
100 marker_vertices: None,
101 marker_gpu_vertices: None,
102 marker_dirty: true,
103 })
104 }
105
106 pub fn from_gpu_buffer(
108 color: Vec4,
109 buffer: GpuVertexBuffer,
110 vertex_count: usize,
111 bounds: BoundingBox,
112 ) -> Self {
113 Self {
114 x: Vec::new(),
115 y: Vec::new(),
116 color,
117 line_width: 1.0,
118 label: None,
119 visible: true,
120 vertices: None,
121 bounds: None,
122 dirty: false,
123 gpu_vertices: Some(buffer),
124 gpu_vertex_count: Some(vertex_count),
125 gpu_bounds: Some(bounds),
126 gpu_inputs: None,
127 marker: None,
128 marker_vertices: None,
129 marker_gpu_vertices: None,
130 marker_dirty: true,
131 }
132 }
133
134 pub fn with_style(mut self, color: Vec4, line_width: f32) -> Self {
135 self.color = color;
136 self.line_width = line_width.max(0.5);
137 self.dirty = true;
138 self.marker_dirty = true;
139 self.drop_gpu_render_cache();
140 self
141 }
142
143 pub fn with_gpu_source_inputs(mut self, inputs: StairsGpuInputs) -> Self {
144 self.gpu_inputs = Some(inputs);
145 self
146 }
147 pub fn with_label<S: Into<String>>(mut self, label: S) -> Self {
148 self.label = Some(label.into());
149 self
150 }
151 pub fn set_visible(&mut self, v: bool) {
152 self.visible = v;
153 }
154
155 pub fn set_marker(&mut self, marker: Option<LineMarkerAppearance>) {
156 self.marker = marker;
157 self.marker_dirty = true;
158 if self.marker.is_none() {
159 self.marker_vertices = None;
160 self.marker_gpu_vertices = None;
161 }
162 }
163
164 pub fn set_marker_gpu_vertices(&mut self, buffer: Option<GpuVertexBuffer>) {
165 let has_gpu = buffer.is_some();
166 self.marker_gpu_vertices = buffer;
167 if has_gpu {
168 self.marker_vertices = None;
169 }
170 }
171
172 fn drop_gpu_render_cache(&mut self) {
173 self.gpu_vertices = None;
174 self.gpu_vertex_count = None;
175 self.gpu_bounds = None;
176 self.marker_gpu_vertices = None;
177 }
178 pub fn generate_vertices(&mut self) -> &Vec<Vertex> {
179 if self.gpu_vertices.is_some() {
180 if self.vertices.is_none() {
181 self.vertices = Some(Vec::new());
182 }
183 return self.vertices.as_ref().unwrap();
184 }
185 if self.dirty || self.vertices.is_none() {
186 let mut verts = Vec::new();
187 for i in 0..self.x.len().saturating_sub(1) {
188 let x0 = self.x[i] as f32;
189 let y0 = self.y[i] as f32;
190 let x1 = self.x[i + 1] as f32;
191 let y1 = self.y[i + 1] as f32;
192 if !x0.is_finite() || !y0.is_finite() || !x1.is_finite() || !y1.is_finite() {
193 continue;
194 }
195 verts.push(Vertex::new(Vec3::new(x0, y0, 0.0), self.color));
197 verts.push(Vertex::new(Vec3::new(x1, y0, 0.0), self.color));
198 verts.push(Vertex::new(Vec3::new(x1, y0, 0.0), self.color));
200 verts.push(Vertex::new(Vec3::new(x1, y1, 0.0), self.color));
201 }
202 self.vertices = Some(verts);
203 self.dirty = false;
204 }
205 self.vertices.as_ref().unwrap()
206 }
207 pub fn bounds(&mut self) -> BoundingBox {
208 if let Some(bounds) = self.gpu_bounds {
209 return bounds;
210 }
211 if self.dirty || self.bounds.is_none() {
212 let mut min = Vec3::new(f32::INFINITY, f32::INFINITY, 0.0);
213 let mut max = Vec3::new(f32::NEG_INFINITY, f32::NEG_INFINITY, 0.0);
214 for (&x, &y) in self.x.iter().zip(self.y.iter()) {
215 let (x, y) = (x as f32, y as f32);
216 if !x.is_finite() || !y.is_finite() {
217 continue;
218 }
219 min.x = min.x.min(x);
220 max.x = max.x.max(x);
221 min.y = min.y.min(y);
222 max.y = max.y.max(y);
223 }
224 if !min.x.is_finite() {
225 min = Vec3::ZERO;
226 max = Vec3::ZERO;
227 }
228 self.bounds = Some(BoundingBox::new(min, max));
229 }
230 self.bounds.unwrap()
231 }
232 pub fn render_data(&mut self) -> RenderData {
233 let using_gpu = self.gpu_vertices.is_some();
234 let bounds = self.bounds();
235 let (vertices, vertex_count, gpu_vertices) = if using_gpu {
236 (
237 Vec::new(),
238 self.gpu_vertex_count.unwrap_or(0),
239 self.gpu_vertices.clone(),
240 )
241 } else {
242 let verts = self.generate_vertices().clone();
243 let count = verts.len();
244 (verts, count, None)
245 };
246 let material = Material {
247 albedo: self.color,
248 ..Default::default()
249 };
250 let draw_call = DrawCall {
251 vertex_offset: 0,
252 vertex_count,
253 index_offset: None,
254 index_count: None,
255 instance_count: 1,
256 };
257 RenderData {
258 pipeline_type: PipelineType::Lines,
259 vertices,
260 indices: None,
261 gpu_vertices,
262 bounds: Some(bounds),
263 material,
264 draw_calls: vec![draw_call],
265 image: None,
266 }
267 }
268
269 pub fn render_data_with_viewport(&mut self, viewport_px: Option<(u32, u32)>) -> RenderData {
270 if self.gpu_vertices.is_some() {
271 return self.render_data();
272 }
273
274 let bounds = self.bounds();
275 let (vertices, vertex_count, pipeline_type) = if self.line_width > 1.0 {
276 let Some(viewport_px) = viewport_px else {
277 return self.render_data();
278 };
279 let data_per_px = crate::core::data_units_per_px(&bounds, viewport_px);
280 let width_data = self.line_width.max(0.1) * data_per_px;
281 let verts = self.generate_vertices().clone();
282 let mut thick = Vec::new();
283 for segment in verts.chunks_exact(2) {
284 let x = [segment[0].position[0] as f64, segment[1].position[0] as f64];
285 let y = [segment[0].position[1] as f64, segment[1].position[1] as f64];
286 let color = Vec4::from_array(segment[0].color);
287 thick.extend(vertex_utils::create_thick_polyline(
288 &x, &y, color, width_data,
289 ));
290 }
291 let count = thick.len();
292 (thick, count, PipelineType::Triangles)
293 } else {
294 let verts = self.generate_vertices().clone();
295 let count = verts.len();
296 (verts, count, PipelineType::Lines)
297 };
298 let material = Material {
299 albedo: self.color,
300 roughness: self.line_width.max(0.0),
301 ..Default::default()
302 };
303 let draw_call = DrawCall {
304 vertex_offset: 0,
305 vertex_count,
306 index_offset: None,
307 index_count: None,
308 instance_count: 1,
309 };
310 RenderData {
311 pipeline_type,
312 vertices,
313 indices: None,
314 gpu_vertices: None,
315 bounds: Some(bounds),
316 material,
317 draw_calls: vec![draw_call],
318 image: None,
319 }
320 }
321
322 pub fn marker_render_data(&mut self) -> Option<RenderData> {
323 let marker = self.marker.clone()?;
324 if let Some(gpu_vertices) = self.marker_gpu_vertices.clone() {
325 let vertex_count = gpu_vertices.vertex_count;
326 if vertex_count == 0 {
327 return None;
328 }
329 let draw_call = DrawCall {
330 vertex_offset: 0,
331 vertex_count,
332 index_offset: None,
333 index_count: None,
334 instance_count: 1,
335 };
336 let material = Self::marker_material(&marker);
337 return Some(RenderData {
338 pipeline_type: PipelineType::Points,
339 vertices: Vec::new(),
340 indices: None,
341 gpu_vertices: Some(gpu_vertices),
342 bounds: None,
343 material,
344 draw_calls: vec![draw_call],
345 image: None,
346 });
347 }
348
349 let vertices = self.marker_vertices_slice(&marker)?;
350 if vertices.is_empty() {
351 return None;
352 }
353 let draw_call = DrawCall {
354 vertex_offset: 0,
355 vertex_count: vertices.len(),
356 index_offset: None,
357 index_count: None,
358 instance_count: 1,
359 };
360 let material = Self::marker_material(&marker);
361 Some(RenderData {
362 pipeline_type: PipelineType::Points,
363 vertices: vertices.to_vec(),
364 indices: None,
365 gpu_vertices: None,
366 bounds: None,
367 material,
368 draw_calls: vec![draw_call],
369 image: None,
370 })
371 }
372
373 fn marker_material(marker: &LineMarkerAppearance) -> Material {
374 let mut material = Material {
375 albedo: marker.face_color,
376 ..Default::default()
377 };
378 if !marker.filled {
379 material.albedo.w = 0.0;
380 }
381 material.emissive = marker.edge_color;
382 material.roughness = 1.0;
383 material.metallic = 0.0;
384 material.alpha_mode = AlphaMode::Blend;
385 material
386 }
387
388 fn marker_vertices_slice(&mut self, marker: &LineMarkerAppearance) -> Option<&[Vertex]> {
389 if self.x.len() != self.y.len() || self.x.is_empty() {
390 return None;
391 }
392 if self.marker_vertices.is_none() || self.marker_dirty {
393 let mut verts = Vec::with_capacity(self.x.len());
394 for (&x, &y) in self.x.iter().zip(self.y.iter()) {
395 let mut vertex = Vertex::new(Vec3::new(x as f32, y as f32, 0.0), marker.face_color);
396 vertex.normal[2] = marker.size.max(1.0);
397 verts.push(vertex);
398 }
399 self.marker_vertices = Some(verts);
400 self.marker_dirty = false;
401 }
402 self.marker_vertices.as_deref()
403 }
404 pub fn estimated_memory_usage(&self) -> usize {
405 self.vertices
406 .as_ref()
407 .map_or(0, |v| v.len() * std::mem::size_of::<Vertex>())
408 }
409}
410
411#[cfg(test)]
412mod tests {
413 use super::*;
414
415 #[test]
416 fn thick_stairs_use_viewport_aware_triangles() {
417 let mut plot = StairsPlot::new(vec![0.0, 1.0, 2.0], vec![1.0, 2.0, 1.5])
418 .unwrap()
419 .with_style(Vec4::ONE, 2.0);
420 let render = plot.render_data_with_viewport(Some((600, 400)));
421 assert_eq!(render.pipeline_type, PipelineType::Triangles);
422 assert!(!render.vertices.is_empty());
423 }
424}