Skip to main content

dear_implot/plots/
stems.rs

1//! Stem plot implementation
2
3use super::{
4    Plot, PlotDataLayout, PlotDataOffset, PlotDataStride, PlotError, PlotItemStyle,
5    plot_spec_with_style, validate_data_lengths, with_plot_str_or_empty,
6};
7use crate::{ItemFlags, StemsFlags, sys};
8
9/// Builder for stem plots (lollipop charts)
10pub struct StemPlot<'a> {
11    label: &'a str,
12    x_data: &'a [f64],
13    y_data: &'a [f64],
14    style: PlotItemStyle,
15    y_ref: f64,
16    flags: StemsFlags,
17    item_flags: ItemFlags,
18    layout: PlotDataLayout,
19}
20
21impl<'a> super::PlotItemStyled for StemPlot<'a> {
22    fn style_mut(&mut self) -> &mut PlotItemStyle {
23        &mut self.style
24    }
25}
26
27impl<'a> StemPlot<'a> {
28    /// Create a new stem plot with the given label and data
29    pub fn new(label: &'a str, x_data: &'a [f64], y_data: &'a [f64]) -> Self {
30        Self {
31            label,
32            x_data,
33            y_data,
34            style: PlotItemStyle::default(),
35            y_ref: 0.0, // Default reference line at Y=0
36            flags: StemsFlags::NONE,
37            item_flags: ItemFlags::NONE,
38            layout: PlotDataLayout::DEFAULT,
39        }
40    }
41
42    /// Set the reference Y value for stems
43    /// Stems will be drawn from this Y value to the data points
44    pub fn with_y_ref(mut self, y_ref: f64) -> Self {
45        self.y_ref = y_ref;
46        self
47    }
48
49    /// Set stem flags for customization
50    pub fn with_flags(mut self, flags: StemsFlags) -> Self {
51        self.flags = flags;
52        self
53    }
54
55    /// Set common item flags for this plot item (applies to all plot types)
56    pub fn with_item_flags(mut self, flags: ItemFlags) -> Self {
57        self.item_flags = flags;
58        self
59    }
60
61    /// Set the data layout used to read X/Y samples.
62    ///
63    /// # Safety
64    ///
65    /// Every sample address computed from `layout` must refer to an initialized, properly aligned
66    /// `f64` within both coordinate allocations retained by this builder.
67    pub unsafe fn with_data_layout(mut self, layout: PlotDataLayout) -> Self {
68        self.layout = layout;
69        self
70    }
71
72    /// Set the sample-index offset used to read X/Y samples.
73    pub fn with_offset(mut self, offset: PlotDataOffset) -> Self {
74        self.layout = self.layout.with_offset(offset);
75        self
76    }
77
78    /// Set the byte stride used to read X/Y samples.
79    ///
80    /// # Safety
81    ///
82    /// Every strided sample read must remain initialized, aligned, and within both coordinate
83    /// allocations retained by this builder.
84    pub unsafe fn with_stride(mut self, stride: PlotDataStride) -> Self {
85        self.layout = self.layout.with_stride(stride);
86        self
87    }
88
89    /// Validate the plot data
90    pub fn validate(&self) -> Result<(), PlotError> {
91        validate_data_lengths(self.x_data, self.y_data)
92    }
93}
94
95impl<'a> Plot for StemPlot<'a> {
96    fn plot(&self, plot_ui: &crate::PlotUi<'_>) {
97        if self.validate().is_err() {
98            return;
99        }
100        let Ok(count) = i32::try_from(self.x_data.len()) else {
101            return;
102        };
103
104        plot_ui.with_bound_context(|| {
105            with_plot_str_or_empty(self.label, |label_ptr| unsafe {
106                let spec = plot_spec_with_style(
107                    self.style,
108                    self.flags.bits() | self.item_flags.bits(),
109                    self.layout,
110                );
111                sys::ImPlot_PlotStems_doublePtrdoublePtr(
112                    label_ptr,
113                    self.x_data.as_ptr(),
114                    self.y_data.as_ptr(),
115                    count,
116                    self.y_ref,
117                    spec,
118                );
119            })
120        })
121    }
122
123    fn label(&self) -> &str {
124        self.label
125    }
126}
127
128/// Simple stem plot for quick plotting without builder pattern
129pub struct SimpleStemPlot<'a> {
130    label: &'a str,
131    values: &'a [f64],
132    style: PlotItemStyle,
133    y_ref: f64,
134    flags: StemsFlags,
135    item_flags: ItemFlags,
136    x_scale: f64,
137    x_start: f64,
138}
139
140impl<'a> super::PlotItemStyled for SimpleStemPlot<'a> {
141    fn style_mut(&mut self) -> &mut PlotItemStyle {
142        &mut self.style
143    }
144}
145
146impl<'a> SimpleStemPlot<'a> {
147    /// Create a simple stem plot with Y values only (X will be indices)
148    pub fn new(label: &'a str, values: &'a [f64]) -> Self {
149        Self {
150            label,
151            values,
152            style: PlotItemStyle::default(),
153            y_ref: 0.0,
154            flags: StemsFlags::NONE,
155            item_flags: ItemFlags::NONE,
156            x_scale: 1.0,
157            x_start: 0.0,
158        }
159    }
160
161    /// Set the reference Y value for stems
162    pub fn with_y_ref(mut self, y_ref: f64) -> Self {
163        self.y_ref = y_ref;
164        self
165    }
166
167    /// Set stem flags for customization
168    pub fn with_flags(mut self, flags: StemsFlags) -> Self {
169        self.flags = flags;
170        self
171    }
172
173    /// Set common item flags for this plot item (applies to all plot types)
174    pub fn with_item_flags(mut self, flags: ItemFlags) -> Self {
175        self.item_flags = flags;
176        self
177    }
178
179    /// Set X scale factor
180    pub fn with_x_scale(mut self, scale: f64) -> Self {
181        self.x_scale = scale;
182        self
183    }
184
185    /// Set X start value
186    pub fn with_x_start(mut self, start: f64) -> Self {
187        self.x_start = start;
188        self
189    }
190}
191
192impl<'a> Plot for SimpleStemPlot<'a> {
193    fn plot(&self, plot_ui: &crate::PlotUi<'_>) {
194        if self.values.is_empty() {
195            return;
196        }
197        let Ok(count) = i32::try_from(self.values.len()) else {
198            return;
199        };
200
201        plot_ui.with_bound_context(|| {
202            with_plot_str_or_empty(self.label, |label_ptr| unsafe {
203                let spec = plot_spec_with_style(
204                    self.style,
205                    self.flags.bits() | self.item_flags.bits(),
206                    PlotDataLayout::DEFAULT,
207                );
208                sys::ImPlot_PlotStems_doublePtrInt(
209                    label_ptr,
210                    self.values.as_ptr(),
211                    count,
212                    self.y_ref,
213                    self.x_scale,
214                    self.x_start,
215                    spec,
216                );
217            })
218        })
219    }
220
221    fn label(&self) -> &str {
222        self.label
223    }
224}
225
226/// Convenience functions for quick stem plotting
227impl<'ui> crate::PlotUi<'ui> {
228    /// Plot a stem plot with X and Y data
229    pub fn stem_plot(&self, label: &str, x_data: &[f64], y_data: &[f64]) -> Result<(), PlotError> {
230        let plot = StemPlot::new(label, x_data, y_data);
231        plot.validate()?;
232        plot.plot(self);
233        Ok(())
234    }
235
236    /// Plot a stem plot with custom reference Y value
237    pub fn stem_plot_with_ref(
238        &self,
239        label: &str,
240        x_data: &[f64],
241        y_data: &[f64],
242        y_ref: f64,
243    ) -> Result<(), PlotError> {
244        let plot = StemPlot::new(label, x_data, y_data).with_y_ref(y_ref);
245        plot.validate()?;
246        plot.plot(self);
247        Ok(())
248    }
249
250    /// Plot a simple stem plot with Y values only (X will be indices)
251    pub fn simple_stem_plot(&self, label: &str, values: &[f64]) -> Result<(), PlotError> {
252        if values.is_empty() {
253            return Err(PlotError::EmptyData);
254        }
255        let plot = SimpleStemPlot::new(label, values);
256        plot.plot(self);
257        Ok(())
258    }
259
260    /// Plot a simple stem plot with custom reference Y value
261    pub fn simple_stem_plot_with_ref(
262        &self,
263        label: &str,
264        values: &[f64],
265        y_ref: f64,
266    ) -> Result<(), PlotError> {
267        if values.is_empty() {
268            return Err(PlotError::EmptyData);
269        }
270        let plot = SimpleStemPlot::new(label, values).with_y_ref(y_ref);
271        plot.plot(self);
272        Ok(())
273    }
274}
275
276#[cfg(test)]
277mod tests {
278    use super::*;
279
280    #[test]
281    fn test_stem_plot_creation() {
282        let x_data = [1.0, 2.0, 3.0, 4.0];
283        let y_data = [1.0, 4.0, 2.0, 3.0];
284
285        let plot = StemPlot::new("test", &x_data, &y_data);
286        assert_eq!(plot.label(), "test");
287        assert!(plot.validate().is_ok());
288    }
289
290    #[test]
291    fn test_stem_plot_validation() {
292        let x_data = [1.0, 2.0, 3.0];
293        let y_data = [1.0, 4.0]; // Different length
294
295        let plot = StemPlot::new("test", &x_data, &y_data);
296        assert!(plot.validate().is_err());
297    }
298
299    #[test]
300    fn test_simple_stem_plot() {
301        let values = [1.0, 2.0, 3.0, 4.0];
302        let plot = SimpleStemPlot::new("test", &values)
303            .with_flags(StemsFlags::HORIZONTAL)
304            .with_item_flags(ItemFlags::NO_LEGEND);
305        assert_eq!(plot.label(), "test");
306        assert_eq!(plot.flags.bits(), StemsFlags::HORIZONTAL.bits());
307        assert_eq!(plot.item_flags, ItemFlags::NO_LEGEND);
308    }
309
310    #[test]
311    fn test_stem_plot_with_ref() {
312        let x_data = [1.0, 2.0, 3.0, 4.0];
313        let y_data = [1.0, 4.0, 2.0, 3.0];
314
315        let plot = StemPlot::new("test", &x_data, &y_data).with_y_ref(1.0);
316        assert_eq!(plot.y_ref, 1.0);
317        assert!(plot.validate().is_ok());
318    }
319}