Skip to main content

flow_ngin/
pick.rs

1//! Object picking and selection.
2//!
3//! This module implements GPU-based object picking: rendering scene objects with
4//! unique IDs to an offscreen texture, then reading the pixel under the mouse cursor
5//! to determine which object was clicked. Supports picking for both 3D (instanced)
6//! and 2D (GUI/flat) objects.
7//!
8//! The picking pipeline works as follows:
9//! 1. Render all objects to an offscreen texture using unique IDs as RGBA values for the fragment shader
10//! 2. Read the pixel at the mouse cursor position (scaled according to platform limitations on texture sizes)
11//! 3. Map the pick ID back to the flow that owns the object (determined by the render tree)
12//! 4. Return the selected object ID and owning flows
13//!
14//! Especially step 4 makes sure that only those flows are invoked that were responsible for selected object.
15
16use std::{
17    collections::{HashMap, HashSet},
18    iter,
19};
20
21use crate::{
22    context::{Context, MouseState},
23    data_structures::model::DrawModel,
24    flow::GraphicsFlow,
25    render::{Flat, Geometry, Instanced},
26    resources::pick::{load_pick_model, load_pick_texture},
27};
28
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
30pub struct PickId(pub u32);
31
32#[cfg(target_arch = "wasm32")]
33use crate::flow::FlowEvent;
34
35/// Render all flows to pick texture and determine which object was clicked.
36///
37/// # Arguments
38///
39/// * `async_runtime` using the tokio runtime for async resource loading if not on WASM
40/// * `flows` represent all active graphics flows with their renderable objects
41/// * `ctx` is the rendering context
42/// * `mouse_state` is required for getting the mouse coordinates at the time of picking
43/// * `proxy` WASM futures can only resolve using the winit event loop proxy by sending events
44///
45/// # Returns
46///
47/// `Some((pick_id, flow_ids))` if an object was picked, or `None` picking is done via the event loop.
48pub(crate) fn draw_to_pick_buffer<State, Event: Send>(
49    #[cfg(not(target_arch = "wasm32"))] async_runtime: &tokio::runtime::Runtime,
50    flows: &mut Vec<Box<dyn GraphicsFlow<State, Event>>>,
51    ctx: &Context,
52    mouse_state: &MouseState,
53    #[cfg(target_arch = "wasm32")] proxy: winit::event_loop::EventLoopProxy<
54        crate::flow::FlowEvent<State, Event>,
55    >,
56) -> Option<(u32, HashSet<usize>)> {
57    // Prepare data for picking:
58    let u32_size = std::mem::size_of::<u32>() as u32;
59    // The img lib requires divisibility of 256...
60    let width = ctx.config.width;
61    let height = ctx.config.height;
62    let width_offset = 256 - (width % 256);
63    let height_offset = 256 - (height % 256);
64    // TODO: if on wasm max at 2048 and keep ratio
65    let width_factor = (f64::from(width) + f64::from(width_offset)) / f64::from(width);
66    let height_factor = (f64::from(height) + f64::from(height_offset)) / f64::from(height);
67    let width = width + width_offset;
68    let height = height + height_offset;
69
70    let extent3d = wgpu::Extent3d {
71        width: width,
72        height: height,
73        depth_or_array_layers: 1,
74    };
75
76    let pick_texture = &ctx.device.create_texture(&wgpu::TextureDescriptor {
77        label: Some("Pick texture"),
78        size: extent3d,
79        mip_level_count: 1,
80        sample_count: 1,
81        dimension: wgpu::TextureDimension::D2,
82        format: wgpu::TextureFormat::R32Uint,
83        usage: wgpu::TextureUsages::COPY_SRC | wgpu::TextureUsages::RENDER_ATTACHMENT,
84        view_formats: &[],
85    });
86
87    let pick_depth_texture = &ctx.device.create_texture(&wgpu::TextureDescriptor {
88        label: Some("Pick depth texture"),
89        size: extent3d,
90        mip_level_count: 1,
91        sample_count: 1,
92        dimension: wgpu::TextureDimension::D2,
93        format: wgpu::TextureFormat::Depth24Plus,
94        usage: wgpu::TextureUsages::COPY_SRC | wgpu::TextureUsages::RENDER_ATTACHMENT,
95        view_formats: &[],
96    });
97
98    let mut encoder = ctx
99        .device
100        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
101            label: Some("Pick Encoder"),
102        });
103    let mut translation: HashMap<PickId, HashSet<usize>> = HashMap::new();
104
105    {
106        let mut render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
107            label: Some("Render Pass"),
108            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
109                view: &pick_texture.create_view(&wgpu::TextureViewDescriptor {
110                    label: Some("Render texture"),
111                    format: Some(wgpu::TextureFormat::R32Uint),
112                    dimension: Some(wgpu::TextureViewDimension::D2),
113                    usage: None,
114                    aspect: wgpu::TextureAspect::All,
115                    base_mip_level: 0,
116                    mip_level_count: None,
117                    base_array_layer: 0,
118                    array_layer_count: None,
119                }),
120                resolve_target: None,
121                ops: wgpu::Operations {
122                    load: wgpu::LoadOp::Clear(wgpu::Color::TRANSPARENT),
123                    store: wgpu::StoreOp::Store,
124                },
125                depth_slice: None,
126            })],
127            depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
128                view: &pick_depth_texture.create_view(&wgpu::TextureViewDescriptor {
129                    label: Some("Stencil texture"),
130                    format: Some(wgpu::TextureFormat::Depth24Plus),
131                    dimension: Some(wgpu::TextureViewDimension::D2),
132                    usage: None,
133                    aspect: wgpu::TextureAspect::All,
134                    base_mip_level: 0,
135                    mip_level_count: None,
136                    base_array_layer: 0,
137                    array_layer_count: None,
138                }),
139                depth_ops: Some(wgpu::Operations {
140                    load: wgpu::LoadOp::Clear(1.0),
141                    store: wgpu::StoreOp::Store,
142                }),
143                stencil_ops: None,
144            }),
145            ..Default::default()
146        });
147
148        let mut basics: Vec<Instanced> = Vec::new();
149        let mut flats: Vec<Flat> = Vec::new();
150        let mut geoms: Vec<Geometry> = Vec::new();
151        /*
152           We support graphics flow that handle pick IDs internally. Thus, we store the
153           correspondance of the flow index and the model picked so that each flow only
154           gets invoked if one of the IDs it manages was picked.
155
156           Example:
157           flow1 at index 0 owns the pick IDs [1, 2, 3, 4, 5]
158           flow2 at index 1 owns the pick IDs [5, 6, 7, 8, 9]
159
160           Warning: Overlapping ID responsibility may not be the best design choice.
161
162           On pick result 2 we invoke flow1.on_pick(2).
163           On pick result 5 we invoke flow1.on_pick(5) followed by flow2.on_pick(5).
164        */
165        flows.iter_mut().enumerate().for_each(|(idx, flow)| {
166            let render = flow.on_render();
167            render.map_ids(idx, &mut translation);
168            render.set_pick_pipelines(&ctx, &mut render_pass, &mut basics, &mut flats, &mut geoms);
169        });
170
171        render_pass.set_pipeline(&ctx.pipelines.pick);
172        for instanced in basics.iter_mut() {
173            if instanced.amount == 0 || instanced.instance.size() == 0 {
174                log::debug!("Cannot pick empty render.");
175                continue;
176            }
177            let pick_model =
178                load_pick_model(&ctx.device, instanced.id, instanced.model.meshes.clone()).unwrap();
179            render_pass.set_vertex_buffer(1, instanced.instance.slice(..));
180            let amount: Result<u32, _> = instanced.amount.try_into();
181            match amount {
182                Err(e) => log::error!(
183                    "Failed to render flat object with id {:?}. Maximum amount of supported instances is {}. Error: {}",
184                    instanced.id,
185                    u32::MAX,
186                    e
187                ),
188                Ok(amount) => render_pass.draw_model_instanced(
189                    &pick_model,
190                    0..amount,
191                    &ctx.camera.bind_group,
192                    &ctx.light.bind_group,
193                ),
194            }
195        }
196
197        render_pass.set_pipeline(&ctx.pipelines.flat_pick);
198        render_pass.set_bind_group(1, &ctx.screen_size.bind_group, &[]);
199        for flat in flats {
200            let pick_group = load_pick_texture(flat.id, &ctx.device);
201            render_pass.set_bind_group(0, &pick_group, &[]);
202            render_pass.set_vertex_buffer(0, flat.vertex.slice(..));
203            render_pass.set_index_buffer(flat.index.slice(..), wgpu::IndexFormat::Uint16);
204            let amount: Result<u32, _> = flat.amount.try_into();
205            match amount {
206                Err(e) => log::error!(
207                    "Failed to render flat object with id {:?}. Maximum amount of supported instances is {}. Error: {}",
208                    flat.id,
209                    u32::MAX,
210                    e
211                ),
212                Ok(amount) => render_pass.draw_indexed(0..amount, 0, 0..1),
213            }
214        }
215    }
216
217    let output_buffer_size = (u32_size * (width) * (height)) as wgpu::BufferAddress;
218    let output_buffer_desc = wgpu::BufferDescriptor {
219        size: output_buffer_size,
220        usage: wgpu::BufferUsages::COPY_DST
221                    // this tells wpgu that we want to read this buffer from the cpu
222                    | wgpu::BufferUsages::MAP_READ,
223        label: None,
224        mapped_at_creation: false,
225    };
226    let output_buffer = ctx.device.create_buffer(&output_buffer_desc);
227
228    encoder.copy_texture_to_buffer(
229        wgpu::TexelCopyTextureInfo {
230            aspect: wgpu::TextureAspect::All,
231            texture: &pick_texture,
232            mip_level: 0,
233            origin: wgpu::Origin3d::ZERO,
234        },
235        wgpu::TexelCopyBufferInfo {
236            buffer: &output_buffer,
237            layout: wgpu::TexelCopyBufferLayout {
238                offset: 0,
239                bytes_per_row: Some(u32_size * (width)),
240                rows_per_image: Some(height),
241            },
242        },
243        extent3d,
244    );
245
246    ctx.queue.submit(iter::once(encoder.finish()));
247    let device = ctx.device.clone();
248    let mouse_coords = mouse_state.coords.clone();
249    #[cfg(target_arch = "wasm32")]
250    wasm_bindgen_futures::spawn_local(async move {
251        let buffer_slice = output_buffer.slice(..);
252        let future_id = read_texture_buffer(
253            buffer_slice,
254            &device,
255            width_factor,
256            height_factor,
257            width,
258            height,
259            mouse_coords,
260        );
261        let id = future_id.await;
262        if let Some(flow_ids) = translation.get(&id) {
263            assert!(
264                proxy
265                    .send_event(FlowEvent::Id((id, flow_ids.clone())))
266                    .is_ok()
267            );
268            output_buffer.unmap();
269        };
270    });
271    #[cfg(target_arch = "wasm32")]
272    return None;
273    #[cfg(not(target_arch = "wasm32"))]
274    {
275        let buffer_slice = output_buffer.slice(..);
276        let future_id = read_texture_buffer(
277            buffer_slice,
278            &device,
279            width_factor,
280            height_factor,
281            width,
282            height,
283            mouse_coords,
284        );
285        // Depending on the average timing this hould not block but rather always send an event
286        let id = async_runtime.block_on(future_id);
287        return translation.get(&PickId(id)).map(|flow_ids| (id, flow_ids.clone()));
288    }
289}
290
291async fn read_texture_buffer(
292    buffer_slice: wgpu::BufferSlice<'_>,
293    device: &wgpu::Device,
294    width_factor: f64,
295    height_factor: f64,
296    width: u32,
297    _height: u32,
298    mouse_coords: winit::dpi::PhysicalPosition<f64>,
299) -> u32 {
300    // NOTE: We have to create the mapping THEN device.poll() before await
301    // the future. Otherwise the application will freeze.
302    let (tx, rx) = futures_intrusive::channel::shared::oneshot_channel();
303    buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
304        tx.send(result).unwrap();
305    });
306    #[cfg(target_arch = "wasm32")]
307    device.poll(wgpu::PollType::Poll).unwrap();
308    #[cfg(not(target_arch = "wasm32"))]
309    device
310        .poll(wgpu::PollType::Wait {
311            submission_index: None,
312            timeout: None,
313        })
314        .unwrap();
315    rx.receive().await.unwrap().unwrap();
316
317    let data = buffer_slice.get_mapped_range();
318    // [(0, 0, 0, 0), (0`, 255, 0, 255), (0, 0, 0, 0),
319    // (0, 0, 0, 0), (0, 255, 0, 255), (0, 0, 0, 0)]
320    let x = mouse_coords.x * width_factor;
321    let y = mouse_coords.y * height_factor;
322    let bytes_per_pixel = 4;
323    let pick_index = (y as usize * width as usize + x as usize) * bytes_per_pixel;
324    // TODO: bounds check.
325    let r = data[pick_index];
326    let g = data[pick_index + 1];
327    let b = data[pick_index + 2];
328    let a = data[pick_index + 3];
329
330    let rgba_u32 = u32::from(r) | u32::from(g) << 8 | u32::from(b) << 16 | u32::from(a) << 24;
331
332    // This is great for debugging. I'll keep it as I need it often.
333    /*use image::{ImageBuffer, Rgba};
334    let buffer = ImageBuffer::<Rgba<u8>, _>::from_raw(width, height, data).unwrap();
335    buffer.save("image.png").unwrap();*/
336
337    log::info!("Selected obj with id {}", rgba_u32);
338    rgba_u32
339}