Skip to main content

dear_implot3d/plots/
surface.rs

1use std::borrow::Cow;
2
3use super::{Plot3D, Plot3DError};
4use crate::{Plot3DDataLayout, Plot3DUi, Surface3DFlags, debug_before_plot, plot3d_spec_from, sys};
5
6const SURFACE_LABEL_ERROR: &str = "surface label contains NUL";
7
8#[derive(Debug, Clone, Copy)]
9pub(crate) struct SurfaceLabel<'a>(&'a str);
10
11impl<'a> SurfaceLabel<'a> {
12    pub(crate) fn checked(label: &'a str) -> Result<Self, Plot3DError> {
13        if label.contains('\0') {
14            Err(Plot3DError::StringConversion(SURFACE_LABEL_ERROR))
15        } else {
16            Ok(Self(label))
17        }
18    }
19
20    const fn as_str(self) -> &'a str {
21        self.0
22    }
23}
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub(crate) struct SurfaceGridShape {
27    x_count: i32,
28    y_count: i32,
29    point_count: usize,
30}
31
32impl SurfaceGridShape {
33    pub(crate) fn checked(x_count: usize, y_count: usize) -> Result<Self, Plot3DError> {
34        if x_count == 0 || y_count == 0 {
35            return Err(Plot3DError::EmptyData);
36        }
37
38        let point_count = x_count
39            .checked_mul(y_count)
40            .ok_or(Plot3DError::GridSizeOverflow { x_count, y_count })?;
41        if point_count > i32::MAX as usize {
42            return Err(Plot3DError::GridPointCountOutOfRange {
43                x_count,
44                y_count,
45                point_count,
46            });
47        }
48
49        // Non-empty dimensions cannot exceed the already-checked point count.
50        let range_error = || Plot3DError::GridPointCountOutOfRange {
51            x_count,
52            y_count,
53            point_count,
54        };
55        let x_count_i32 = i32::try_from(x_count).map_err(|_| range_error())?;
56        let y_count_i32 = i32::try_from(y_count).map_err(|_| range_error())?;
57
58        Ok(Self {
59            x_count: x_count_i32,
60            y_count: y_count_i32,
61            point_count,
62        })
63    }
64
65    fn validate_z_len(self, z_len: usize) -> Result<Self, Plot3DError> {
66        if z_len == self.point_count {
67            Ok(self)
68        } else {
69            Err(Plot3DError::GridSizeMismatch {
70                x_count: self.x_count as usize,
71                y_count: self.y_count as usize,
72                expected: self.point_count,
73                z_len,
74            })
75        }
76    }
77
78    pub(crate) const fn counts_i32(self) -> (i32, i32) {
79        (self.x_count, self.y_count)
80    }
81
82    const fn point_count(self) -> usize {
83        self.point_count
84    }
85}
86
87#[derive(Debug)]
88pub(crate) struct SurfaceGrid<'a> {
89    shape: SurfaceGridShape,
90    xs: Cow<'a, [f32]>,
91    ys: Cow<'a, [f32]>,
92    zs: &'a [f32],
93}
94
95impl<'a> SurfaceGrid<'a> {
96    pub(crate) fn from_axes(
97        xs: &'a [f32],
98        ys: &'a [f32],
99        zs: &'a [f32],
100    ) -> Result<Self, Plot3DError> {
101        let shape = SurfaceGridShape::checked(xs.len(), ys.len())?.validate_z_len(zs.len())?;
102        let mut xs_flat = Vec::with_capacity(shape.point_count());
103        let mut ys_flat = Vec::with_capacity(shape.point_count());
104
105        for &y in ys {
106            for &x in xs {
107                xs_flat.push(x);
108                ys_flat.push(y);
109            }
110        }
111
112        Ok(Self {
113            shape,
114            xs: Cow::Owned(xs_flat),
115            ys: Cow::Owned(ys_flat),
116            zs,
117        })
118    }
119
120    pub(crate) fn from_flattened(
121        xs: &'a [f32],
122        ys: &'a [f32],
123        zs: &'a [f32],
124        x_count: usize,
125        y_count: usize,
126    ) -> Result<Self, Plot3DError> {
127        let shape = SurfaceGridShape::checked(x_count, y_count)?.validate_z_len(zs.len())?;
128        validate_coordinate_len(xs.len(), shape.point_count(), "surface x coordinates")?;
129        validate_coordinate_len(ys.len(), shape.point_count(), "surface y coordinates")?;
130
131        Ok(Self {
132            shape,
133            xs: Cow::Borrowed(xs),
134            ys: Cow::Borrowed(ys),
135            zs,
136        })
137    }
138
139    pub(crate) fn shape(&self) -> SurfaceGridShape {
140        self.shape
141    }
142
143    pub(crate) fn xs(&self) -> &[f32] {
144        self.xs.as_ref()
145    }
146
147    pub(crate) fn ys(&self) -> &[f32] {
148        self.ys.as_ref()
149    }
150
151    pub(crate) fn zs(&self) -> &[f32] {
152        self.zs
153    }
154}
155
156fn validate_coordinate_len(
157    actual: usize,
158    expected: usize,
159    what: &'static str,
160) -> Result<(), Plot3DError> {
161    if actual == expected {
162        Ok(())
163    } else {
164        Err(Plot3DError::DataLengthMismatch {
165            a: actual,
166            b: expected,
167            what,
168        })
169    }
170}
171
172pub(crate) fn submit_surface_grid(
173    ui: &Plot3DUi<'_>,
174    label: SurfaceLabel<'_>,
175    grid: &SurfaceGrid<'_>,
176    scale_min: f64,
177    scale_max: f64,
178    make_spec: impl FnOnce() -> sys::ImPlot3DSpec_c,
179) {
180    let label = label.as_str();
181    let (x_count, y_count) = grid.shape().counts_i32();
182
183    ui.with_bound_context(|| {
184        debug_before_plot();
185        let spec = make_spec();
186        dear_imgui_rs::with_scratch_txt(label, |label_ptr| unsafe {
187            #[cfg(all(test, not(target_arch = "wasm32")))]
188            sys::surface_test_probe::dear_implot3d_surface_probe_plot(
189                label_ptr,
190                grid.xs().as_ptr(),
191                grid.ys().as_ptr(),
192                grid.zs().as_ptr(),
193                x_count,
194                y_count,
195                scale_min,
196                scale_max,
197                spec,
198            );
199            #[cfg(not(all(test, not(target_arch = "wasm32"))))]
200            sys::ImPlot3D_PlotSurface_FloatPtr(
201                label_ptr,
202                grid.xs().as_ptr(),
203                grid.ys().as_ptr(),
204                grid.zs().as_ptr(),
205                x_count,
206                y_count,
207                scale_min,
208                scale_max,
209                spec,
210            );
211        });
212    });
213}
214
215pub(crate) unsafe fn submit_surface_raw(
216    ui: &Plot3DUi<'_>,
217    label: SurfaceLabel<'_>,
218    xs: &[f32],
219    ys: &[f32],
220    zs: &[f32],
221    shape: SurfaceGridShape,
222    scale_min: f64,
223    scale_max: f64,
224    flags: Surface3DFlags,
225    layout: Plot3DDataLayout,
226) {
227    let label = label.as_str();
228    let (x_count, y_count) = shape.counts_i32();
229
230    ui.with_bound_context(|| {
231        debug_before_plot();
232        let spec = plot3d_spec_from(flags.bits(), layout);
233        dear_imgui_rs::with_scratch_txt(label, |label_ptr| unsafe {
234            sys::ImPlot3D_PlotSurface_FloatPtr(
235                label_ptr,
236                xs.as_ptr(),
237                ys.as_ptr(),
238                zs.as_ptr(),
239                x_count,
240                y_count,
241                scale_min,
242                scale_max,
243                spec,
244            );
245        });
246    });
247}
248
249/// A surface whose X and Y inputs describe grid axes.
250///
251/// Custom offsets are intentionally unavailable on this safe type. Use the unsafe
252/// [`Plot3DUi::surface_f32_raw`] escape hatch for arbitrary native layouts.
253///
254/// ```compile_fail
255/// use dear_implot3d::{Plot3DDataOffset, Surface3D};
256///
257/// let xs = [0.0, 1.0];
258/// let ys = [0.0, 1.0];
259/// let zs = [0.0, 1.0, 1.0, 2.0];
260/// let _ = Surface3D::new("surface", &xs, &ys, &zs)
261///     .offset(Plot3DDataOffset::samples(1));
262/// ```
263pub struct Surface3D<'a> {
264    pub label: &'a str,
265    pub xs: &'a [f32],
266    pub ys: &'a [f32],
267    pub zs: &'a [f32],
268    pub scale_min: f64,
269    pub scale_max: f64,
270    pub flags: Surface3DFlags,
271}
272
273impl<'a> Surface3D<'a> {
274    pub fn new(label: &'a str, xs: &'a [f32], ys: &'a [f32], zs: &'a [f32]) -> Self {
275        Self {
276            label,
277            xs,
278            ys,
279            zs,
280            scale_min: f64::NAN,
281            scale_max: f64::NAN,
282            flags: Surface3DFlags::NONE,
283        }
284    }
285
286    pub fn scale(mut self, min: f64, max: f64) -> Self {
287        self.scale_min = min;
288        self.scale_max = max;
289        self
290    }
291
292    pub fn flags(mut self, flags: Surface3DFlags) -> Self {
293        self.flags = flags;
294        self
295    }
296}
297
298impl<'a> Plot3D for Surface3D<'a> {
299    fn label(&self) -> &str {
300        self.label
301    }
302
303    fn try_plot(&self, ui: &Plot3DUi<'_>) -> Result<(), Plot3DError> {
304        let label = SurfaceLabel::checked(self.label)?;
305        let grid = SurfaceGrid::from_axes(self.xs, self.ys, self.zs)?;
306        submit_surface_grid(ui, label, &grid, self.scale_min, self.scale_max, || {
307            plot3d_spec_from(self.flags.bits(), Plot3DDataLayout::DEFAULT)
308        });
309        Ok(())
310    }
311}
312
313#[cfg(test)]
314mod tests {
315    use std::borrow::Cow;
316
317    #[cfg(not(target_arch = "wasm32"))]
318    use std::sync::{Mutex, OnceLock};
319
320    use super::{Plot3D, Surface3D, SurfaceGrid, SurfaceGridShape, SurfaceLabel};
321    use crate::{plots::Plot3DError, sys};
322
323    #[test]
324    fn axis_grid_flattens_in_row_major_order() {
325        let xs = [10.0, 20.0];
326        let ys = [1.0, 2.0];
327        let zs = [11.0, 21.0, 12.0, 22.0];
328
329        let grid = SurfaceGrid::from_axes(&xs, &ys, &zs).unwrap();
330
331        assert_eq!(grid.xs(), &[10.0, 20.0, 10.0, 20.0]);
332        assert_eq!(grid.ys(), &[1.0, 1.0, 2.0, 2.0]);
333        assert_eq!(grid.zs(), &zs);
334        assert_eq!(grid.shape().counts_i32(), (2, 2));
335    }
336
337    #[cfg(not(target_arch = "wasm32"))]
338    #[test]
339    fn cpp_probe_receives_equal_length_row_major_surface_arrays() {
340        static GUARD: OnceLock<Mutex<()>> = OnceLock::new();
341        let _guard = GUARD
342            .get_or_init(|| Mutex::new(()))
343            .lock()
344            .unwrap_or_else(std::sync::PoisonError::into_inner);
345
346        let mut imgui = dear_imgui_rs::Context::create();
347        let io = imgui.io_mut();
348        io.set_display_size([800.0, 600.0]);
349        io.set_delta_time(1.0 / 60.0);
350        imgui
351            .font_atlas()
352            .try_claim_legacy_renderer()
353            .expect("headless test requires the legacy font-atlas capability")
354            .build();
355        let plot_context = crate::Plot3DContext::create(&imgui);
356        let frame = imgui.begin_frame();
357        let plot_ui = plot_context.get_plot_ui(frame.ui());
358        let _plot = plot_ui
359            .begin_plot("surface C++ capture")
360            .build()
361            .expect("surface capture plot should begin");
362
363        let xs = [10.0, 20.0];
364        let ys = [1.0, 2.0];
365        let zs = [11.0, 21.0, 12.0, 22.0];
366        unsafe { sys::surface_test_probe::dear_implot3d_surface_probe_reset() };
367        Surface3D::new("captured surface", &xs, &ys, &zs)
368            .try_plot(&plot_ui)
369            .unwrap();
370
371        let mut captured_xs = [0.0; 4];
372        let mut captured_ys = [0.0; 4];
373        let mut captured_zs = [0.0; 4];
374        let mut x_count = 0;
375        let mut y_count = 0;
376        let point_count = unsafe {
377            sys::surface_test_probe::dear_implot3d_surface_probe_read(
378                captured_xs.as_mut_ptr(),
379                captured_ys.as_mut_ptr(),
380                captured_zs.as_mut_ptr(),
381                captured_xs.len() as i32,
382                &mut x_count,
383                &mut y_count,
384            )
385        };
386
387        assert_eq!((x_count, y_count, point_count), (2, 2, 4));
388        assert_eq!(captured_xs, [10.0, 20.0, 10.0, 20.0]);
389        assert_eq!(captured_ys, [1.0, 1.0, 2.0, 2.0]);
390        assert_eq!(captured_zs, zs);
391    }
392
393    #[test]
394    fn one_by_n_axis_grid_is_valid() {
395        let xs = [3.0];
396        let ys = [1.0, 2.0, 3.0];
397        let zs = [4.0, 5.0, 6.0];
398
399        let grid = SurfaceGrid::from_axes(&xs, &ys, &zs).unwrap();
400
401        assert_eq!(grid.xs(), &[3.0, 3.0, 3.0]);
402        assert_eq!(grid.ys(), &ys);
403        assert_eq!(grid.zs(), &zs);
404    }
405
406    #[test]
407    fn empty_axes_are_rejected() {
408        assert_eq!(SurfaceGridShape::checked(0, 1), Err(Plot3DError::EmptyData));
409        assert_eq!(SurfaceGridShape::checked(1, 0), Err(Plot3DError::EmptyData));
410    }
411
412    #[test]
413    fn axis_grid_rejects_z_length_mismatch() {
414        assert_eq!(
415            SurfaceGrid::from_axes(&[0.0, 1.0], &[0.0, 1.0], &[0.0, 1.0, 2.0]).unwrap_err(),
416            Plot3DError::GridSizeMismatch {
417                x_count: 2,
418                y_count: 2,
419                expected: 4,
420                z_len: 3,
421            }
422        );
423    }
424
425    #[test]
426    fn checked_shape_rejects_usize_multiplication_overflow() {
427        let x_count = usize::MAX / 2 + 1;
428        let y_count = 2;
429
430        assert_eq!(
431            SurfaceGridShape::checked(x_count, y_count),
432            Err(Plot3DError::GridSizeOverflow { x_count, y_count })
433        );
434    }
435
436    #[test]
437    fn checked_shape_rejects_cpp_int_overflow_without_allocating() {
438        let x_count = i32::MAX as usize;
439        let y_count = 2;
440        let point_count = x_count * y_count;
441
442        assert_eq!(
443            SurfaceGridShape::checked(x_count, y_count),
444            Err(Plot3DError::GridPointCountOutOfRange {
445                x_count,
446                y_count,
447                point_count,
448            })
449        );
450    }
451
452    #[test]
453    fn flattened_grid_requires_all_coordinate_arrays_to_match() {
454        let error = SurfaceGrid::from_flattened(&[0.0; 4], &[0.0; 3], &[0.0; 4], 2, 2).unwrap_err();
455
456        assert_eq!(
457            error,
458            Plot3DError::DataLengthMismatch {
459                a: 3,
460                b: 4,
461                what: "surface y coordinates",
462            }
463        );
464    }
465
466    #[test]
467    fn flattened_grid_borrows_contiguous_coordinates() {
468        let xs = [0.0; 4];
469        let ys = [0.0; 4];
470        let zs = [0.0; 4];
471        let grid = SurfaceGrid::from_flattened(&xs, &ys, &zs, 2, 2).unwrap();
472
473        assert!(matches!(grid.xs, Cow::Borrowed(_)));
474        assert!(matches!(grid.ys, Cow::Borrowed(_)));
475    }
476
477    #[test]
478    fn axis_and_flattened_inputs_share_z_shape_errors() {
479        let axis_error = SurfaceGrid::from_axes(&[0.0, 1.0], &[0.0, 1.0], &[0.0; 3]).unwrap_err();
480        let flat_error =
481            SurfaceGrid::from_flattened(&[0.0; 4], &[0.0; 4], &[0.0; 3], 2, 2).unwrap_err();
482
483        assert_eq!(axis_error, flat_error);
484    }
485
486    #[test]
487    fn surface_labels_reject_embedded_nul_consistently() {
488        assert_eq!(
489            SurfaceLabel::checked("invalid\0label").map(|_| ()),
490            Err(Plot3DError::StringConversion("surface label contains NUL"))
491        );
492        assert_eq!(SurfaceLabel::checked("valid label").map(|_| ()), Ok(()));
493    }
494
495    #[test]
496    fn grid_mismatch_display_uses_validated_expected_count() {
497        let error = Plot3DError::GridSizeMismatch {
498            x_count: usize::MAX,
499            y_count: usize::MAX,
500            expected: 4,
501            z_len: 3,
502        };
503
504        assert_eq!(
505            error.to_string(),
506            format!(
507                "grid mismatch: x={} y={} => expected z_len=4, got 3",
508                usize::MAX,
509                usize::MAX
510            )
511        );
512    }
513}