Skip to main content

viewport_lib/renderer/
picking.rs

1use super::*;
2
3impl ViewportRenderer {
4    // -----------------------------------------------------------------------
5    // Phase K : GPU object-ID picking
6    // -----------------------------------------------------------------------
7
8    /// GPU object-ID pick: renders the scene to an offscreen `R32Uint` texture
9    /// and reads back the single pixel under `cursor`.
10    ///
11    /// This is O(1) in mesh complexity : every object is rendered with a flat
12    /// `u32` ID, and only one pixel is read back. For triangle-level queries
13    /// (barycentric scalar probe, exact world position), use the CPU
14    /// [`crate::interaction::picking::pick_scene_cpu`] path instead.
15    ///
16    /// The pipeline is lazily initialized on first call : zero overhead when
17    /// this method is never invoked.
18    ///
19    /// # Arguments
20    /// * `device` : wgpu device
21    /// * `queue` : wgpu queue
22    /// * `cursor` : cursor position in viewport-local pixels (top-left origin)
23    /// * `frame` : current grouped frame data (camera, scene surfaces, viewport size)
24    ///
25    /// # Returns
26    /// `Some(GpuPickHit)` if an object is under the cursor, `None` if empty space.
27    pub fn pick_scene_gpu(
28        &mut self,
29        device: &wgpu::Device,
30        queue: &wgpu::Queue,
31        cursor: glam::Vec2,
32        frame: &FrameData,
33    ) -> Option<crate::interaction::picking::GpuPickHit> {
34        // In Playback mode, throttle picking to every 4th frame to reduce overhead
35        // during animation. Interactive, Paused, and Capture modes always pick.
36        if self.runtime_mode == crate::renderer::stats::RuntimeMode::Playback
37            && self.frame_counter % 4 != 0
38        {
39            return None;
40        }
41
42        // Read scene items from the surface submission.
43        let scene_items: &[SceneRenderItem] = match &frame.scene.surfaces {
44            SurfaceSubmission::Flat(items) => items.as_ref(),
45        };
46
47        let vp_w = frame.camera.viewport_size[0] as u32;
48        let vp_h = frame.camera.viewport_size[1] as u32;
49
50        // --- bounds check ---
51        if cursor.x < 0.0
52            || cursor.y < 0.0
53            || cursor.x >= frame.camera.viewport_size[0]
54            || cursor.y >= frame.camera.viewport_size[1]
55            || vp_w == 0
56            || vp_h == 0
57        {
58            return None;
59        }
60
61        // --- lazy pipeline init ---
62        self.resources.ensure_pick_pipeline(device);
63
64        // --- build PickInstance data ---
65        // Only surfaces with a nonzero pick_id participate in picking.
66        // Clear value 0 means "no hit" (or non-pickable surface).
67        let pickable_items: Vec<&SceneRenderItem> = scene_items
68            .iter()
69            .filter(|item| item.visible && item.pick_id != PickId::NONE)
70            .collect();
71
72        let pick_instances: Vec<PickInstance> = pickable_items
73            .iter()
74            .map(|item| {
75                let m = item.model;
76                PickInstance {
77                    model_c0: m[0],
78                    model_c1: m[1],
79                    model_c2: m[2],
80                    model_c3: m[3],
81                    object_id: item.pick_id.0 as u32,
82                    _pad: [0; 3],
83                }
84            })
85            .collect();
86
87        if pick_instances.is_empty() {
88            return None;
89        }
90
91        // --- pick instance storage buffer + bind group ---
92        let pick_instance_bytes = bytemuck::cast_slice(&pick_instances);
93        let pick_instance_buf = device.create_buffer(&wgpu::BufferDescriptor {
94            label: Some("pick_instance_buf"),
95            size: pick_instance_bytes.len().max(80) as u64,
96            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
97            mapped_at_creation: false,
98        });
99        queue.write_buffer(&pick_instance_buf, 0, pick_instance_bytes);
100
101        let pick_instance_bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
102            label: Some("pick_instance_bg"),
103            layout: self
104                .resources
105                .pick_bind_group_layout_1
106                .as_ref()
107                .expect("ensure_pick_pipeline must be called first"),
108            entries: &[wgpu::BindGroupEntry {
109                binding: 0,
110                resource: pick_instance_buf.as_entire_binding(),
111            }],
112        });
113
114        // --- pick camera uniform buffer + bind group ---
115        let camera_uniform = frame.camera.render_camera.camera_uniform();
116        let camera_bytes = bytemuck::bytes_of(&camera_uniform);
117        let pick_camera_buf = device.create_buffer(&wgpu::BufferDescriptor {
118            label: Some("pick_camera_buf"),
119            size: std::mem::size_of::<CameraUniform>() as u64,
120            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
121            mapped_at_creation: false,
122        });
123        queue.write_buffer(&pick_camera_buf, 0, camera_bytes);
124
125        let pick_camera_bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
126            label: Some("pick_camera_bg"),
127            layout: self
128                .resources
129                .pick_camera_bgl
130                .as_ref()
131                .expect("ensure_pick_pipeline must be called first"),
132            entries: &[wgpu::BindGroupEntry {
133                binding: 0,
134                resource: pick_camera_buf.as_entire_binding(),
135            }],
136        });
137
138        // --- offscreen pick textures (R32Uint + R32Float) + depth ---
139        let pick_id_texture = device.create_texture(&wgpu::TextureDescriptor {
140            label: Some("pick_id_texture"),
141            size: wgpu::Extent3d {
142                width: vp_w,
143                height: vp_h,
144                depth_or_array_layers: 1,
145            },
146            mip_level_count: 1,
147            sample_count: 1,
148            dimension: wgpu::TextureDimension::D2,
149            format: wgpu::TextureFormat::R32Uint,
150            usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
151            view_formats: &[],
152        });
153        let pick_id_view = pick_id_texture.create_view(&wgpu::TextureViewDescriptor::default());
154
155        let pick_depth_texture = device.create_texture(&wgpu::TextureDescriptor {
156            label: Some("pick_depth_color_texture"),
157            size: wgpu::Extent3d {
158                width: vp_w,
159                height: vp_h,
160                depth_or_array_layers: 1,
161            },
162            mip_level_count: 1,
163            sample_count: 1,
164            dimension: wgpu::TextureDimension::D2,
165            format: wgpu::TextureFormat::R32Float,
166            usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
167            view_formats: &[],
168        });
169        let pick_depth_view =
170            pick_depth_texture.create_view(&wgpu::TextureViewDescriptor::default());
171
172        let depth_stencil_texture = device.create_texture(&wgpu::TextureDescriptor {
173            label: Some("pick_ds_texture"),
174            size: wgpu::Extent3d {
175                width: vp_w,
176                height: vp_h,
177                depth_or_array_layers: 1,
178            },
179            mip_level_count: 1,
180            sample_count: 1,
181            dimension: wgpu::TextureDimension::D2,
182            format: wgpu::TextureFormat::Depth24PlusStencil8,
183            usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
184            view_formats: &[],
185        });
186        let depth_stencil_view =
187            depth_stencil_texture.create_view(&wgpu::TextureViewDescriptor::default());
188
189        // --- render pass ---
190        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
191            label: Some("pick_pass_encoder"),
192        });
193        {
194            let mut pick_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
195                label: Some("pick_pass"),
196                color_attachments: &[
197                    Some(wgpu::RenderPassColorAttachment {
198                        view: &pick_id_view,
199                        resolve_target: None,
200                        depth_slice: None,
201                        ops: wgpu::Operations {
202                            load: wgpu::LoadOp::Clear(wgpu::Color {
203                                r: 0.0,
204                                g: 0.0,
205                                b: 0.0,
206                                a: 0.0,
207                            }),
208                            store: wgpu::StoreOp::Store,
209                        },
210                    }),
211                    Some(wgpu::RenderPassColorAttachment {
212                        view: &pick_depth_view,
213                        resolve_target: None,
214                        depth_slice: None,
215                        ops: wgpu::Operations {
216                            load: wgpu::LoadOp::Clear(wgpu::Color {
217                                r: 1.0,
218                                g: 0.0,
219                                b: 0.0,
220                                a: 0.0,
221                            }),
222                            store: wgpu::StoreOp::Store,
223                        },
224                    }),
225                ],
226                depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
227                    view: &depth_stencil_view,
228                    depth_ops: Some(wgpu::Operations {
229                        load: wgpu::LoadOp::Clear(1.0),
230                        store: wgpu::StoreOp::Store,
231                    }),
232                    stencil_ops: None,
233                }),
234                timestamp_writes: None,
235                occlusion_query_set: None,
236            });
237
238            pick_pass.set_pipeline(
239                self.resources
240                    .pick_pipeline
241                    .as_ref()
242                    .expect("ensure_pick_pipeline must be called first"),
243            );
244            pick_pass.set_bind_group(0, &pick_camera_bg, &[]);
245            pick_pass.set_bind_group(1, &pick_instance_bg, &[]);
246
247            // Draw each pickable item with its instance slot.
248            // Instance index in the storage buffer = position in pick_instances vec.
249            for (instance_slot, item) in pickable_items.iter().enumerate() {
250                let Some(mesh) = self
251                    .resources
252                    .mesh_store
253                    .get(item.mesh_id)
254                else {
255                    continue;
256                };
257                pick_pass.set_vertex_buffer(0, mesh.vertex_buffer.slice(..));
258                pick_pass.set_index_buffer(mesh.index_buffer.slice(..), wgpu::IndexFormat::Uint32);
259                let slot = instance_slot as u32;
260                pick_pass.draw_indexed(0..mesh.index_count, 0, slot..slot + 1);
261            }
262        }
263
264        // --- copy 1×1 pixels to staging buffers ---
265        // R32Uint: 4 bytes per pixel, min bytes_per_row = 256 (wgpu alignment)
266        let bytes_per_row_aligned = 256u32; // wgpu requires multiples of 256
267
268        let id_staging = device.create_buffer(&wgpu::BufferDescriptor {
269            label: Some("pick_id_staging"),
270            size: bytes_per_row_aligned as u64,
271            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
272            mapped_at_creation: false,
273        });
274        let depth_staging = device.create_buffer(&wgpu::BufferDescriptor {
275            label: Some("pick_depth_staging"),
276            size: bytes_per_row_aligned as u64,
277            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
278            mapped_at_creation: false,
279        });
280
281        let px = cursor.x as u32;
282        let py = cursor.y as u32;
283
284        encoder.copy_texture_to_buffer(
285            wgpu::TexelCopyTextureInfo {
286                texture: &pick_id_texture,
287                mip_level: 0,
288                origin: wgpu::Origin3d { x: px, y: py, z: 0 },
289                aspect: wgpu::TextureAspect::All,
290            },
291            wgpu::TexelCopyBufferInfo {
292                buffer: &id_staging,
293                layout: wgpu::TexelCopyBufferLayout {
294                    offset: 0,
295                    bytes_per_row: Some(bytes_per_row_aligned),
296                    rows_per_image: Some(1),
297                },
298            },
299            wgpu::Extent3d {
300                width: 1,
301                height: 1,
302                depth_or_array_layers: 1,
303            },
304        );
305        encoder.copy_texture_to_buffer(
306            wgpu::TexelCopyTextureInfo {
307                texture: &pick_depth_texture,
308                mip_level: 0,
309                origin: wgpu::Origin3d { x: px, y: py, z: 0 },
310                aspect: wgpu::TextureAspect::All,
311            },
312            wgpu::TexelCopyBufferInfo {
313                buffer: &depth_staging,
314                layout: wgpu::TexelCopyBufferLayout {
315                    offset: 0,
316                    bytes_per_row: Some(bytes_per_row_aligned),
317                    rows_per_image: Some(1),
318                },
319            },
320            wgpu::Extent3d {
321                width: 1,
322                height: 1,
323                depth_or_array_layers: 1,
324            },
325        );
326
327        queue.submit(std::iter::once(encoder.finish()));
328
329        // --- map and read ---
330        let (tx_id, rx_id) = std::sync::mpsc::channel::<Result<(), wgpu::BufferAsyncError>>();
331        let (tx_dep, rx_dep) = std::sync::mpsc::channel::<Result<(), wgpu::BufferAsyncError>>();
332        id_staging
333            .slice(..)
334            .map_async(wgpu::MapMode::Read, move |r| {
335                let _ = tx_id.send(r);
336            });
337        depth_staging
338            .slice(..)
339            .map_async(wgpu::MapMode::Read, move |r| {
340                let _ = tx_dep.send(r);
341            });
342        device
343            .poll(wgpu::PollType::Wait {
344                submission_index: None,
345                timeout: Some(std::time::Duration::from_secs(5)),
346            })
347            .unwrap();
348        let _ = rx_id.recv().unwrap_or(Err(wgpu::BufferAsyncError));
349        let _ = rx_dep.recv().unwrap_or(Err(wgpu::BufferAsyncError));
350
351        let object_id = {
352            let data = id_staging.slice(..).get_mapped_range();
353            u32::from_le_bytes([data[0], data[1], data[2], data[3]])
354        };
355        id_staging.unmap();
356
357        let depth = {
358            let data = depth_staging.slice(..).get_mapped_range();
359            f32::from_le_bytes([data[0], data[1], data[2], data[3]])
360        };
361        depth_staging.unmap();
362
363        // 0 = miss (clear color or non-pickable surface).
364        if object_id == 0 {
365            return None;
366        }
367
368        Some(crate::interaction::picking::GpuPickHit {
369            object_id: PickId(object_id as u64),
370            depth,
371        })
372    }
373}