1use crate::context::shared_wgpu_context;
4use crate::core::{
5 vertex_utils, BoundingBox, DrawCall, GpuVertexBuffer, Material, PipelineType, RenderData,
6 Vertex,
7};
8use crate::gpu::util::copy_readback_bytes;
9use bytemuck::cast_slice;
10use glam::{Vec3, Vec4};
11
12#[derive(Debug, Clone)]
13pub struct ContourPlot {
14 pub base_z: f32,
15 pub force_3d: bool,
16 pub label: Option<String>,
17 pub visible: bool,
18 pub line_width: f32,
19 vertices: Option<Vec<Vertex>>,
20 gpu_vertices: Option<GpuVertexBuffer>,
21 vertex_count: usize,
22 bounds: Option<BoundingBox>,
23}
24
25impl ContourPlot {
26 pub async fn export_scene_vertices(&self) -> Result<Vec<Vertex>, String> {
27 if let Some(vertices) = &self.vertices {
28 return Ok(vertices.clone());
29 }
30
31 if let Some(gpu_vertices) = &self.gpu_vertices {
32 let context = shared_wgpu_context().ok_or_else(|| {
33 "contour plot has GPU vertices but no shared WGPU context is installed".to_string()
34 })?;
35 let vertex_count = self.vertex_count.min(gpu_vertices.vertex_count);
36 let byte_len = vertex_count * std::mem::size_of::<Vertex>();
37 let bytes = copy_readback_bytes(
38 &context.device,
39 &context.queue,
40 &gpu_vertices.buffer,
41 byte_len,
42 )
43 .await?;
44 let vertices: &[Vertex] = cast_slice(&bytes);
45 return Ok(vertices.to_vec());
46 }
47
48 Ok(Vec::new())
49 }
50
51 pub fn from_vertices(vertices: Vec<Vertex>, base_z: f32, bounds: BoundingBox) -> Self {
53 Self {
54 base_z,
55 force_3d: false,
56 label: None,
57 visible: true,
58 line_width: 1.0,
59 vertex_count: vertices.len(),
60 vertices: Some(vertices),
61 gpu_vertices: None,
62 bounds: Some(bounds),
63 }
64 }
65
66 pub fn from_gpu_buffer(
68 buffer: GpuVertexBuffer,
69 vertex_count: usize,
70 base_z: f32,
71 bounds: BoundingBox,
72 ) -> Self {
73 Self {
74 base_z,
75 force_3d: false,
76 label: None,
77 visible: true,
78 line_width: 1.0,
79 vertex_count,
80 vertices: None,
81 gpu_vertices: Some(buffer),
82 bounds: Some(bounds),
83 }
84 }
85
86 pub fn with_label<S: Into<String>>(mut self, label: S) -> Self {
87 self.label = Some(label.into());
88 self
89 }
90
91 pub fn with_force_3d(mut self, force_3d: bool) -> Self {
92 self.force_3d = force_3d;
93 self
94 }
95
96 pub fn is_3d(&self) -> bool {
97 self.force_3d || (self.bounds().max.z - self.bounds().min.z).abs() > f32::EPSILON
98 }
99
100 pub fn set_visible(&mut self, visible: bool) {
101 self.visible = visible;
102 }
103
104 pub fn with_line_width(mut self, line_width: f32) -> Self {
105 self.line_width = line_width.max(0.5);
106 self
107 }
108
109 pub fn vertices(&mut self) -> &Vec<Vertex> {
110 if self.vertices.is_none() {
111 self.vertices = Some(Vec::new());
112 }
113 self.vertices.as_ref().unwrap()
114 }
115
116 pub fn bounds(&self) -> BoundingBox {
117 self.bounds.unwrap_or_default()
118 }
119
120 pub fn cpu_vertices(&self) -> Option<&[Vertex]> {
121 self.vertices.as_deref()
122 }
123
124 pub fn render_data_with_viewport(&mut self, viewport_px: Option<(u32, u32)>) -> RenderData {
125 if self.gpu_vertices.is_some() {
126 return self.render_data();
127 }
128
129 let bounds = self.bounds();
130 let (vertices, vertex_count, pipeline_type, render_bounds) = if self.line_width > 1.0 {
131 let Some(viewport_px) = viewport_px else {
132 return self.render_data();
133 };
134 let data_per_px = crate::core::data_units_per_px(&bounds, viewport_px);
135 let width_data = self.line_width.max(0.1) * data_per_px;
136 let verts = self.vertices().clone();
137 let mut thick = Vec::new();
138 for segment in verts.chunks_exact(2) {
139 let color = Vec4::from_array(segment[0].color);
140 if self.is_3d() {
141 thick.extend(create_thick_segment_3d(
142 Vec3::from_array(segment[0].position),
143 Vec3::from_array(segment[1].position),
144 color,
145 width_data * 0.5,
146 ));
147 } else {
148 let x = [segment[0].position[0] as f64, segment[1].position[0] as f64];
149 let y = [segment[0].position[1] as f64, segment[1].position[1] as f64];
150 thick.extend(vertex_utils::create_thick_polyline(
151 &x, &y, color, width_data,
152 ));
153 }
154 }
155 let count = thick.len();
156 let render_bounds = expanded_bounds_for_vertices(bounds, &thick);
157 (thick, count, PipelineType::Triangles, render_bounds)
158 } else {
159 let verts = self.vertices().clone();
160 let count = verts.len();
161 (verts, count, PipelineType::Lines, bounds)
162 };
163
164 let material = Material {
165 albedo: Vec4::ONE,
166 roughness: self.line_width.max(0.0),
167 ..Default::default()
168 };
169
170 let draw_call = DrawCall {
171 vertex_offset: 0,
172 vertex_count,
173 index_offset: None,
174 index_count: None,
175 instance_count: 1,
176 };
177
178 RenderData {
179 pipeline_type,
180 vertices,
181 indices: None,
182 gpu_vertices: None,
183 bounds: Some(render_bounds),
184 material,
185 draw_calls: vec![draw_call],
186 image: None,
187 }
188 }
189
190 pub fn render_data(&mut self) -> RenderData {
191 let using_gpu = self.gpu_vertices.is_some();
192 let bounds = self.bounds();
193 let (vertices, vertex_count, gpu_vertices, pipeline_type, render_bounds) = if using_gpu {
194 (
195 Vec::new(),
196 self.vertex_count,
197 self.gpu_vertices.clone(),
198 PipelineType::Lines,
199 bounds,
200 )
201 } else {
202 let verts = self.vertices().clone();
203 if self.line_width > 1.0 {
204 let mut thick = Vec::new();
205 for segment in verts.chunks_exact(2) {
206 let color = Vec4::from_array(segment[0].color);
207 if self.is_3d() {
208 thick.extend(create_thick_segment_3d(
209 Vec3::from_array(segment[0].position),
210 Vec3::from_array(segment[1].position),
211 color,
212 self.line_width.max(0.5) * 0.01,
213 ));
214 } else {
215 let x = [segment[0].position[0] as f64, segment[1].position[0] as f64];
216 let y = [segment[0].position[1] as f64, segment[1].position[1] as f64];
217 thick.extend(vertex_utils::create_thick_polyline(
218 &x,
219 &y,
220 color,
221 self.line_width,
222 ));
223 }
224 }
225 let count = thick.len();
226 let render_bounds = expanded_bounds_for_vertices(bounds, &thick);
227 (thick, count, None, PipelineType::Triangles, render_bounds)
228 } else {
229 let count = verts.len();
230 (verts, count, None, PipelineType::Lines, bounds)
231 }
232 };
233
234 let material = Material {
235 albedo: Vec4::ONE,
236 roughness: self.line_width.max(0.0),
237 ..Default::default()
238 };
239
240 let draw_call = DrawCall {
241 vertex_offset: 0,
242 vertex_count,
243 index_offset: None,
244 index_count: None,
245 instance_count: 1,
246 };
247
248 RenderData {
249 pipeline_type,
250 vertices,
251 indices: None,
252 gpu_vertices,
253 bounds: Some(render_bounds),
254 material,
255 draw_calls: vec![draw_call],
256 image: None,
257 }
258 }
259
260 pub fn estimated_memory_usage(&self) -> usize {
261 self.vertices
262 .as_ref()
263 .map(|v| v.len() * std::mem::size_of::<Vertex>())
264 .unwrap_or(0)
265 }
266}
267
268pub fn contour_bounds(x_min: f32, x_max: f32, y_min: f32, y_max: f32, base_z: f32) -> BoundingBox {
269 BoundingBox::new(
270 Vec3::new(x_min, y_min, base_z),
271 Vec3::new(x_max, y_max, base_z),
272 )
273}
274
275pub fn contour_bounds_3d(
276 x_min: f32,
277 x_max: f32,
278 y_min: f32,
279 y_max: f32,
280 z_min: f32,
281 z_max: f32,
282) -> BoundingBox {
283 BoundingBox::new(
284 Vec3::new(x_min, y_min, z_min),
285 Vec3::new(x_max, y_max, z_max),
286 )
287}
288
289fn expanded_bounds_for_vertices(mut bounds: BoundingBox, vertices: &[Vertex]) -> BoundingBox {
290 for vertex in vertices {
291 bounds.expand(Vec3::from_array(vertex.position));
292 }
293 bounds
294}
295
296fn create_thick_segment_3d(start: Vec3, end: Vec3, color: Vec4, half_width: f32) -> Vec<Vertex> {
297 let dir = (end - start).normalize_or_zero();
298 let mut normal = dir.cross(Vec3::Z);
299 if normal.length_squared() < 1e-6 {
300 normal = dir.cross(Vec3::X);
301 }
302 let normal = normal.normalize_or_zero() * half_width.max(0.0001);
303 let v0 = start + normal;
304 let v1 = end + normal;
305 let v2 = end - normal;
306 let v3 = start - normal;
307 vec![
308 Vertex::new(v0, color),
309 Vertex::new(v1, color),
310 Vertex::new(v2, color),
311 Vertex::new(v0, color),
312 Vertex::new(v2, color),
313 Vertex::new(v3, color),
314 ]
315}
316
317#[cfg(test)]
318mod tests {
319 use super::*;
320
321 fn test_vertex(x: f32, y: f32, z: f32) -> Vertex {
322 Vertex::new(Vec3::new(x, y, z), Vec4::ONE)
323 }
324
325 #[test]
326 fn viewport_thick_contour_bounds_include_extruded_geometry() {
327 let bounds = BoundingBox::new(Vec3::new(0.0, 0.0, 0.0), Vec3::new(1.0, 0.0, 0.0));
328 let mut contour = ContourPlot::from_vertices(
329 vec![test_vertex(0.0, 0.0, 0.0), test_vertex(1.0, 0.0, 0.0)],
330 0.0,
331 bounds,
332 )
333 .with_line_width(2.0);
334
335 let render_data = contour.render_data_with_viewport(Some((100, 100)));
336 let render_bounds = render_data.bounds.expect("bounds");
337
338 assert!(render_bounds.min.y < bounds.min.y);
339 assert!(render_bounds.max.y > bounds.max.y);
340 }
341
342 #[test]
343 fn non_viewport_thick_3d_contour_bounds_include_extruded_geometry() {
344 let bounds = BoundingBox::new(Vec3::new(0.0, 0.0, 1.0), Vec3::new(0.0, 1.0, 1.0));
345 let mut contour = ContourPlot::from_vertices(
346 vec![test_vertex(0.0, 0.0, 1.0), test_vertex(0.0, 1.0, 1.0)],
347 0.0,
348 bounds,
349 )
350 .with_force_3d(true)
351 .with_line_width(2.0);
352
353 let render_data = contour.render_data();
354 let render_bounds = render_data.bounds.expect("bounds");
355
356 assert!(render_bounds.min.x < bounds.min.x);
357 assert!(render_bounds.max.x > bounds.max.x);
358 }
359
360 #[test]
361 fn viewport_thick_3d_contour_uses_half_width_data() {
362 let bounds = BoundingBox::new(Vec3::new(0.0, 0.0, 1.0), Vec3::new(1.0, 1.0, 1.0));
363 let mut contour = ContourPlot::from_vertices(
364 vec![test_vertex(0.0, 0.0, 1.0), test_vertex(1.0, 0.0, 1.0)],
365 0.0,
366 bounds,
367 )
368 .with_force_3d(true)
369 .with_line_width(4.0);
370
371 let render_data = contour.render_data_with_viewport(Some((100, 100)));
372 let render_bounds = render_data.bounds.expect("bounds");
373
374 assert!((render_bounds.min.y - -0.02).abs() < 1e-6);
375 }
376}