const_combinations/
combinations.rs

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

const_combinations/
combinations.rs
