Skip to main content

runmat_plot/plots/
stairs.rs

1//! Stairs (step) plot implementation
2
3use 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    /// Build a stairs plot backed directly by a GPU vertex buffer.
107    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                // Horizontal segment
196                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                // Vertical jump
199                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}