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

molgfx_render/engine/
picking.rs
