const_combinations/
permutations.rs

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

const_combinations/
permutations.rs
