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

ferritin_core/data/
ordered_set.rs
