molframe_core/selection/query/
representation.rs

1//! Selection shapes, construction, membership and iteration.
2
3use crate::column::{BitVec, Ones};
4use smallvec::SmallVec;
5use std::ops::Range;
6
7/// A set of atoms, addressed by position in the structure's flat atom order.
8///
9/// # Examples
10///
11/// ```
12/// use molframe_core::AtomSelection;
13///
14/// let chain = AtomSelection::range(0..100);
15/// let metals = AtomSelection::from_sorted(vec![5, 50, 500]);
16/// assert_eq!(chain.intersect(&metals).len(), 2);
17/// ```
18#[derive(Clone, PartialEq, Eq, Debug, Default)]
19pub enum AtomSelection {
20    /// No atoms.
21    #[default]
22    Empty,
23    /// Every atom of a structure with this many.
24    All(u32),
25    /// One contiguous run.
26    Range(Range<u32>),
27    /// Several contiguous runs, ascending and disjoint.
28    Ranges(SmallVec<[Range<u32>; 4]>),
29    /// Scattered positions, ascending and unique.
30    Sparse(Vec<u32>),
31    /// A bit per position, for a set dense enough to make that cheaper.
32    Dense(BitVec),
33}
34
35const DENSE_FRACTION: u32 = 16;
36
37impl AtomSelection {
38    /// A selection covering one contiguous run.
39    #[must_use]
40    pub fn range(atoms: Range<u32>) -> Self {
41        if atoms.start >= atoms.end {
42            Self::Empty
43        } else {
44            Self::Range(atoms)
45        }
46    }
47
48    /// A selection of runs; empty runs are dropped, runs are put in order, and
49    /// runs that touch or overlap are merged.
50    ///
51    /// Residue- and chain-level predicates decide once per residue, so their
52    /// result is a handful of runs rather than one position per atom.
53    #[must_use]
54    pub fn from_runs(runs: impl IntoIterator<Item = Range<u32>>) -> Self {
55        let mut ordered: Vec<Range<u32>> =
56            runs.into_iter().filter(|run| run.start < run.end).collect();
57        // Already ascending in the common case, which the stable sort handles in one pass.
58        ordered.sort_by_key(|run| run.start);
59        let mut merged = SmallVec::<[Range<u32>; 4]>::new();
60        for run in ordered {
61            match merged.last_mut() {
62                Some(last) if run.start <= last.end => last.end = last.end.max(run.end),
63                _ => merged.push(run),
64            }
65        }
66        match merged.len() {
67            0 => Self::Empty,
68            1 => merged.pop().map_or(Self::Empty, Self::Range),
69            _ => Self::Ranges(merged),
70        }
71    }
72
73    /// Chooses the cheapest shape for ascending unique positions.
74    ///
75    /// # Panics
76    /// Panics if `positions` is not ascending and unique, or if its storage
77    /// exceeds the platform's addressable range. Collect arbitrary positions
78    /// into `AtomSelection` to sort and deduplicate them first.
79    #[must_use]
80    pub fn from_sorted(positions: Vec<u32>) -> Self {
81        let (Some(first), Some(last)) = (positions.first().copied(), positions.last().copied())
82        else {
83            return Self::Empty;
84        };
85        let span = u64::from(last - first) + 1;
86        if last < u32::MAX && span == positions.len() as u64 {
87            return Self::Range(first..last + 1);
88        }
89        if let Some(runs) = inline_runs(&positions) {
90            let covered: u64 = runs.iter().map(|run| u64::from(run.end - run.start)).sum();
91            if covered == positions.len() as u64 && runs.len() * 2 < positions.len() {
92                return Self::Ranges(runs);
93            }
94        }
95        if let Some(mask_len) = last.checked_add(1)
96            && positions.len() as u64 * u64::from(DENSE_FRACTION) >= u64::from(mask_len)
97        {
98            return Self::Dense(mask_from_positions(mask_len, &positions));
99        }
100        Self::Sparse(positions)
101    }
102
103    /// The number of atoms selected.
104    ///
105    /// # Panics
106    /// Panics if manually constructed storage violates normalized `u32` ranges.
107    #[must_use]
108    pub fn len(&self) -> u64 {
109        selection_len(self)
110    }
111
112    /// Returns true when nothing is selected.
113    #[must_use]
114    pub fn is_empty(&self) -> bool {
115        self.len() == 0
116    }
117
118    /// Returns true when `atom` is selected.
119    #[must_use]
120    pub fn contains(&self, atom: u32) -> bool {
121        match self {
122            Self::Empty => false,
123            Self::All(count) => atom < *count,
124            Self::Range(run) => run.contains(&atom),
125            Self::Ranges(runs) => {
126                let candidate = runs.partition_point(|run| run.end <= atom);
127                runs.get(candidate).is_some_and(|run| run.contains(&atom))
128            }
129            Self::Sparse(positions) => positions.binary_search(&atom).is_ok(),
130            Self::Dense(mask) => mask.test(atom),
131        }
132    }
133
134    /// One past the largest position this selection could name.
135    ///
136    /// Derived from the shape rather than the positions, so a whole-structure
137    /// selection answers from its own count and a run from its end. A caller
138    /// that has to prove every position is inside a structure compares this once
139    /// instead of asking about each atom — which matters because the cheapest
140    /// selections, `All` and `Empty`, are exactly the ones a per-atom walk would
141    /// read as a full traversal of the structure.
142    ///
143    /// The bound is exact: no position in the selection reaches it, and every
144    /// shape reports the smallest bound that says so.
145    #[must_use]
146    pub fn position_bound(&self) -> u32 {
147        match self {
148            Self::Empty => 0,
149            Self::All(count) => *count,
150            Self::Range(run) => run.end,
151            Self::Ranges(runs) => match runs.last() {
152                Some(run) => run.end,
153                None => 0,
154            },
155            Self::Sparse(positions) => match positions.last() {
156                Some(last) => last.saturating_add(1),
157                None => 0,
158            },
159            Self::Dense(mask) => mask.len(),
160        }
161    }
162
163    /// The selected positions, ascending.
164    ///
165    /// A concrete iterator rather than a boxed one: the boxed form cost a heap
166    /// allocation per call and a virtual call per position, which prevented the
167    /// consuming loop from inlining anything. Selection walks are the innermost
168    /// loop of every spatial and analysis pass.
169    #[must_use]
170    pub fn iter(&self) -> SelectionIter<'_> {
171        match self {
172            Self::Empty => SelectionIter::Run(0..0),
173            Self::All(count) => SelectionIter::Run(0..*count),
174            Self::Range(run) => SelectionIter::Run(run.start..run.end),
175            Self::Ranges(runs) => SelectionIter::Runs {
176                runs,
177                next: 0,
178                current: 0..0,
179            },
180            Self::Sparse(positions) => SelectionIter::Sparse(positions.iter()),
181            Self::Dense(mask) => SelectionIter::Dense(mask.ones()),
182        }
183    }
184}
185
186/// The positions of an [`AtomSelection`], ascending.
187#[derive(Clone, Debug)]
188pub enum SelectionIter<'a> {
189    /// One contiguous run, which also covers the empty and whole-structure
190    /// selections.
191    Run(Range<u32>),
192    /// Several runs, walked one after another.
193    Runs {
194        /// The remaining runs, including the one being walked.
195        runs: &'a [Range<u32>],
196        /// The index of the next run to start.
197        next: usize,
198        /// The unfinished part of the current run.
199        current: Range<u32>,
200    },
201    /// Scattered positions read straight from their slice.
202    Sparse(std::slice::Iter<'a, u32>),
203    /// Set bits of a dense mask.
204    Dense(Ones<'a>),
205}
206
207impl Iterator for SelectionIter<'_> {
208    type Item = u32;
209
210    fn next(&mut self) -> Option<u32> {
211        match self {
212            Self::Run(run) => run.next(),
213            Self::Runs {
214                runs,
215                next,
216                current,
217            } => loop {
218                if let Some(position) = current.next() {
219                    return Some(position);
220                }
221                let run = runs.get(*next)?;
222                *next = next.checked_add(1)?;
223                *current = run.start..run.end;
224            },
225            Self::Sparse(positions) => positions.next().copied(),
226            Self::Dense(ones) => ones.next(),
227        }
228    }
229
230    fn size_hint(&self) -> (usize, Option<usize>) {
231        match self {
232            Self::Run(run) => run.size_hint(),
233            Self::Runs {
234                runs,
235                next,
236                current,
237            } => {
238                let remaining = match runs.get(*next..) {
239                    Some(tail) => tail.iter().map(ExactSizeIterator::len).sum::<usize>(),
240                    None => 0,
241                };
242                let total = current.len() + remaining;
243                (total, Some(total))
244            }
245            Self::Sparse(positions) => positions.size_hint(),
246            Self::Dense(ones) => ones.size_hint(),
247        }
248    }
249}
250
251impl ExactSizeIterator for SelectionIter<'_> {}
252
253impl std::iter::FusedIterator for SelectionIter<'_> {}
254
255impl FromIterator<u32> for AtomSelection {
256    fn from_iter<T: IntoIterator<Item = u32>>(iter: T) -> Self {
257        let mut positions: Vec<u32> = iter.into_iter().collect();
258        positions.sort_unstable();
259        positions.dedup();
260        Self::from_sorted(positions)
261    }
262}
263
264impl<'a> IntoIterator for &'a AtomSelection {
265    type Item = u32;
266    type IntoIter = SelectionIter<'a>;
267
268    fn into_iter(self) -> Self::IntoIter {
269        self.iter()
270    }
271}
272
273fn inline_runs(positions: &[u32]) -> Option<SmallVec<[Range<u32>; 4]>> {
274    let mut runs = SmallVec::<[Range<u32>; 4]>::new();
275    for position in positions.iter().copied() {
276        let end = position.checked_add(1)?;
277        if let Some(run) = runs.last_mut()
278            && run.end == position
279        {
280            run.end = end;
281            continue;
282        }
283        if runs.len() == runs.inline_size() {
284            return None;
285        }
286        runs.push(position..end);
287    }
288    Some(runs)
289}
290
291fn mask_from_positions(len: u32, positions: &[u32]) -> BitVec {
292    let mut mask = BitVec::repeat(false, len);
293    for position in positions {
294        mask.set(*position, true);
295    }
296    mask
297}
298
299fn selection_len(selection: &AtomSelection) -> u64 {
300    match selection {
301        AtomSelection::Empty => 0,
302        AtomSelection::All(count) => u64::from(*count),
303        AtomSelection::Range(run) => u64::from(run.end - run.start),
304        AtomSelection::Ranges(runs) => runs.iter().map(|run| u64::from(run.end - run.start)).sum(),
305        AtomSelection::Sparse(positions) => positions.len() as u64,
306        AtomSelection::Dense(mask) => u64::from(mask.count_ones()),
307    }
308}
molframe_core/selection/query/representation.rs

molframe_core/selection/query/
representation.rs
