1use std::fmt;
2use std::iter::FusedIterator;
3
4use super::lazy_buffer::LazyBuffer;
5use alloc::vec::Vec;
6
7/// An iterator to iterate through all the `k`-length combinations in an iterator.
8///
9/// See [`.combinations()`](crate::Itertools::combinations) for more information.
10#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
11pub struct Combinations<I: Iterator> {
12 indices: Vec<usize>,
13 pool: LazyBuffer<I>,
14 first: bool,
15}
16
17impl<I> Clone for Combinations<I>
18 where I: Clone + Iterator,
19 I::Item: Clone,
20{
21 clone_fields!(indices, pool, first);
22}
23
24impl<I> fmt::Debug for Combinations<I>
25 where I: Iterator + fmt::Debug,
26 I::Item: fmt::Debug,
27{
28 debug_fmt_fields!(Combinations, indices, pool, first);
29}
30
31/// Create a new `Combinations` from a clonable iterator.
32pub fn combinations<I>(iter: I, k: usize) -> Combinations<I>
33 where I: Iterator
34{
35 let mut pool = LazyBuffer::new(iter);
36 pool.prefill(k);
37
38 Combinations {
39 indices: (0..k).collect(),
40 pool,
41 first: true,
42 }
43}
44
45impl<I: Iterator> Combinations<I> {
46 /// Returns the length of a combination produced by this iterator.
47 #[inline]
48 pub fn k(&self) -> usize { self.indices.len() }
49
50 /// Returns the (current) length of the pool from which combination elements are
51 /// selected. This value can change between invocations of [`next`](Combinations::next).
52 #[inline]
53 pub fn n(&self) -> usize { self.pool.len() }
54
55 /// Returns a reference to the source iterator.
56 #[inline]
57 pub(crate) fn src(&self) -> &I { &self.pool.it }
58
59 /// Resets this `Combinations` back to an initial state for combinations of length
60 /// `k` over the same pool data source. If `k` is larger than the current length
61 /// of the data pool an attempt is made to prefill the pool so that it holds `k`
62 /// elements.
63 pub(crate) fn reset(&mut self, k: usize) {
64 self.first = true;
65
66 if k < self.indices.len() {
67 self.indices.truncate(k);
68 for i in 0..k {
69 self.indices[i] = i;
70 }
71
72 } else {
73 for i in 0..self.indices.len() {
74 self.indices[i] = i;
75 }
76 self.indices.extend(self.indices.len()..k);
77 self.pool.prefill(k);
78 }
79 }
80}
81
82impl<I> Iterator for Combinations<I>
83 where I: Iterator,
84 I::Item: Clone
85{
86 type Item = Vec<I::Item>;
87 fn next(&mut self) -> Option<Self::Item> {
88 if self.first {
89 if self.k() > self.n() {
90 return None;
91 }
92 self.first = false;
93 } else if self.indices.is_empty() {
94 return None;
95 } else {
96 // Scan from the end, looking for an index to increment
97 let mut i: usize = self.indices.len() - 1;
98
99 // Check if we need to consume more from the iterator
100 if self.indices[i] == self.pool.len() - 1 {
101 self.pool.get_next(); // may change pool size
102 }
103
104 while self.indices[i] == i + self.pool.len() - self.indices.len() {
105 if i > 0 {
106 i -= 1;
107 } else {
108 // Reached the last combination
109 return None;
110 }
111 }
112
113 // Increment index, and reset the ones to its right
114 self.indices[i] += 1;
115 for j in i+1..self.indices.len() {
116 self.indices[j] = self.indices[j - 1] + 1;
117 }
118 }
119
120 // Create result vector based on the indices
121 Some(self.indices.iter().map(|i| self.pool[*i].clone()).collect())
122 }
123}
124
125impl<I> FusedIterator for Combinations<I>
126 where I: Iterator,
127 I::Item: Clone
128{}
129