const_combinations/
combinations.rs

1use alloc::vec::Vec;
2use core::iter::{FusedIterator, Iterator};
3
4#[derive(Clone)]
5pub struct LazyCombinationGenerator<const K: usize> {
6    indices: [usize; K],
7    done: bool,
8}
9
10impl<const K: usize> LazyCombinationGenerator<K> {
11    pub fn new() -> Self {
12        Self {
13            indices: core::array::from_fn(|i| i),
14            done: false,
15        }
16    }
17
18    pub fn max_index(&self) -> Option<usize> {
19        self.indices.last().copied()
20    }
21
22    pub fn is_done(&self, item_count: usize) -> bool {
23        self.done || self.max_index() >= Some(item_count)
24    }
25
26    pub fn indices(&self) -> &[usize; K] {
27        &self.indices
28    }
29
30    pub fn step(&mut self) {
31        if K == 0 {
32            self.done = true;
33        } else {
34            let mut i = 0;
35            // Reset consecutive indices
36            while i + 1 < K && self.indices[i] + 1 == self.indices[i + 1] {
37                self.indices[i] = i;
38                i += 1;
39            }
40            // Increment the last consecutive index
41            self.indices[i] += 1;
42        }
43    }
44}
45
46#[derive(Clone)]
47struct State<const K: usize> {
48    gen: LazyCombinationGenerator<K>,
49}
50
51impl<const K: usize> State<K> {
52    fn new() -> Self {
53        Self {
54            gen: LazyCombinationGenerator::new(),
55        }
56    }
57
58    fn max_index(&self) -> Option<usize> {
59        self.gen.max_index()
60    }
61
62    fn get_and_step<'a, T, O, F>(&mut self, items: &'a [T], f: F) -> Option<[O; K]>
63    where
64        F: Fn(&'a T) -> O,
65        O: 'a,
66    {
67        if self.gen.is_done(items.len()) {
68            None
69        } else {
70            let indices = self.gen.indices();
71            let res = core::array::from_fn(|i| f(&items[indices[i]]));
72            self.gen.step();
73            Some(res)
74        }
75    }
76}
77
78/// An iterator that returns k-length combinations of values from `iter`.
79///
80/// This `struct` is created by the [`combinations`] method on [`IterExt`]. See its
81/// documentation for more.
82///
83/// [`combinations`]: super::IterExt::combinations
84/// [`IterExt`]: super::IterExt
85#[derive(Clone)]
86#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
87pub struct Combinations<I, const K: usize>
88where
89    I: Iterator,
90{
91    iter: I,
92    items: Vec<I::Item>,
93    state: State<K>,
94}
95
96impl<I, const K: usize> Combinations<I, K>
97where
98    I: Iterator,
99{
100    pub(crate) fn new(iter: I) -> Self {
101        Self {
102            iter,
103            items: Vec::new(),
104            state: State::new(),
105        }
106    }
107}
108
109impl<I, const K: usize> Iterator for Combinations<I, K>
110where
111    I: Iterator,
112    I::Item: Clone,
113{
114    type Item = [I::Item; K];
115
116    fn next(&mut self) -> Option<[I::Item; K]> {
117        if K > 0 {
118            let max_index = self.state.max_index().unwrap();
119            let missing_count = (max_index + 1).saturating_sub(self.items.len());
120            if missing_count > 0 {
121                // Try to fill the buffer
122                self.items.extend(self.iter.by_ref().take(missing_count));
123            }
124        }
125        self.state.get_and_step(&self.items, |t| t.clone())
126    }
127}
128
129impl<I, const K: usize> FusedIterator for Combinations<I, K>
130where
131    I: FusedIterator,
132    I::Item: Clone,
133{
134}
135
136/// An iterator that returns k-length combinations of values from `slice`.
137#[derive(Clone)]
138#[must_use = "iterator does nothing unless consumed"]
139pub struct SliceCombinations<'a, T, const K: usize> {
140    items: &'a [T],
141    state: State<K>,
142}
143
144impl<'a, T, const K: usize> SliceCombinations<'a, T, K> {
145    pub(crate) fn new(items: &'a [T]) -> Self {
146        Self {
147            items,
148            state: State::new(),
149        }
150    }
151}
152
153impl<'a, T, const K: usize> Iterator for SliceCombinations<'a, T, K> {
154    type Item = [&'a T; K];
155
156    fn next(&mut self) -> Option<[&'a T; K]> {
157        self.state.get_and_step(self.items, |t| t)
158    }
159}
160
161impl<T, const K: usize> FusedIterator for SliceCombinations<'_, T, K> {}
162
163#[cfg(test)]
164mod test {
165    use super::*;
166    use crate::IterExt;
167    use core::sync::atomic::{AtomicUsize, Ordering};
168
169    #[test]
170    fn order() {
171        let mut combinations = (1..6).combinations();
172        assert_eq!(combinations.next(), Some([1, 2, 3]));
173        assert_eq!(combinations.next(), Some([1, 2, 4]));
174        assert_eq!(combinations.next(), Some([1, 3, 4]));
175        assert_eq!(combinations.next(), Some([2, 3, 4]));
176        assert_eq!(combinations.next(), Some([1, 2, 5]));
177        assert_eq!(combinations.next(), Some([1, 3, 5]));
178        assert_eq!(combinations.next(), Some([2, 3, 5]));
179        assert_eq!(combinations.next(), Some([1, 4, 5]));
180        assert_eq!(combinations.next(), Some([2, 4, 5]));
181        assert_eq!(combinations.next(), Some([3, 4, 5]));
182        assert_eq!(combinations.next(), None);
183        assert_eq!(combinations.next(), None);
184    }
185
186    #[test]
187    fn none_on_size_too_big() {
188        let mut combinations = (1..2).combinations::<2>();
189        assert_eq!(combinations.next(), None);
190        assert_eq!(combinations.next(), None);
191    }
192
193    #[test]
194    fn empty_arr_on_n_zero() {
195        let mut combinations = (1..5).combinations();
196        assert_eq!(combinations.next(), Some([]));
197        assert_eq!(combinations.next(), None);
198        assert_eq!(combinations.next(), None);
199    }
200
201    #[test]
202    fn fused_propagation() {
203        let fused = [1, 2, 3].iter().fuse();
204        let combinations = fused.combinations::<2>();
205
206        fn is_fused<T: FusedIterator>(_: T) {}
207        is_fused(combinations);
208    }
209
210    #[test]
211    fn resume_after_none() {
212        struct ResumeIter<'l, 'a, T>
213        where
214            T: Copy,
215        {
216            items: &'a [T],
217            i: usize,
218            len: &'l AtomicUsize,
219        }
220
221        impl<T> Iterator for ResumeIter<'_, '_, T>
222        where
223            T: Copy,
224        {
225            type Item = T;
226            fn next(&mut self) -> Option<T> {
227                if self.i >= self.len.load(Ordering::SeqCst) {
228                    None
229                } else {
230                    self.i += 1;
231                    Some(self.items[self.i - 1])
232                }
233            }
234        }
235
236        let len = AtomicUsize::new(0);
237        let mut combinations = ResumeIter {
238            items: &[1, 2, 3, 4],
239            len: &len,
240            i: 0,
241        }
242        .combinations();
243
244        assert_eq!(combinations.next(), None);
245
246        len.fetch_add(1, Ordering::SeqCst);
247        assert_eq!(combinations.next(), None);
248
249        len.fetch_add(1, Ordering::SeqCst);
250        assert_eq!(combinations.next(), None);
251
252        len.fetch_add(1, Ordering::SeqCst);
253        assert_eq!(combinations.next(), Some([1, 2, 3]));
254        assert_eq!(combinations.next(), None);
255
256        len.fetch_add(1, Ordering::SeqCst);
257        assert_eq!(combinations.next(), Some([1, 2, 4]));
258        assert_eq!(combinations.next(), Some([1, 3, 4]));
259        assert_eq!(combinations.next(), Some([2, 3, 4]));
260        assert_eq!(combinations.next(), None);
261        assert_eq!(combinations.next(), None);
262    }
263}
264
265#[cfg(test)]
266mod slice_test {
267    use crate::SliceExt;
268
269    #[test]
270    fn order() {
271        let mut combinations = [1, 2, 3, 4, 5].combinations();
272        assert_eq!(combinations.next(), Some([&1, &2, &3]));
273        assert_eq!(combinations.next(), Some([&1, &2, &4]));
274        assert_eq!(combinations.next(), Some([&1, &3, &4]));
275        assert_eq!(combinations.next(), Some([&2, &3, &4]));
276        assert_eq!(combinations.next(), Some([&1, &2, &5]));
277        assert_eq!(combinations.next(), Some([&1, &3, &5]));
278        assert_eq!(combinations.next(), Some([&2, &3, &5]));
279        assert_eq!(combinations.next(), Some([&1, &4, &5]));
280        assert_eq!(combinations.next(), Some([&2, &4, &5]));
281        assert_eq!(combinations.next(), Some([&3, &4, &5]));
282        assert_eq!(combinations.next(), None);
283        assert_eq!(combinations.next(), None);
284    }
285
286    #[test]
287    fn none_on_size_too_big() {
288        let mut combinations = [1].combinations::<2>();
289        assert_eq!(combinations.next(), None);
290        assert_eq!(combinations.next(), None);
291    }
292
293    #[test]
294    fn empty_arr_on_n_zero() {
295        let mut combinations = [1, 2, 3, 4].combinations();
296        assert_eq!(combinations.next(), Some([]));
297        assert_eq!(combinations.next(), None);
298        assert_eq!(combinations.next(), None);
299    }
300}
const_combinations/combinations.rs

const_combinations/
combinations.rs
