1use alloc::boxed::Box;
2use alloc::vec::Vec;
3use std::fmt;
4use std::iter::once;
5use std::iter::FusedIterator;
6
7use super::lazy_buffer::LazyBuffer;
8use crate::size_hint::{self, SizeHint};
9
10/// An iterator adaptor that iterates through all the `k`-permutations of the
11/// elements from an iterator.
12///
13/// See [`.permutations()`](crate::Itertools::permutations) for
14/// more information.
15#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
16pub struct Permutations<I: Iterator> {
17 vals: LazyBuffer<I>,
18 state: PermutationState,
19}
20
21impl<I> Clone for Permutations<I>
22where
23 I: Clone + Iterator,
24 I::Item: Clone,
25{
26 clone_fields!(vals, state);
27}
28
29#[derive(Clone, Debug)]
30enum PermutationState {
31 /// No permutation generated yet.
32 Start { k: usize },
33 /// Values from the iterator are not fully loaded yet so `n` is still unknown.
34 Buffered { k: usize, min_n: usize },
35 /// All values from the iterator are known so `n` is known.
36 Loaded {
37 indices: Box<[usize]>,
38 cycles: Box<[usize]>,
39 },
40 /// No permutation left to generate.
41 End,
42}
43
44impl<I> fmt::Debug for Permutations<I>
45where
46 I: Iterator + fmt::Debug,
47 I::Item: fmt::Debug,
48{
49 debug_fmt_fields!(Permutations, vals, state);
50}
51
52pub fn permutations<I: Iterator>(iter: I, k: usize) -> Permutations<I> {
53 Permutations {
54 vals: LazyBuffer::new(it:iter),
55 state: PermutationState::Start { k },
56 }
57}
58
59impl<I> Iterator for Permutations<I>
60where
61 I: Iterator,
62 I::Item: Clone,
63{
64 type Item = Vec<I::Item>;
65
66 fn next(&mut self) -> Option<Self::Item> {
67 let Self { vals, state } = self;
68 match state {
69 PermutationState::Start { k: 0 } => {
70 *state = PermutationState::End;
71 Some(Vec::new())
72 }
73 &mut PermutationState::Start { k } => {
74 vals.prefill(k);
75 if vals.len() != k {
76 *state = PermutationState::End;
77 return None;
78 }
79 *state = PermutationState::Buffered { k, min_n: k };
80 Some(vals[0..k].to_vec())
81 }
82 PermutationState::Buffered { ref k, min_n } => {
83 if vals.get_next() {
84 let item = (0..*k - 1)
85 .chain(once(*min_n))
86 .map(|i| vals[i].clone())
87 .collect();
88 *min_n += 1;
89 Some(item)
90 } else {
91 let n = *min_n;
92 let prev_iteration_count = n - *k + 1;
93 let mut indices: Box<[_]> = (0..n).collect();
94 let mut cycles: Box<[_]> = (n - k..n).rev().collect();
95 // Advance the state to the correct point.
96 for _ in 0..prev_iteration_count {
97 if advance(&mut indices, &mut cycles) {
98 *state = PermutationState::End;
99 return None;
100 }
101 }
102 let item = indices[0..*k].iter().map(|&i| vals[i].clone()).collect();
103 *state = PermutationState::Loaded { indices, cycles };
104 Some(item)
105 }
106 }
107 PermutationState::Loaded { indices, cycles } => {
108 if advance(indices, cycles) {
109 *state = PermutationState::End;
110 return None;
111 }
112 let k = cycles.len();
113 Some(indices[0..k].iter().map(|&i| vals[i].clone()).collect())
114 }
115 PermutationState::End => None,
116 }
117 }
118
119 fn count(self) -> usize {
120 let Self { vals, state } = self;
121 let n = vals.count();
122 state.size_hint_for(n).1.unwrap()
123 }
124
125 fn size_hint(&self) -> SizeHint {
126 let (mut low, mut upp) = self.vals.size_hint();
127 low = self.state.size_hint_for(low).0;
128 upp = upp.and_then(|n| self.state.size_hint_for(n).1);
129 (low, upp)
130 }
131}
132
133impl<I> FusedIterator for Permutations<I>
134where
135 I: Iterator,
136 I::Item: Clone,
137{
138}
139
140fn advance(indices: &mut [usize], cycles: &mut [usize]) -> bool {
141 let n: usize = indices.len();
142 let k: usize = cycles.len();
143 // NOTE: if `cycles` are only zeros, then we reached the last permutation.
144 for i: usize in (0..k).rev() {
145 if cycles[i] == 0 {
146 cycles[i] = n - i - 1;
147 indices[i..].rotate_left(mid:1);
148 } else {
149 let swap_index: usize = n - cycles[i];
150 indices.swap(a:i, b:swap_index);
151 cycles[i] -= 1;
152 return false;
153 }
154 }
155 true
156}
157
158impl PermutationState {
159 fn size_hint_for(&self, n: usize) -> SizeHint {
160 // At the beginning, there are `n!/(n-k)!` items to come.
161 let at_start = |n, k| {
162 debug_assert!(n >= k);
163 let total = (n - k + 1..=n).try_fold(1usize, |acc, i| acc.checked_mul(i));
164 (total.unwrap_or(usize::MAX), total)
165 };
166 match *self {
167 Self::Start { k } if n < k => (0, Some(0)),
168 Self::Start { k } => at_start(n, k),
169 Self::Buffered { k, min_n } => {
170 // Same as `Start` minus the previously generated items.
171 size_hint::sub_scalar(at_start(n, k), min_n - k + 1)
172 }
173 Self::Loaded {
174 ref indices,
175 ref cycles,
176 } => {
177 let count = cycles.iter().enumerate().try_fold(0usize, |acc, (i, &c)| {
178 acc.checked_mul(indices.len() - i)
179 .and_then(|count| count.checked_add(c))
180 });
181 (count.unwrap_or(usize::MAX), count)
182 }
183 Self::End => (0, Some(0)),
184 }
185 }
186}
187