1use alloc::vec::Vec;
2use std::fmt;
3use std::iter::once;
4
5use super::lazy_buffer::LazyBuffer;
6
7/// An iterator adaptor that iterates through all the `k`-permutations of the
8/// elements from an iterator.
9///
10/// See [`.permutations()`](crate::Itertools::permutations) for
11/// more information.
12#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
13pub struct Permutations<I: Iterator> {
14 vals: LazyBuffer<I>,
15 state: PermutationState,
16}
17
18impl<I> Clone for Permutations<I>
19 where I: Clone + Iterator,
20 I::Item: Clone,
21{
22 clone_fields!(vals, state);
23}
24
25#[derive(Clone, Debug)]
26enum PermutationState {
27 StartUnknownLen {
28 k: usize,
29 },
30 OngoingUnknownLen {
31 k: usize,
32 min_n: usize,
33 },
34 Complete(CompleteState),
35 Empty,
36}
37
38#[derive(Clone, Debug)]
39enum CompleteState {
40 Start {
41 n: usize,
42 k: usize,
43 },
44 Ongoing {
45 indices: Vec<usize>,
46 cycles: Vec<usize>,
47 }
48}
49
50enum CompleteStateRemaining {
51 Known(usize),
52 Overflow,
53}
54
55impl<I> fmt::Debug for Permutations<I>
56 where I: Iterator + fmt::Debug,
57 I::Item: fmt::Debug,
58{
59 debug_fmt_fields!(Permutations, vals, state);
60}
61
62pub fn permutations<I: Iterator>(iter: I, k: usize) -> Permutations<I> {
63 let mut vals = LazyBuffer::new(iter);
64
65 if k == 0 {
66 // Special case, yields single empty vec; `n` is irrelevant
67 let state = PermutationState::Complete(CompleteState::Start { n: 0, k: 0 });
68
69 return Permutations {
70 vals,
71 state
72 };
73 }
74
75 let mut enough_vals = true;
76
77 while vals.len() < k {
78 if !vals.get_next() {
79 enough_vals = false;
80 break;
81 }
82 }
83
84 let state = if enough_vals {
85 PermutationState::StartUnknownLen { k }
86 } else {
87 PermutationState::Empty
88 };
89
90 Permutations {
91 vals,
92 state
93 }
94}
95
96impl<I> Iterator for Permutations<I>
97where
98 I: Iterator,
99 I::Item: Clone
100{
101 type Item = Vec<I::Item>;
102
103 fn next(&mut self) -> Option<Self::Item> {
104 self.advance();
105
106 let &mut Permutations { ref vals, ref state } = self;
107
108 match *state {
109 PermutationState::StartUnknownLen { .. } => panic!("unexpected iterator state"),
110 PermutationState::OngoingUnknownLen { k, min_n } => {
111 let latest_idx = min_n - 1;
112 let indices = (0..(k - 1)).chain(once(latest_idx));
113
114 Some(indices.map(|i| vals[i].clone()).collect())
115 }
116 PermutationState::Complete(CompleteState::Ongoing { ref indices, ref cycles }) => {
117 let k = cycles.len();
118 Some(indices[0..k].iter().map(|&i| vals[i].clone()).collect())
119 },
120 PermutationState::Complete(CompleteState::Start { .. }) | PermutationState::Empty => None
121 }
122 }
123
124 fn count(self) -> usize {
125 fn from_complete(complete_state: CompleteState) -> usize {
126 match complete_state.remaining() {
127 CompleteStateRemaining::Known(count) => count,
128 CompleteStateRemaining::Overflow => {
129 panic!("Iterator count greater than usize::MAX");
130 }
131 }
132 }
133
134 let Permutations { vals, state } = self;
135 match state {
136 PermutationState::StartUnknownLen { k } => {
137 let n = vals.len() + vals.it.count();
138 let complete_state = CompleteState::Start { n, k };
139
140 from_complete(complete_state)
141 }
142 PermutationState::OngoingUnknownLen { k, min_n } => {
143 let prev_iteration_count = min_n - k + 1;
144 let n = vals.len() + vals.it.count();
145 let complete_state = CompleteState::Start { n, k };
146
147 from_complete(complete_state) - prev_iteration_count
148 },
149 PermutationState::Complete(state) => from_complete(state),
150 PermutationState::Empty => 0
151 }
152 }
153
154 fn size_hint(&self) -> (usize, Option<usize>) {
155 match self.state {
156 PermutationState::StartUnknownLen { .. } |
157 PermutationState::OngoingUnknownLen { .. } => (0, None), // TODO can we improve this lower bound?
158 PermutationState::Complete(ref state) => match state.remaining() {
159 CompleteStateRemaining::Known(count) => (count, Some(count)),
160 CompleteStateRemaining::Overflow => (::std::usize::MAX, None)
161 }
162 PermutationState::Empty => (0, Some(0))
163 }
164 }
165}
166
167impl<I> Permutations<I>
168where
169 I: Iterator,
170 I::Item: Clone
171{
172 fn advance(&mut self) {
173 let &mut Permutations { ref mut vals, ref mut state } = self;
174
175 *state = match *state {
176 PermutationState::StartUnknownLen { k } => {
177 PermutationState::OngoingUnknownLen { k, min_n: k }
178 }
179 PermutationState::OngoingUnknownLen { k, min_n } => {
180 if vals.get_next() {
181 PermutationState::OngoingUnknownLen { k, min_n: min_n + 1 }
182 } else {
183 let n = min_n;
184 let prev_iteration_count = n - k + 1;
185 let mut complete_state = CompleteState::Start { n, k };
186
187 // Advance the complete-state iterator to the correct point
188 for _ in 0..(prev_iteration_count + 1) {
189 complete_state.advance();
190 }
191
192 PermutationState::Complete(complete_state)
193 }
194 }
195 PermutationState::Complete(ref mut state) => {
196 state.advance();
197
198 return;
199 }
200 PermutationState::Empty => { return; }
201 };
202 }
203}
204
205impl CompleteState {
206 fn advance(&mut self) {
207 *self = match *self {
208 CompleteState::Start { n, k } => {
209 let indices = (0..n).collect();
210 let cycles = ((n - k)..n).rev().collect();
211
212 CompleteState::Ongoing {
213 cycles,
214 indices
215 }
216 },
217 CompleteState::Ongoing { ref mut indices, ref mut cycles } => {
218 let n = indices.len();
219 let k = cycles.len();
220
221 for i in (0..k).rev() {
222 if cycles[i] == 0 {
223 cycles[i] = n - i - 1;
224
225 let to_push = indices.remove(i);
226 indices.push(to_push);
227 } else {
228 let swap_index = n - cycles[i];
229 indices.swap(i, swap_index);
230
231 cycles[i] -= 1;
232 return;
233 }
234 }
235
236 CompleteState::Start { n, k }
237 }
238 }
239 }
240
241 fn remaining(&self) -> CompleteStateRemaining {
242 use self::CompleteStateRemaining::{Known, Overflow};
243
244 match *self {
245 CompleteState::Start { n, k } => {
246 if n < k {
247 return Known(0);
248 }
249
250 let count: Option<usize> = (n - k + 1..n + 1).fold(Some(1), |acc, i| {
251 acc.and_then(|acc| acc.checked_mul(i))
252 });
253
254 match count {
255 Some(count) => Known(count),
256 None => Overflow
257 }
258 }
259 CompleteState::Ongoing { ref indices, ref cycles } => {
260 let mut count: usize = 0;
261
262 for (i, &c) in cycles.iter().enumerate() {
263 let radix = indices.len() - i;
264 let next_count = count.checked_mul(radix)
265 .and_then(|count| count.checked_add(c));
266
267 count = match next_count {
268 Some(count) => count,
269 None => { return Overflow; }
270 };
271 }
272
273 Known(count)
274 }
275 }
276 }
277}
278