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`] 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        // Resolve scene items from the SurfaceSubmission seam.
35        let scene_items: &[SceneRenderItem] = match &frame.scene.surfaces {
36            SurfaceSubmission::Flat(items) => items,
37        };
38
39        let vp_w = frame.camera.viewport_size[0] as u32;
40        let vp_h = frame.camera.viewport_size[1] as u32;
41
42        // --- bounds check ---
43        if cursor.x < 0.0
44            || cursor.y < 0.0
45            || cursor.x >= frame.camera.viewport_size[0]
46            || cursor.y >= frame.camera.viewport_size[1]
47            || vp_w == 0
48            || vp_h == 0
49        {
50            return None;
51        }
52
53        // --- lazy pipeline init ---
54        self.resources.ensure_pick_pipeline(device);
55
56        // --- build PickInstance data ---
57        // Sentinel scheme: object_id stored = (scene_items_index + 1) so that
58        // clear value 0 unambiguously means "no hit".
59        let pick_instances: Vec<PickInstance> = scene_items
60            .iter()
61            .enumerate()
62            .filter(|(_, item)| item.visible)
63            .map(|(idx, item)| {
64                let m = item.model;
65                PickInstance {
66                    model_c0: m[0],
67                    model_c1: m[1],
68                    model_c2: m[2],
69                    model_c3: m[3],
70                    object_id: (idx + 1) as u32,
71                    _pad: [0; 3],
72                }
73            })
74            .collect();
75
76        if pick_instances.is_empty() {
77            return None;
78        }
79
80        // Build a mapping from sentinel object_id -> original scene_items index.
81        // Also track which scene_items are visible and their scene_items indices
82        // so we can issue the right draw calls.
83        let visible_items: Vec<(usize, &SceneRenderItem)> = scene_items
84            .iter()
85            .enumerate()
86            .filter(|(_, item)| item.visible)
87            .collect();
88
89        // --- pick instance storage buffer + bind group ---
90        let pick_instance_bytes = bytemuck::cast_slice(&pick_instances);
91        let pick_instance_buf = device.create_buffer(&wgpu::BufferDescriptor {
92            label: Some("pick_instance_buf"),
93            size: pick_instance_bytes.len().max(80) as u64,
94            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
95            mapped_at_creation: false,
96        });
97        queue.write_buffer(&pick_instance_buf, 0, pick_instance_bytes);
98
99        let pick_instance_bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
100            label: Some("pick_instance_bg"),
101            layout: self
102                .resources
103                .pick_bind_group_layout_1
104                .as_ref()
105                .expect("ensure_pick_pipeline must be called first"),
106            entries: &[wgpu::BindGroupEntry {
107                binding: 0,
108                resource: pick_instance_buf.as_entire_binding(),
109            }],
110        });
111
112        // --- pick camera uniform buffer + bind group ---
113        let camera_uniform = frame.camera.render_camera.camera_uniform();
114        let camera_bytes = bytemuck::bytes_of(&camera_uniform);
115        let pick_camera_buf = device.create_buffer(&wgpu::BufferDescriptor {
116            label: Some("pick_camera_buf"),
117            size: std::mem::size_of::<CameraUniform>() as u64,
118            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
119            mapped_at_creation: false,
120        });
121        queue.write_buffer(&pick_camera_buf, 0, camera_bytes);
122
123        let pick_camera_bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
124            label: Some("pick_camera_bg"),
125            layout: self
126                .resources
127                .pick_camera_bgl
128                .as_ref()
129                .expect("ensure_pick_pipeline must be called first"),
130            entries: &[wgpu::BindGroupEntry {
131                binding: 0,
132                resource: pick_camera_buf.as_entire_binding(),
133            }],
134        });
135
136        // --- offscreen pick textures (R32Uint + R32Float) + depth ---
137        let pick_id_texture = device.create_texture(&wgpu::TextureDescriptor {
138            label: Some("pick_id_texture"),
139            size: wgpu::Extent3d {
140                width: vp_w,
141                height: vp_h,
142                depth_or_array_layers: 1,
143            },
144            mip_level_count: 1,
145            sample_count: 1,
146            dimension: wgpu::TextureDimension::D2,
147            format: wgpu::TextureFormat::R32Uint,
148            usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
149            view_formats: &[],
150        });
151        let pick_id_view = pick_id_texture.create_view(&wgpu::TextureViewDescriptor::default());
152
153        let pick_depth_texture = device.create_texture(&wgpu::TextureDescriptor {
154            label: Some("pick_depth_color_texture"),
155            size: wgpu::Extent3d {
156                width: vp_w,
157                height: vp_h,
158                depth_or_array_layers: 1,
159            },
160            mip_level_count: 1,
161            sample_count: 1,
162            dimension: wgpu::TextureDimension::D2,
163            format: wgpu::TextureFormat::R32Float,
164            usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
165            view_formats: &[],
166        });
167        let pick_depth_view =
168            pick_depth_texture.create_view(&wgpu::TextureViewDescriptor::default());
169
170        let depth_stencil_texture = device.create_texture(&wgpu::TextureDescriptor {
171            label: Some("pick_ds_texture"),
172            size: wgpu::Extent3d {
173                width: vp_w,
174                height: vp_h,
175                depth_or_array_layers: 1,
176            },
177            mip_level_count: 1,
178            sample_count: 1,
179            dimension: wgpu::TextureDimension::D2,
180            format: wgpu::TextureFormat::Depth24PlusStencil8,
181            usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
182            view_formats: &[],
183        });
184        let depth_stencil_view =
185            depth_stencil_texture.create_view(&wgpu::TextureViewDescriptor::default());
186
187        // --- render pass ---
188        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
189            label: Some("pick_pass_encoder"),
190        });
191        {
192            let mut pick_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
193                label: Some("pick_pass"),
194                color_attachments: &[
195                    Some(wgpu::RenderPassColorAttachment {
196                        view: &pick_id_view,
197                        resolve_target: None,
198                        depth_slice: None,
199                        ops: wgpu::Operations {
200                            load: wgpu::LoadOp::Clear(wgpu::Color {
201                                r: 0.0,
202                                g: 0.0,
203                                b: 0.0,
204                                a: 0.0,
205                            }),
206                            store: wgpu::StoreOp::Store,
207                        },
208                    }),
209                    Some(wgpu::RenderPassColorAttachment {
210                        view: &pick_depth_view,
211                        resolve_target: None,
212                        depth_slice: None,
213                        ops: wgpu::Operations {
214                            load: wgpu::LoadOp::Clear(wgpu::Color {
215                                r: 1.0,
216                                g: 0.0,
217                                b: 0.0,
218                                a: 0.0,
219                            }),
220                            store: wgpu::StoreOp::Store,
221                        },
222                    }),
223                ],
224                depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
225                    view: &depth_stencil_view,
226                    depth_ops: Some(wgpu::Operations {
227                        load: wgpu::LoadOp::Clear(1.0),
228                        store: wgpu::StoreOp::Store,
229                    }),
230                    stencil_ops: None,
231                }),
232                timestamp_writes: None,
233                occlusion_query_set: None,
234            });
235
236            pick_pass.set_pipeline(
237                self.resources
238                    .pick_pipeline
239                    .as_ref()
240                    .expect("ensure_pick_pipeline must be called first"),
241            );
242            pick_pass.set_bind_group(0, &pick_camera_bg, &[]);
243            pick_pass.set_bind_group(1, &pick_instance_bg, &[]);
244
245            // Draw each visible item with its instance slot.
246            // Instance index in the storage buffer = position in pick_instances vec.
247            for (instance_slot, (_, item)) in visible_items.iter().enumerate() {
248                let Some(mesh) = self
249                    .resources
250                    .mesh_store
251                    .get(crate::resources::mesh_store::MeshId(item.mesh_index))
252                else {
253                    continue;
254                };
255                pick_pass.set_vertex_buffer(0, mesh.vertex_buffer.slice(..));
256                pick_pass.set_index_buffer(mesh.index_buffer.slice(..), wgpu::IndexFormat::Uint32);
257                let slot = instance_slot as u32;
258                pick_pass.draw_indexed(0..mesh.index_count, 0, slot..slot + 1);
259            }
260        }
261
262        // --- copy 1×1 pixels to staging buffers ---
263        // R32Uint: 4 bytes per pixel, min bytes_per_row = 256 (wgpu alignment)
264        let bytes_per_row_aligned = 256u32; // wgpu requires multiples of 256
265
266        let id_staging = device.create_buffer(&wgpu::BufferDescriptor {
267            label: Some("pick_id_staging"),
268            size: bytes_per_row_aligned as u64,
269            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
270            mapped_at_creation: false,
271        });
272        let depth_staging = device.create_buffer(&wgpu::BufferDescriptor {
273            label: Some("pick_depth_staging"),
274            size: bytes_per_row_aligned as u64,
275            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
276            mapped_at_creation: false,
277        });
278
279        let px = cursor.x as u32;
280        let py = cursor.y as u32;
281
282        encoder.copy_texture_to_buffer(
283            wgpu::TexelCopyTextureInfo {
284                texture: &pick_id_texture,
285                mip_level: 0,
286                origin: wgpu::Origin3d { x: px, y: py, z: 0 },
287                aspect: wgpu::TextureAspect::All,
288            },
289            wgpu::TexelCopyBufferInfo {
290                buffer: &id_staging,
291                layout: wgpu::TexelCopyBufferLayout {
292                    offset: 0,
293                    bytes_per_row: Some(bytes_per_row_aligned),
294                    rows_per_image: Some(1),
295                },
296            },
297            wgpu::Extent3d {
298                width: 1,
299                height: 1,
300                depth_or_array_layers: 1,
301            },
302        );
303        encoder.copy_texture_to_buffer(
304            wgpu::TexelCopyTextureInfo {
305                texture: &pick_depth_texture,
306                mip_level: 0,
307                origin: wgpu::Origin3d { x: px, y: py, z: 0 },
308                aspect: wgpu::TextureAspect::All,
309            },
310            wgpu::TexelCopyBufferInfo {
311                buffer: &depth_staging,
312                layout: wgpu::TexelCopyBufferLayout {
313                    offset: 0,
314                    bytes_per_row: Some(bytes_per_row_aligned),
315                    rows_per_image: Some(1),
316                },
317            },
318            wgpu::Extent3d {
319                width: 1,
320                height: 1,
321                depth_or_array_layers: 1,
322            },
323        );
324
325        queue.submit(std::iter::once(encoder.finish()));
326
327        // --- map and read ---
328        let (tx_id, rx_id) = std::sync::mpsc::channel::<Result<(), wgpu::BufferAsyncError>>();
329        let (tx_dep, rx_dep) = std::sync::mpsc::channel::<Result<(), wgpu::BufferAsyncError>>();
330        id_staging
331            .slice(..)
332            .map_async(wgpu::MapMode::Read, move |r| {
333                let _ = tx_id.send(r);
334            });
335        depth_staging
336            .slice(..)
337            .map_async(wgpu::MapMode::Read, move |r| {
338                let _ = tx_dep.send(r);
339            });
340        device
341            .poll(wgpu::PollType::Wait {
342                submission_index: None,
343                timeout: Some(std::time::Duration::from_secs(5)),
344            })
345            .unwrap();
346        let _ = rx_id.recv().unwrap_or(Err(wgpu::BufferAsyncError));
347        let _ = rx_dep.recv().unwrap_or(Err(wgpu::BufferAsyncError));
348
349        let object_id = {
350            let data = id_staging.slice(..).get_mapped_range();
351            u32::from_le_bytes([data[0], data[1], data[2], data[3]])
352        };
353        id_staging.unmap();
354
355        let depth = {
356            let data = depth_staging.slice(..).get_mapped_range();
357            f32::from_le_bytes([data[0], data[1], data[2], data[3]])
358        };
359        depth_staging.unmap();
360
361        // --- decode sentinel ---
362        // 0 = miss (clear color); anything else is (scene_items_index + 1).
363        if object_id == 0 {
364            return None;
365        }
366
367        Some(crate::interaction::picking::GpuPickHit {
368            object_id: (object_id - 1) as u64,
369            depth,
370        })
371    }
372}