molgfx_render/engine/
picking.rs

1//! Constant-time picking from the exact integer gbuffer silhouette.
2
3use super::Engine;
4use crate::error::RenderError;
5use crate::passes::{
6    ENTITY_RESOURCE, SEGMENT_LABEL_RESOURCE, SEGMENT_VOLUME_RESOURCE, STRUCTURE_RESOURCE,
7};
8use molgfx_core::{
9    AtomSelection, EntityKind, GlobalPickIdentity, GpuPickToken, PickPageTicket, VolumeSegmentRef,
10};
11use molgfx_gpu::{BufferDesc, BufferUsage, CommandEncoder as _, Device, Queue as _, Readback as _};
12
13const READBACK_BYTES: u32 = 256;
14const PICK_FIELDS: u64 = 4;
15
16/// Bytes a pick readback resolves: one field per identity domain.
17pub const PICK_READBACK_BYTES: u64 = READBACK_BYTES as u64 * PICK_FIELDS;
18
19/// A resolved visible entity and its convenient single-atom selection.
20#[derive(Clone, Debug)]
21pub struct Pick {
22    /// Exact scene entity written by the visible fragment.
23    pub entity: PickEntity,
24    /// A one-atom selection for atom-backed entities; empty otherwise.
25    pub selection: AtomSelection,
26}
27
28/// The two identity domains that can be visible in a frame.
29#[derive(Clone, Copy, PartialEq, Eq, Debug)]
30pub enum PickEntity {
31    /// A molecular entity from the opaque or transparent structural path.
32    Structure(GlobalPickIdentity),
33    /// A caller-supplied categorical volume label.
34    VolumeSegment(VolumeSegmentRef),
35}
36
37#[derive(Debug)]
38pub(crate) struct Picker<D: Device> {
39    readback: D::Buffer,
40    submission: Box<[Option<PickPageTicket>]>,
41}
42
43impl<D: Device> Picker<D> {
44    pub(crate) fn new(device: &D, page_capacity: u32) -> Result<Self, RenderError> {
45        let capacity = usize::try_from(page_capacity)
46            .map_err(|_| molgfx_core::PickingError::CapacityTooLarge)?;
47        let mut submission = Vec::new();
48        submission
49            .try_reserve_exact(capacity)
50            .map_err(|_| molgfx_core::PickingError::AllocationFailed)?;
51        submission.resize(capacity, None);
52        Ok(Self {
53            readback: readback(device, "packed pick readback")?,
54            submission: submission.into_boxed_slice(),
55        })
56    }
57}
58
59impl<D: Device> Engine<D> {
60    /// Asynchronously resolves the exact visible entity at one
61    /// top-left-origin pixel. This is the portable browser entry point.
62    ///
63    /// # Errors
64    ///
65    /// Returns a typed device error when readback fails. A pixel outside the
66    /// target or over the background resolves to `Ok(None)`.
67    pub async fn pick_async(&mut self, x: u32, y: u32) -> Result<Option<Pick>, RenderError> {
68        let Some(readback) = self.begin_pick(x, y)? else {
69            return Ok(None);
70        };
71        let packed = readback
72            .resolve(0, u64::from(READBACK_BYTES) * PICK_FIELDS)
73            .await?;
74        self.finish_pick(&packed)
75    }
76
77    /// Records one pick and returns a detached readback handle.
78    ///
79    /// The handle borrows neither the engine nor its device, so on the
80    /// browser's single JavaScript thread a frame may render between this
81    /// call and [`Self::finish_pick`]. `None` means the pixel is outside the
82    /// target or nothing is drawable.
83    ///
84    /// # Errors
85    ///
86    /// Returns a typed error when the pick copies cannot be recorded.
87    pub fn begin_pick(&mut self, x: u32, y: u32) -> Result<Option<D::Readback>, RenderError> {
88        if !self.record_pick(x, y)? {
89            return Ok(None);
90        }
91        Ok(Some(
92            self.queue.readback(&self.device, &self.picker.readback),
93        ))
94    }
95
96    /// Resolves the packed bytes a [`Self::begin_pick`] produced.
97    ///
98    /// # Errors
99    ///
100    /// Returns a typed error when the packed identity cannot be resolved.
101    pub fn finish_pick(&self, packed: &[u8]) -> Result<Option<Pick>, RenderError> {
102        self.resolve_pick(packed)
103    }
104
105    /// Resolves the exact visible entity at one top-left-origin pixel in
106    /// constant time with respect to scene size.
107    ///
108    /// # Errors
109    ///
110    /// Returns a typed device error when readback fails. A pixel outside the
111    /// target or over the background resolves to `Ok(None)`.
112    #[cfg(not(target_arch = "wasm32"))]
113    pub fn pick(&mut self, x: u32, y: u32) -> Result<Option<Pick>, RenderError> {
114        if !self.record_pick(x, y)? {
115            return Ok(None);
116        }
117        let packed = self.queue.read_buffer_blocking(
118            &self.device,
119            &self.picker.readback,
120            0,
121            PICK_READBACK_BYTES,
122        )?;
123        self.finish_pick(&packed)
124    }
125
126    fn record_pick(&mut self, x: u32, y: u32) -> Result<bool, RenderError> {
127        if x >= self.width || y >= self.height {
128            return Ok(false);
129        }
130        let Some(pool) = &self.pool else {
131            return Ok(false);
132        };
133        let (
134            Some(entity_texture),
135            Some(structure_texture),
136            Some(segment_volume_texture),
137            Some(segment_label_texture),
138        ) = (
139            pool.texture(ENTITY_RESOURCE),
140            pool.texture(STRUCTURE_RESOURCE),
141            pool.texture(SEGMENT_VOLUME_RESOURCE),
142            pool.texture(SEGMENT_LABEL_RESOURCE),
143        )
144        else {
145            return Ok(false);
146        };
147        self.scene_gpu
148            .capture_pick_submission(&mut self.picker.submission)?;
149        let mut encoder = self.device.create_command_encoder();
150        encoder.copy_texture_to_buffer(
151            entity_texture,
152            (x, y),
153            (1, 1),
154            READBACK_BYTES,
155            0,
156            &self.picker.readback,
157        );
158        encoder.copy_texture_to_buffer(
159            structure_texture,
160            (x, y),
161            (1, 1),
162            READBACK_BYTES,
163            u64::from(READBACK_BYTES),
164            &self.picker.readback,
165        );
166        encoder.copy_texture_to_buffer(
167            segment_volume_texture,
168            (x, y),
169            (1, 1),
170            READBACK_BYTES,
171            u64::from(READBACK_BYTES) * 2,
172            &self.picker.readback,
173        );
174        encoder.copy_texture_to_buffer(
175            segment_label_texture,
176            (x, y),
177            (1, 1),
178            READBACK_BYTES,
179            u64::from(READBACK_BYTES) * 3,
180            &self.picker.readback,
181        );
182        self.queue.submit(encoder);
183        Ok(true)
184    }
185
186    fn resolve_pick(&self, packed: &[u8]) -> Result<Option<Pick>, RenderError> {
187        let local_row = pick_field(packed, 0);
188        let resident_page = pick_field(packed, 1);
189        let segment_volume = pick_field(packed, 2);
190        let segment_label = pick_field(packed, 3);
191        if let (Some(source_id), Some(label)) = (read_u32(segment_volume), read_u32(segment_label))
192            && source_id != u32::MAX
193            && let Some(segment) = self.scene_gpu.resolve_segment(source_id, label)
194        {
195            return Ok(Some(Pick {
196                entity: PickEntity::VolumeSegment(segment),
197                selection: AtomSelection::Empty,
198            }));
199        }
200        let (Some(local_row), Some(resident_page)) = (read_u32(local_row), read_u32(resident_page))
201        else {
202            return Ok(None);
203        };
204        let token = GpuPickToken::new(resident_page, local_row);
205        if token == GpuPickToken::NONE {
206            return Ok(None);
207        }
208        let identity = self
209            .scene_gpu
210            .resolve_global_pick(token, &self.picker.submission)?;
211        let selection = match (identity.kind(), u32::try_from(identity.row().get())) {
212            (EntityKind::Atom, Ok(row)) => match row.checked_add(1) {
213                Some(end) => AtomSelection::Range(row..end),
214                None => AtomSelection::Empty,
215            },
216            _ => AtomSelection::Empty,
217        };
218        Ok(Some(Pick {
219            entity: PickEntity::Structure(identity),
220            selection,
221        }))
222    }
223
224    #[cfg(test)]
225    pub(crate) fn capture_pick_submission_for_test(&mut self) -> Result<(), RenderError> {
226        self.scene_gpu
227            .capture_pick_submission(&mut self.picker.submission)
228    }
229
230    #[cfg(test)]
231    pub(crate) fn resolve_pick_token_for_test(
232        &self,
233        token: GpuPickToken,
234    ) -> Result<GlobalPickIdentity, RenderError> {
235        self.scene_gpu
236            .resolve_global_pick(token, &self.picker.submission)
237    }
238}
239
240fn readback<D: Device>(device: &D, label: &'static str) -> Result<D::Buffer, RenderError> {
241    Ok(device.create_buffer(&BufferDesc {
242        label,
243        size: u64::from(READBACK_BYTES) * PICK_FIELDS,
244        usage: BufferUsage::COPY_DST.union(BufferUsage::MAP_READ),
245    })?)
246}
247
248fn pick_field(bytes: &[u8], index: usize) -> &[u8] {
249    let width = READBACK_BYTES as usize;
250    let start = index.saturating_mul(width);
251    let end = start.saturating_add(width).min(bytes.len());
252    match bytes.get(start..end) {
253        Some(field) => field,
254        None => &[],
255    }
256}
257
258fn read_u32(bytes: &[u8]) -> Option<u32> {
259    let slice = bytes.get(..std::mem::size_of::<u32>())?;
260    let mut array = [0; std::mem::size_of::<u32>()];
261    array.copy_from_slice(slice);
262    Some(u32::from_le_bytes(array))
263}
molgfx_render/engine/picking.rs

molgfx_render/engine/
picking.rs
