polyscope_render/engine/
pick.rs

1use std::num::NonZeroU64;
2
3use super::RenderEngine;
4
5/// A contiguous range of global pick indices assigned to a structure.
6#[derive(Debug, Clone)]
7pub struct PickRange {
8    /// The starting global index for this structure's elements.
9    pub global_start: u32,
10    /// Number of elements in this range.
11    pub count: u32,
12    /// Structure type name (e.g., `PointCloud`, `SurfaceMesh`).
13    pub type_name: String,
14    /// Structure name.
15    pub name: String,
16}
17
18impl RenderEngine {
19    // ========== Pick System - Range-Based ID Management ==========
20
21    /// Assigns a contiguous range of global pick indices to a structure.
22    ///
23    /// Returns the `global_start` index. The structure owns indices
24    /// `[global_start, global_start + num_elements)`.
25    ///
26    /// Index 0 is reserved as background (no hit), so all ranges start from >= 1.
27    pub fn assign_pick_range(&mut self, type_name: &str, name: &str, num_elements: u32) -> u32 {
28        let key = (type_name.to_string(), name.to_string());
29
30        // If already assigned, return existing start
31        if let Some(range) = self.pick_ranges.get(&key) {
32            return range.global_start;
33        }
34
35        let global_start = self.next_global_index;
36        self.next_global_index += num_elements;
37
38        let range = PickRange {
39            global_start,
40            count: num_elements,
41            type_name: type_name.to_string(),
42            name: name.to_string(),
43        };
44
45        self.pick_ranges.insert(key, range);
46
47        global_start
48    }
49
50    /// Removes a structure's pick range.
51    ///
52    /// The range is freed but the global index counter is not decremented
53    /// (monotonic allocation — no fragmentation complexity).
54    pub fn remove_pick_range(&mut self, type_name: &str, name: &str) {
55        let key = (type_name.to_string(), name.to_string());
56        self.pick_ranges.remove(&key);
57    }
58
59    /// Looks up which structure owns a global pick index.
60    ///
61    /// Returns `(type_name, name, local_element_index)` or None if no structure
62    /// owns this index (background or freed range).
63    pub fn lookup_global_index(&self, global_index: u32) -> Option<(&str, &str, u32)> {
64        for range in self.pick_ranges.values() {
65            if global_index >= range.global_start && global_index < range.global_start + range.count
66            {
67                let local = global_index - range.global_start;
68                return Some((&range.type_name, &range.name, local));
69            }
70        }
71        None
72    }
73
74    /// Gets the `global_start` for a structure's pick range, if assigned.
75    pub fn get_pick_range_start(&self, type_name: &str, name: &str) -> Option<u32> {
76        let key = (type_name.to_string(), name.to_string());
77        self.pick_ranges.get(&key).map(|r| r.global_start)
78    }
79
80    // ========== Pick System - GPU Resources ==========
81
82    /// Creates or recreates pick buffer textures to match viewport size.
83    pub fn init_pick_buffers(&mut self, width: u32, height: u32) {
84        // Skip if size unchanged
85        if self.pick_buffer_size == (width, height) && self.pick_texture.is_some() {
86            return;
87        }
88
89        let device = &self.device;
90
91        // Create pick color texture (Rgba8Unorm for exact values)
92        let pick_texture = device.create_texture(&wgpu::TextureDescriptor {
93            label: Some("Pick Texture"),
94            size: wgpu::Extent3d {
95                width,
96                height,
97                depth_or_array_layers: 1,
98            },
99            mip_level_count: 1,
100            sample_count: 1,
101            dimension: wgpu::TextureDimension::D2,
102            format: wgpu::TextureFormat::Rgba8Unorm,
103            usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
104            view_formats: &[],
105        });
106        let pick_texture_view = pick_texture.create_view(&wgpu::TextureViewDescriptor::default());
107
108        // Create pick depth texture
109        let pick_depth_texture = device.create_texture(&wgpu::TextureDescriptor {
110            label: Some("Pick Depth Texture"),
111            size: wgpu::Extent3d {
112                width,
113                height,
114                depth_or_array_layers: 1,
115            },
116            mip_level_count: 1,
117            sample_count: 1,
118            dimension: wgpu::TextureDimension::D2,
119            format: wgpu::TextureFormat::Depth24Plus,
120            usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
121            view_formats: &[],
122        });
123        let pick_depth_view =
124            pick_depth_texture.create_view(&wgpu::TextureViewDescriptor::default());
125
126        // Create staging buffer for single pixel readback (4 bytes RGBA)
127        // Buffer size must be aligned to COPY_BYTES_PER_ROW_ALIGNMENT (256)
128        let pick_staging_buffer = device.create_buffer(&wgpu::BufferDescriptor {
129            label: Some("Pick Staging Buffer"),
130            size: 256, // Minimum aligned size, we only read 4 bytes
131            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
132            mapped_at_creation: false,
133        });
134
135        self.pick_texture = Some(pick_texture);
136        self.pick_texture_view = Some(pick_texture_view);
137        self.pick_depth_texture = Some(pick_depth_texture);
138        self.pick_depth_view = Some(pick_depth_view);
139        self.pick_staging_buffer = Some(pick_staging_buffer);
140        self.pick_buffer_size = (width, height);
141    }
142
143    /// Initializes the pick pipeline for point clouds.
144    pub(crate) fn init_pick_pipeline(&mut self) {
145        let shader_source = include_str!("../shaders/pick.wgsl");
146        let shader = self
147            .device
148            .create_shader_module(wgpu::ShaderModuleDescriptor {
149                label: Some("Pick Shader"),
150                source: wgpu::ShaderSource::Wgsl(shader_source.into()),
151            });
152
153        // Pick bind group layout: camera, pick uniforms, positions
154        let bind_group_layout =
155            self.device
156                .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
157                    label: Some("Pick Bind Group Layout"),
158                    entries: &[
159                        // Camera uniforms
160                        wgpu::BindGroupLayoutEntry {
161                            binding: 0,
162                            visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
163                            ty: wgpu::BindingType::Buffer {
164                                ty: wgpu::BufferBindingType::Uniform,
165                                has_dynamic_offset: false,
166                                min_binding_size: NonZeroU64::new(272),
167                            },
168                            count: None,
169                        },
170                        // Pick uniforms
171                        wgpu::BindGroupLayoutEntry {
172                            binding: 1,
173                            visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
174                            ty: wgpu::BindingType::Buffer {
175                                ty: wgpu::BufferBindingType::Uniform,
176                                has_dynamic_offset: false,
177                                min_binding_size: NonZeroU64::new(16),
178                            },
179                            count: None,
180                        },
181                        // Position storage buffer
182                        wgpu::BindGroupLayoutEntry {
183                            binding: 2,
184                            visibility: wgpu::ShaderStages::VERTEX,
185                            ty: wgpu::BindingType::Buffer {
186                                ty: wgpu::BufferBindingType::Storage { read_only: true },
187                                has_dynamic_offset: false,
188                                min_binding_size: None,
189                            },
190                            count: None,
191                        },
192                    ],
193                });
194
195        let pipeline_layout = self
196            .device
197            .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
198                label: Some("Pick Pipeline Layout"),
199                bind_group_layouts: &[&bind_group_layout],
200                push_constant_ranges: &[],
201            });
202
203        let pipeline = self
204            .device
205            .create_render_pipeline(&wgpu::RenderPipelineDescriptor {
206                label: Some("PointCloud Pick Pipeline"),
207                layout: Some(&pipeline_layout),
208                vertex: wgpu::VertexState {
209                    module: &shader,
210                    entry_point: Some("vs_main"),
211                    buffers: &[],
212                    compilation_options: wgpu::PipelineCompilationOptions::default(),
213                },
214                fragment: Some(wgpu::FragmentState {
215                    module: &shader,
216                    entry_point: Some("fs_main"),
217                    targets: &[Some(wgpu::ColorTargetState {
218                        format: wgpu::TextureFormat::Rgba8Unorm,
219                        blend: None, // No blending for pick buffer
220                        write_mask: wgpu::ColorWrites::ALL,
221                    })],
222                    compilation_options: wgpu::PipelineCompilationOptions::default(),
223                }),
224                primitive: wgpu::PrimitiveState {
225                    topology: wgpu::PrimitiveTopology::TriangleList,
226                    ..wgpu::PrimitiveState::default()
227                },
228                depth_stencil: Some(wgpu::DepthStencilState {
229                    format: wgpu::TextureFormat::Depth24Plus,
230                    depth_write_enabled: true,
231                    depth_compare: wgpu::CompareFunction::Less,
232                    stencil: wgpu::StencilState::default(),
233                    bias: wgpu::DepthBiasState::default(),
234                }),
235                multisample: wgpu::MultisampleState::default(),
236                multiview: None,
237                cache: None,
238            });
239
240        self.point_pick_pipeline = Some(pipeline);
241        self.pick_bind_group_layout = Some(bind_group_layout);
242    }
243
244    /// Gets the pick bind group layout.
245    pub fn pick_bind_group_layout(&self) -> &wgpu::BindGroupLayout {
246        self.pick_bind_group_layout
247            .as_ref()
248            .expect("pick pipeline not initialized")
249    }
250
251    /// Gets the point cloud pick pipeline.
252    pub fn point_pick_pipeline(&self) -> &wgpu::RenderPipeline {
253        self.point_pick_pipeline
254            .as_ref()
255            .expect("pick pipeline not initialized")
256    }
257
258    /// Gets the curve network pick pipeline.
259    pub fn curve_network_pick_pipeline(&self) -> &wgpu::RenderPipeline {
260        self.curve_network_pick_pipeline
261            .as_ref()
262            .expect("curve network pick pipeline not initialized")
263    }
264
265    /// Initializes the curve network pick pipeline.
266    pub fn init_curve_network_pick_pipeline(&mut self) {
267        let shader = self
268            .device
269            .create_shader_module(wgpu::ShaderModuleDescriptor {
270                label: Some("CurveNetwork Pick Shader"),
271                source: wgpu::ShaderSource::Wgsl(include_str!("../shaders/pick_curve.wgsl").into()),
272            });
273
274        // Reuse the pick bind group layout from point cloud pick
275        let bind_group_layout = self
276            .pick_bind_group_layout
277            .as_ref()
278            .expect("pick bind group layout not initialized");
279
280        let pipeline_layout = self
281            .device
282            .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
283                label: Some("CurveNetwork Pick Pipeline Layout"),
284                bind_group_layouts: &[bind_group_layout],
285                push_constant_ranges: &[],
286            });
287
288        let pipeline = self
289            .device
290            .create_render_pipeline(&wgpu::RenderPipelineDescriptor {
291                label: Some("CurveNetwork Pick Pipeline"),
292                layout: Some(&pipeline_layout),
293                vertex: wgpu::VertexState {
294                    module: &shader,
295                    entry_point: Some("vs_main"),
296                    buffers: &[],
297                    compilation_options: wgpu::PipelineCompilationOptions::default(),
298                },
299                fragment: Some(wgpu::FragmentState {
300                    module: &shader,
301                    entry_point: Some("fs_main"),
302                    targets: &[Some(wgpu::ColorTargetState {
303                        format: wgpu::TextureFormat::Rgba8Unorm,
304                        blend: None, // No blending for pick buffer
305                        write_mask: wgpu::ColorWrites::ALL,
306                    })],
307                    compilation_options: wgpu::PipelineCompilationOptions::default(),
308                }),
309                primitive: wgpu::PrimitiveState {
310                    topology: wgpu::PrimitiveTopology::LineList,
311                    ..wgpu::PrimitiveState::default()
312                },
313                depth_stencil: Some(wgpu::DepthStencilState {
314                    format: wgpu::TextureFormat::Depth24Plus,
315                    depth_write_enabled: true,
316                    depth_compare: wgpu::CompareFunction::Less,
317                    stencil: wgpu::StencilState::default(),
318                    bias: wgpu::DepthBiasState::default(),
319                }),
320                multisample: wgpu::MultisampleState::default(),
321                multiview: None,
322                cache: None,
323            });
324
325        self.curve_network_pick_pipeline = Some(pipeline);
326    }
327
328    /// Returns whether the curve network pick pipeline is initialized.
329    pub fn has_curve_network_pick_pipeline(&self) -> bool {
330        self.curve_network_pick_pipeline.is_some()
331    }
332
333    /// Initializes the curve network tube pick pipeline (uses ray-cylinder intersection).
334    pub fn init_curve_network_tube_pick_pipeline(&mut self) {
335        let shader = self
336            .device
337            .create_shader_module(wgpu::ShaderModuleDescriptor {
338                label: Some("CurveNetwork Tube Pick Shader"),
339                source: wgpu::ShaderSource::Wgsl(
340                    include_str!("../shaders/pick_curve_tube.wgsl").into(),
341                ),
342            });
343
344        // Create bind group layout for tube picking
345        let bind_group_layout =
346            self.device
347                .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
348                    label: Some("CurveNetwork Tube Pick Bind Group Layout"),
349                    entries: &[
350                        // Camera uniforms
351                        wgpu::BindGroupLayoutEntry {
352                            binding: 0,
353                            visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
354                            ty: wgpu::BindingType::Buffer {
355                                ty: wgpu::BufferBindingType::Uniform,
356                                has_dynamic_offset: false,
357                                min_binding_size: NonZeroU64::new(272),
358                            },
359                            count: None,
360                        },
361                        // Pick uniforms (global_start, radius, min_pick_radius)
362                        wgpu::BindGroupLayoutEntry {
363                            binding: 1,
364                            visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
365                            ty: wgpu::BindingType::Buffer {
366                                ty: wgpu::BufferBindingType::Uniform,
367                                has_dynamic_offset: false,
368                                min_binding_size: NonZeroU64::new(16),
369                            },
370                            count: None,
371                        },
372                        // Edge vertices (for raycast)
373                        wgpu::BindGroupLayoutEntry {
374                            binding: 2,
375                            visibility: wgpu::ShaderStages::FRAGMENT,
376                            ty: wgpu::BindingType::Buffer {
377                                ty: wgpu::BufferBindingType::Storage { read_only: true },
378                                has_dynamic_offset: false,
379                                min_binding_size: None,
380                            },
381                            count: None,
382                        },
383                    ],
384                });
385
386        let pipeline_layout = self
387            .device
388            .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
389                label: Some("CurveNetwork Tube Pick Pipeline Layout"),
390                bind_group_layouts: &[&bind_group_layout],
391                push_constant_ranges: &[],
392            });
393
394        let pipeline = self
395            .device
396            .create_render_pipeline(&wgpu::RenderPipelineDescriptor {
397                label: Some("CurveNetwork Tube Pick Pipeline"),
398                layout: Some(&pipeline_layout),
399                vertex: wgpu::VertexState {
400                    module: &shader,
401                    entry_point: Some("vs_main"),
402                    buffers: &[
403                        // Generated vertex buffer layout (same as tube render)
404                        wgpu::VertexBufferLayout {
405                            array_stride: 32, // vec4<f32> position + vec4<u32> edge_id_and_vertex_id
406                            step_mode: wgpu::VertexStepMode::Vertex,
407                            attributes: &[
408                                wgpu::VertexAttribute {
409                                    format: wgpu::VertexFormat::Float32x4,
410                                    offset: 0,
411                                    shader_location: 0,
412                                },
413                                wgpu::VertexAttribute {
414                                    format: wgpu::VertexFormat::Uint32x4,
415                                    offset: 16,
416                                    shader_location: 1,
417                                },
418                            ],
419                        },
420                    ],
421                    compilation_options: wgpu::PipelineCompilationOptions::default(),
422                },
423                fragment: Some(wgpu::FragmentState {
424                    module: &shader,
425                    entry_point: Some("fs_main"),
426                    targets: &[Some(wgpu::ColorTargetState {
427                        format: wgpu::TextureFormat::Rgba8Unorm,
428                        blend: None, // No blending for pick buffer
429                        write_mask: wgpu::ColorWrites::ALL,
430                    })],
431                    compilation_options: wgpu::PipelineCompilationOptions::default(),
432                }),
433                primitive: wgpu::PrimitiveState {
434                    topology: wgpu::PrimitiveTopology::TriangleList,
435                    ..wgpu::PrimitiveState::default()
436                },
437                depth_stencil: Some(wgpu::DepthStencilState {
438                    format: wgpu::TextureFormat::Depth24Plus,
439                    depth_write_enabled: true,
440                    depth_compare: wgpu::CompareFunction::Less,
441                    stencil: wgpu::StencilState::default(),
442                    bias: wgpu::DepthBiasState::default(),
443                }),
444                multisample: wgpu::MultisampleState::default(),
445                multiview: None,
446                cache: None,
447            });
448
449        self.curve_network_tube_pick_pipeline = Some(pipeline);
450        self.curve_network_tube_pick_bind_group_layout = Some(bind_group_layout);
451    }
452
453    /// Returns whether the curve network tube pick pipeline is initialized.
454    pub fn has_curve_network_tube_pick_pipeline(&self) -> bool {
455        self.curve_network_tube_pick_pipeline.is_some()
456    }
457
458    /// Gets the curve network tube pick pipeline.
459    pub fn curve_network_tube_pick_pipeline(&self) -> &wgpu::RenderPipeline {
460        self.curve_network_tube_pick_pipeline
461            .as_ref()
462            .expect("curve network tube pick pipeline not initialized")
463    }
464
465    /// Gets the curve network tube pick bind group layout.
466    pub fn curve_network_tube_pick_bind_group_layout(&self) -> &wgpu::BindGroupLayout {
467        self.curve_network_tube_pick_bind_group_layout
468            .as_ref()
469            .expect("curve network tube pick bind group layout not initialized")
470    }
471
472    /// Initializes the mesh pick pipeline for surface meshes.
473    pub fn init_mesh_pick_pipeline(&mut self) {
474        let shader = self
475            .device
476            .create_shader_module(wgpu::ShaderModuleDescriptor {
477                label: Some("Mesh Pick Shader"),
478                source: wgpu::ShaderSource::Wgsl(include_str!("../shaders/pick_mesh.wgsl").into()),
479            });
480
481        // Mesh pick bind group layout: camera, mesh pick uniforms, positions, face_indices
482        let bind_group_layout =
483            self.device
484                .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
485                    label: Some("Mesh Pick Bind Group Layout"),
486                    entries: &[
487                        // Camera uniforms
488                        wgpu::BindGroupLayoutEntry {
489                            binding: 0,
490                            visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
491                            ty: wgpu::BindingType::Buffer {
492                                ty: wgpu::BufferBindingType::Uniform,
493                                has_dynamic_offset: false,
494                                min_binding_size: NonZeroU64::new(272),
495                            },
496                            count: None,
497                        },
498                        // Mesh pick uniforms (global_start + model matrix = 80 bytes)
499                        wgpu::BindGroupLayoutEntry {
500                            binding: 1,
501                            visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
502                            ty: wgpu::BindingType::Buffer {
503                                ty: wgpu::BufferBindingType::Uniform,
504                                has_dynamic_offset: false,
505                                min_binding_size: NonZeroU64::new(80),
506                            },
507                            count: None,
508                        },
509                        // Position storage buffer (expanded per-triangle-vertex)
510                        wgpu::BindGroupLayoutEntry {
511                            binding: 2,
512                            visibility: wgpu::ShaderStages::VERTEX,
513                            ty: wgpu::BindingType::Buffer {
514                                ty: wgpu::BufferBindingType::Storage { read_only: true },
515                                has_dynamic_offset: false,
516                                min_binding_size: None,
517                            },
518                            count: None,
519                        },
520                        // Face index mapping buffer (tri_index -> face_index)
521                        wgpu::BindGroupLayoutEntry {
522                            binding: 3,
523                            visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
524                            ty: wgpu::BindingType::Buffer {
525                                ty: wgpu::BufferBindingType::Storage { read_only: true },
526                                has_dynamic_offset: false,
527                                min_binding_size: None,
528                            },
529                            count: None,
530                        },
531                    ],
532                });
533
534        let pipeline_layout = self
535            .device
536            .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
537                label: Some("Mesh Pick Pipeline Layout"),
538                bind_group_layouts: &[&bind_group_layout],
539                push_constant_ranges: &[],
540            });
541
542        let pipeline = self
543            .device
544            .create_render_pipeline(&wgpu::RenderPipelineDescriptor {
545                label: Some("SurfaceMesh Pick Pipeline"),
546                layout: Some(&pipeline_layout),
547                vertex: wgpu::VertexState {
548                    module: &shader,
549                    entry_point: Some("vs_main"),
550                    buffers: &[],
551                    compilation_options: wgpu::PipelineCompilationOptions::default(),
552                },
553                fragment: Some(wgpu::FragmentState {
554                    module: &shader,
555                    entry_point: Some("fs_main"),
556                    targets: &[Some(wgpu::ColorTargetState {
557                        format: wgpu::TextureFormat::Rgba8Unorm,
558                        blend: None,
559                        write_mask: wgpu::ColorWrites::ALL,
560                    })],
561                    compilation_options: wgpu::PipelineCompilationOptions::default(),
562                }),
563                primitive: wgpu::PrimitiveState {
564                    topology: wgpu::PrimitiveTopology::TriangleList,
565                    ..wgpu::PrimitiveState::default()
566                },
567                depth_stencil: Some(wgpu::DepthStencilState {
568                    format: wgpu::TextureFormat::Depth24Plus,
569                    depth_write_enabled: true,
570                    depth_compare: wgpu::CompareFunction::Less,
571                    stencil: wgpu::StencilState::default(),
572                    bias: wgpu::DepthBiasState::default(),
573                }),
574                multisample: wgpu::MultisampleState::default(),
575                multiview: None,
576                cache: None,
577            });
578
579        self.mesh_pick_pipeline = Some(pipeline);
580        self.mesh_pick_bind_group_layout = Some(bind_group_layout);
581    }
582
583    /// Returns whether the mesh pick pipeline is initialized.
584    pub fn has_mesh_pick_pipeline(&self) -> bool {
585        self.mesh_pick_pipeline.is_some()
586    }
587
588    /// Gets the mesh pick pipeline.
589    pub fn mesh_pick_pipeline(&self) -> &wgpu::RenderPipeline {
590        self.mesh_pick_pipeline
591            .as_ref()
592            .expect("mesh pick pipeline not initialized")
593    }
594
595    /// Gets the mesh pick bind group layout.
596    pub fn mesh_pick_bind_group_layout(&self) -> &wgpu::BindGroupLayout {
597        self.mesh_pick_bind_group_layout
598            .as_ref()
599            .expect("mesh pick bind group layout not initialized")
600    }
601
602    /// Initializes the gridcube pick pipeline for volume grid instances.
603    pub fn init_gridcube_pick_pipeline(&mut self) {
604        let shader = self
605            .device
606            .create_shader_module(wgpu::ShaderModuleDescriptor {
607                label: Some("Gridcube Pick Shader"),
608                source: wgpu::ShaderSource::Wgsl(
609                    include_str!("../shaders/pick_gridcube.wgsl").into(),
610                ),
611            });
612
613        // Gridcube pick bind group layout: camera, GridcubePickUniforms, positions
614        let bind_group_layout =
615            self.device
616                .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
617                    label: Some("Gridcube Pick Bind Group Layout"),
618                    entries: &[
619                        // Camera uniforms (binding 0)
620                        wgpu::BindGroupLayoutEntry {
621                            binding: 0,
622                            visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
623                            ty: wgpu::BindingType::Buffer {
624                                ty: wgpu::BufferBindingType::Uniform,
625                                has_dynamic_offset: false,
626                                min_binding_size: NonZeroU64::new(272),
627                            },
628                            count: None,
629                        },
630                        // GridcubePickUniforms (model + global_start + cube_size_factor = 80 bytes) (binding 1)
631                        wgpu::BindGroupLayoutEntry {
632                            binding: 1,
633                            visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
634                            ty: wgpu::BindingType::Buffer {
635                                ty: wgpu::BufferBindingType::Uniform,
636                                has_dynamic_offset: false,
637                                min_binding_size: NonZeroU64::new(80),
638                            },
639                            count: None,
640                        },
641                        // Cube positions storage buffer (template + per-instance) (binding 2)
642                        wgpu::BindGroupLayoutEntry {
643                            binding: 2,
644                            visibility: wgpu::ShaderStages::VERTEX,
645                            ty: wgpu::BindingType::Buffer {
646                                ty: wgpu::BufferBindingType::Storage { read_only: true },
647                                has_dynamic_offset: false,
648                                min_binding_size: None,
649                            },
650                            count: None,
651                        },
652                    ],
653                });
654
655        let pipeline_layout = self
656            .device
657            .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
658                label: Some("Gridcube Pick Pipeline Layout"),
659                bind_group_layouts: &[&bind_group_layout],
660                push_constant_ranges: &[],
661            });
662
663        let pipeline = self
664            .device
665            .create_render_pipeline(&wgpu::RenderPipelineDescriptor {
666                label: Some("Gridcube Pick Pipeline"),
667                layout: Some(&pipeline_layout),
668                vertex: wgpu::VertexState {
669                    module: &shader,
670                    entry_point: Some("vs_main"),
671                    buffers: &[],
672                    compilation_options: wgpu::PipelineCompilationOptions::default(),
673                },
674                fragment: Some(wgpu::FragmentState {
675                    module: &shader,
676                    entry_point: Some("fs_main"),
677                    targets: &[Some(wgpu::ColorTargetState {
678                        format: wgpu::TextureFormat::Rgba8Unorm,
679                        blend: None,
680                        write_mask: wgpu::ColorWrites::ALL,
681                    })],
682                    compilation_options: wgpu::PipelineCompilationOptions::default(),
683                }),
684                primitive: wgpu::PrimitiveState {
685                    topology: wgpu::PrimitiveTopology::TriangleList,
686                    ..wgpu::PrimitiveState::default()
687                },
688                depth_stencil: Some(wgpu::DepthStencilState {
689                    format: wgpu::TextureFormat::Depth24Plus,
690                    depth_write_enabled: true,
691                    depth_compare: wgpu::CompareFunction::Less,
692                    stencil: wgpu::StencilState::default(),
693                    bias: wgpu::DepthBiasState::default(),
694                }),
695                multisample: wgpu::MultisampleState::default(),
696                multiview: None,
697                cache: None,
698            });
699
700        self.gridcube_pick_pipeline = Some(pipeline);
701        self.gridcube_pick_bind_group_layout = Some(bind_group_layout);
702    }
703
704    /// Returns whether the gridcube pick pipeline is initialized.
705    pub fn has_gridcube_pick_pipeline(&self) -> bool {
706        self.gridcube_pick_pipeline.is_some()
707    }
708
709    /// Gets the gridcube pick pipeline.
710    pub fn gridcube_pick_pipeline(&self) -> &wgpu::RenderPipeline {
711        self.gridcube_pick_pipeline
712            .as_ref()
713            .expect("gridcube pick pipeline not initialized")
714    }
715
716    /// Gets the gridcube pick bind group layout.
717    pub fn gridcube_pick_bind_group_layout(&self) -> &wgpu::BindGroupLayout {
718        self.gridcube_pick_bind_group_layout
719            .as_ref()
720            .expect("gridcube pick bind group layout not initialized")
721    }
722
723    /// Reads the pick buffer at (x, y) and returns the decoded global index.
724    ///
725    /// Returns None if picking system not initialized or coordinates out of bounds.
726    /// Returns Some(0) for background clicks.
727    pub fn pick_at(&self, x: u32, y: u32) -> Option<u32> {
728        let pick_texture = self.pick_texture.as_ref()?;
729        let staging_buffer = self.pick_staging_buffer.as_ref()?;
730
731        // Bounds check
732        let (width, height) = self.pick_buffer_size;
733        if x >= width || y >= height {
734            return None;
735        }
736
737        // Create encoder for copy operation
738        let mut encoder = self
739            .device
740            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
741                label: Some("Pick Readback Encoder"),
742            });
743
744        // Copy single pixel from pick texture to staging buffer
745        encoder.copy_texture_to_buffer(
746            wgpu::TexelCopyTextureInfo {
747                texture: pick_texture,
748                mip_level: 0,
749                origin: wgpu::Origin3d { x, y, z: 0 },
750                aspect: wgpu::TextureAspect::All,
751            },
752            wgpu::TexelCopyBufferInfo {
753                buffer: staging_buffer,
754                layout: wgpu::TexelCopyBufferLayout {
755                    offset: 0,
756                    bytes_per_row: Some(256), // Aligned
757                    rows_per_image: Some(1),
758                },
759            },
760            wgpu::Extent3d {
761                width: 1,
762                height: 1,
763                depth_or_array_layers: 1,
764            },
765        );
766
767        self.queue.submit(std::iter::once(encoder.finish()));
768
769        // Map buffer and read pixel
770        let buffer_slice = staging_buffer.slice(..4);
771        let (tx, rx) = std::sync::mpsc::channel();
772        buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
773            tx.send(result).unwrap();
774        });
775
776        let _ = self.device.poll(wgpu::PollType::wait_indefinitely());
777        rx.recv().unwrap().ok()?;
778
779        let data = buffer_slice.get_mapped_range();
780        let pixel: [u8; 4] = [data[0], data[1], data[2], data[3]];
781        drop(data);
782        staging_buffer.unmap();
783
784        // Decode flat 24-bit global index from RGB
785        let global_index = crate::pick::color_to_index(pixel[0], pixel[1], pixel[2]);
786        Some(global_index)
787    }
788
789    /// Returns the pick texture view for external rendering.
790    pub fn pick_texture_view(&self) -> Option<&wgpu::TextureView> {
791        self.pick_texture_view.as_ref()
792    }
793
794    /// Returns the pick depth texture view for external rendering.
795    pub fn pick_depth_view(&self) -> Option<&wgpu::TextureView> {
796        self.pick_depth_view.as_ref()
797    }
798
799    /// Begins a pick render pass. Returns the render pass encoder.
800    ///
801    /// The caller is responsible for rendering structures to this pass
802    /// and then dropping the encoder to finish the pass.
803    pub fn begin_pick_pass<'a>(
804        &'a self,
805        encoder: &'a mut wgpu::CommandEncoder,
806    ) -> Option<wgpu::RenderPass<'a>> {
807        let pick_view = self.pick_texture_view.as_ref()?;
808        let pick_depth = self.pick_depth_view.as_ref()?;
809
810        Some(encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
811            label: Some("Pick Render Pass"),
812            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
813                view: pick_view,
814                resolve_target: None,
815                ops: wgpu::Operations {
816                    load: wgpu::LoadOp::Clear(wgpu::Color::BLACK), // Background = (0,0,0) = index 0
817                    store: wgpu::StoreOp::Store,
818                },
819                depth_slice: None,
820            })],
821            depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
822                view: pick_depth,
823                depth_ops: Some(wgpu::Operations {
824                    load: wgpu::LoadOp::Clear(1.0),
825                    store: wgpu::StoreOp::Store,
826                }),
827                stencil_ops: None,
828            }),
829            ..Default::default()
830        }))
831    }
832}
polyscope_render/engine/pick.rs

polyscope_render/engine/
pick.rs
