ferritin_core/data/
ordered_set.rs

1//! Ordered set of indices for atom/residue selection membership.
2//!
3//! An [`OrderedSet`] represents either a contiguous interval (for dense ranges)
4//! or a sorted array of unique indices (for sparse selections).
5
6use std::sync::Arc;
7
8/// Allocation-free iterator over an [`OrderedSet`].
9pub enum OrderedSetIter<'a> {
10    Interval(std::ops::Range<u32>),
11    Sorted(std::iter::Copied<std::slice::Iter<'a, u32>>),
12}
13
14impl<'a> Iterator for OrderedSetIter<'a> {
15    type Item = u32;
16    fn next(&mut self) -> Option<u32> {
17        match self {
18            OrderedSetIter::Interval(r) => r.next(),
19            OrderedSetIter::Sorted(it) => it.next(),
20        }
21    }
22}
23
24/// Ordered set of indices — either a contiguous interval or sorted array.
25///
26/// The `Interval` variant provides O(1) membership tests for contiguous ranges.
27/// The `Sorted` variant provides O(log n) membership tests for arbitrary selections.
28#[derive(Clone, Debug, PartialEq)]
29pub enum OrderedSet {
30    /// Contiguous range [start, end) — O(1) membership test
31    Interval { start: u32, end: u32 },
32    /// Sorted unique indices — O(log n) membership test
33    Sorted(Arc<[u32]>),
34}
35
36impl OrderedSet {
37    /// Create a contiguous interval [start, end).
38    ///
39    /// # Panics
40    /// Panics if `start > end`.
41    pub fn interval(start: u32, end: u32) -> Self {
42        assert!(start <= end, "interval start must be <= end");
43        OrderedSet::Interval { start, end }
44    }
45
46    /// Create a `Sorted` set from a `Vec<u32>`.
47    ///
48    /// The input must already be sorted and contain no duplicates.
49    /// This function panics if either condition is violated, keeping the
50    /// invariant that `Sorted` always holds valid data.
51    ///
52    /// If you need to accept arbitrary input, sort and dedup before calling:
53    /// ```
54    /// # use ferritin_core::data::ordered_set::OrderedSet;
55    /// let mut v = vec![3u32, 1, 2, 1];
56    /// v.sort_unstable();
57    /// v.dedup();
58    /// let set = OrderedSet::from_sorted(v);
59    /// ```
60    pub fn from_sorted(indices: Vec<u32>) -> Self {
61        // Validate: must be strictly increasing (sorted + no duplicates)
62        for window in indices.windows(2) {
63            assert!(
64                window[0] < window[1],
65                "from_sorted requires strictly sorted (no duplicates) input; \
66                 found {} >= {}",
67                window[0],
68                window[1]
69            );
70        }
71        OrderedSet::Sorted(indices.into())
72    }
73
74    /// Returns the number of elements in the set.
75    pub fn len(&self) -> usize {
76        match self {
77            OrderedSet::Interval { start, end } => (end - start) as usize,
78            OrderedSet::Sorted(v) => v.len(),
79        }
80    }
81
82    /// Returns `true` if the set contains no elements.
83    pub fn is_empty(&self) -> bool {
84        self.len() == 0
85    }
86
87    /// Returns `true` if `idx` is a member of this set.
88    pub fn contains(&self, idx: u32) -> bool {
89        match self {
90            OrderedSet::Interval { start, end } => idx >= *start && idx < *end,
91            OrderedSet::Sorted(v) => v.binary_search(&idx).is_ok(),
92        }
93    }
94
95    /// Iterate over the elements of this set in ascending order.
96    pub fn iter(&self) -> OrderedSetIter<'_> {
97        match self {
98            OrderedSet::Interval { start, end } => OrderedSetIter::Interval(*start..*end),
99            OrderedSet::Sorted(v) => OrderedSetIter::Sorted(v.iter().copied()),
100        }
101    }
102
103    /// Compute the union of this set and `other`.
104    ///
105    /// Uses a merge-style algorithm on two sorted iterators — O(n + m).
106    /// Returns an `OrderedSet::Sorted`.
107    pub fn union(&self, other: &Self) -> Self {
108        let mut result = Vec::with_capacity(self.len() + other.len());
109        let mut left = self.iter().peekable();
110        let mut right = other.iter().peekable();
111
112        loop {
113            match (left.peek(), right.peek()) {
114                (None, None) => break,
115                (Some(_), None) => {
116                    result.extend(left);
117                    break;
118                }
119                (None, Some(_)) => {
120                    result.extend(right);
121                    break;
122                }
123                (Some(&l), Some(&r)) => {
124                    if l < r {
125                        result.push(l);
126                        left.next();
127                    } else if r < l {
128                        result.push(r);
129                        right.next();
130                    } else {
131                        // equal: include once, advance both
132                        result.push(l);
133                        left.next();
134                        right.next();
135                    }
136                }
137            }
138        }
139
140        OrderedSet::Sorted(result.into())
141    }
142
143    /// Compute the intersection of this set and `other`.
144    ///
145    /// Uses a merge-style algorithm — O(n + m).
146    /// Returns an `OrderedSet::Sorted`.
147    pub fn intersection(&self, other: &Self) -> Self {
148        let mut result = Vec::new();
149        let mut left = self.iter().peekable();
150        let mut right = other.iter().peekable();
151
152        loop {
153            match (left.peek(), right.peek()) {
154                (None, _) | (_, None) => break,
155                (Some(&l), Some(&r)) => {
156                    if l == r {
157                        result.push(l);
158                        left.next();
159                        right.next();
160                    } else if l < r {
161                        left.next();
162                    } else {
163                        right.next();
164                    }
165                }
166            }
167        }
168
169        OrderedSet::Sorted(result.into())
170    }
171
172    /// Compute the set difference `self \ other` (elements in self but not other).
173    ///
174    /// Uses a merge-style algorithm — O(n + m).
175    /// Returns an `OrderedSet::Sorted`.
176    pub fn difference(&self, other: &Self) -> Self {
177        let mut result = Vec::new();
178        let mut left = self.iter().peekable();
179        let mut right = other.iter().peekable();
180
181        loop {
182            match left.peek() {
183                None => break,
184                Some(&l) => match right.peek() {
185                    None => {
186                        result.extend(left);
187                        break;
188                    }
189                    Some(&r) => {
190                        if l < r {
191                            result.push(l);
192                            left.next();
193                        } else if l == r {
194                            left.next();
195                            right.next();
196                        } else {
197                            right.next();
198                        }
199                    }
200                },
201            }
202        }
203
204        OrderedSet::Sorted(result.into())
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211
212    #[test]
213    fn test_interval_contains() {
214        let s = OrderedSet::interval(5, 10);
215        assert!(s.contains(5));
216        assert!(s.contains(9));
217        assert!(!s.contains(4));
218        assert!(!s.contains(10));
219    }
220
221    #[test]
222    fn test_interval_len() {
223        let s = OrderedSet::interval(3, 8);
224        assert_eq!(s.len(), 5);
225
226        let empty = OrderedSet::interval(4, 4);
227        assert_eq!(empty.len(), 0);
228        assert!(empty.is_empty());
229    }
230
231    #[test]
232    fn test_sorted_contains() {
233        let s = OrderedSet::from_sorted(vec![1, 3, 5, 7, 9]);
234        assert!(s.contains(1));
235        assert!(s.contains(5));
236        assert!(s.contains(9));
237        assert!(!s.contains(0));
238        assert!(!s.contains(2));
239        assert!(!s.contains(10));
240    }
241
242    #[test]
243    fn test_ordered_set_intersection() {
244        let a = OrderedSet::from_sorted(vec![1, 3, 5, 7, 9]);
245        let b = OrderedSet::from_sorted(vec![2, 3, 5, 8]);
246        let result = a.intersection(&b);
247        assert_eq!(result, OrderedSet::from_sorted(vec![3, 5]));
248    }
249
250    #[test]
251    fn test_ordered_set_union() {
252        let a = OrderedSet::from_sorted(vec![1, 3, 5]);
253        let b = OrderedSet::from_sorted(vec![2, 4, 6]);
254        let result = a.union(&b);
255        assert_eq!(result, OrderedSet::from_sorted(vec![1, 2, 3, 4, 5, 6]));
256    }
257
258    #[test]
259    fn test_ordered_set_difference() {
260        let a = OrderedSet::from_sorted(vec![1, 3, 5, 7, 9]);
261        let b = OrderedSet::from_sorted(vec![3, 5]);
262        let result = a.difference(&b);
263        assert_eq!(result, OrderedSet::from_sorted(vec![1, 7, 9]));
264    }
265
266    #[test]
267    fn test_interval_iter() {
268        let s = OrderedSet::interval(2, 6);
269        let collected: Vec<u32> = s.iter().collect();
270        assert_eq!(collected, vec![2, 3, 4, 5]);
271    }
272
273    #[test]
274    fn test_sorted_iter() {
275        let s = OrderedSet::from_sorted(vec![10, 20, 30]);
276        let collected: Vec<u32> = s.iter().collect();
277        assert_eq!(collected, vec![10, 20, 30]);
278    }
279
280    #[test]
281    fn test_empty_intersection() {
282        let a = OrderedSet::from_sorted(vec![1, 3, 5]);
283        let b = OrderedSet::from_sorted(vec![2, 4, 6]);
284        let result = a.intersection(&b);
285        assert!(result.is_empty());
286    }
287
288    #[test]
289    #[should_panic(expected = "from_sorted requires strictly sorted")]
290    fn test_from_sorted_panics_on_duplicate() {
291        // Documented behaviour: from_sorted PANICS on non-sorted or duplicate input.
292        // Callers that need to accept arbitrary input should sort+dedup first.
293        OrderedSet::from_sorted(vec![1, 2, 2, 3]);
294    }
295
296    #[test]
297    #[should_panic(expected = "from_sorted requires strictly sorted")]
298    fn test_from_sorted_panics_on_unsorted() {
299        OrderedSet::from_sorted(vec![3, 1, 2]);
300    }
301
302    /// Demonstrate the recommended pattern for arbitrary (unsorted/duplicate) input.
303    #[test]
304    fn test_from_sorted_sort_dedup_pattern() {
305        let mut v = vec![3u32, 1, 2, 1, 3];
306        v.sort_unstable();
307        v.dedup();
308        let s = OrderedSet::from_sorted(v);
309        assert_eq!(s.len(), 3);
310        assert!(s.contains(1));
311        assert!(s.contains(2));
312        assert!(s.contains(3));
313    }
314
315    #[test]
316    fn test_union_with_overlap() {
317        let a = OrderedSet::from_sorted(vec![1, 2, 3]);
318        let b = OrderedSet::from_sorted(vec![2, 3, 4]);
319        let result = a.union(&b);
320        assert_eq!(result, OrderedSet::from_sorted(vec![1, 2, 3, 4]));
321    }
322
323    #[test]
324    fn test_interval_and_sorted_intersection() {
325        // Mixed-variant set algebra: Interval ∩ Sorted
326        let a = OrderedSet::interval(3, 8); // {3,4,5,6,7}
327        let b = OrderedSet::from_sorted(vec![1, 5, 7, 9]);
328        let result = a.intersection(&b);
329        assert_eq!(result, OrderedSet::from_sorted(vec![5, 7]));
330    }
331}
ferritin_core/data/ordered_set.rs

ferritin_core/data/
ordered_set.rs
