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 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
249pub 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}