]>
Commit | Line | Data |
---|---|---|
2c00a5a8 XL |
1 | |
2 | use std::ops::Index; | |
3 | use std::fmt; | |
4 | ||
5 | /// An iterator to iterate through all the `n`-length combinations in an iterator. | |
6 | /// | |
7 | /// See [`.combinations()`](../trait.Itertools.html#method.combinations) for more information. | |
8 | #[must_use = "iterator adaptors are lazy and do nothing unless consumed"] | |
9 | pub struct Combinations<I: Iterator> { | |
10 | n: usize, | |
11 | indices: Vec<usize>, | |
12 | pool: LazyBuffer<I>, | |
13 | first: bool, | |
14 | } | |
15 | ||
16 | impl<I> fmt::Debug for Combinations<I> | |
17 | where I: Iterator + fmt::Debug, | |
18 | I::Item: fmt::Debug, | |
19 | { | |
20 | debug_fmt_fields!(Combinations, n, indices, pool, first); | |
21 | } | |
22 | ||
23 | /// Create a new `Combinations` from a clonable iterator. | |
24 | pub fn combinations<I>(iter: I, n: usize) -> Combinations<I> | |
25 | where I: Iterator | |
26 | { | |
27 | let mut indices: Vec<usize> = Vec::with_capacity(n); | |
28 | for i in 0..n { | |
29 | indices.push(i); | |
30 | } | |
31 | let mut pool: LazyBuffer<I> = LazyBuffer::new(iter); | |
32 | ||
33 | for _ in 0..n { | |
34 | if !pool.get_next() { | |
35 | break; | |
36 | } | |
37 | } | |
38 | ||
39 | Combinations { | |
40 | n: n, | |
41 | indices: indices, | |
42 | pool: pool, | |
43 | first: true, | |
44 | } | |
45 | } | |
46 | ||
47 | impl<I> Iterator for Combinations<I> | |
48 | where I: Iterator, | |
49 | I::Item: Clone | |
50 | { | |
51 | type Item = Vec<I::Item>; | |
52 | fn next(&mut self) -> Option<Self::Item> { | |
53 | let mut pool_len = self.pool.len(); | |
54 | if self.pool.is_done() { | |
55 | if pool_len == 0 || self.n > pool_len { | |
56 | return None; | |
57 | } | |
58 | } | |
59 | ||
60 | if self.first { | |
61 | self.first = false; | |
62 | } else if self.n == 0 { | |
63 | return None; | |
64 | } else { | |
65 | // Scan from the end, looking for an index to increment | |
66 | let mut i: usize = self.n - 1; | |
67 | ||
68 | // Check if we need to consume more from the iterator | |
69 | if self.indices[i] == pool_len - 1 && !self.pool.is_done() { | |
70 | if self.pool.get_next() { | |
71 | pool_len += 1; | |
72 | } | |
73 | } | |
74 | ||
75 | while self.indices[i] == i + pool_len - self.n { | |
76 | if i > 0 { | |
77 | i -= 1; | |
78 | } else { | |
79 | // Reached the last combination | |
80 | return None; | |
81 | } | |
82 | } | |
83 | ||
84 | // Increment index, and reset the ones to its right | |
85 | self.indices[i] += 1; | |
86 | let mut j = i + 1; | |
87 | while j < self.n { | |
88 | self.indices[j] = self.indices[j - 1] + 1; | |
89 | j += 1; | |
90 | } | |
91 | } | |
92 | ||
93 | // Create result vector based on the indices | |
94 | let mut result = Vec::with_capacity(self.n); | |
95 | for i in self.indices.iter() { | |
96 | result.push(self.pool[*i].clone()); | |
97 | } | |
98 | Some(result) | |
99 | } | |
100 | } | |
101 | ||
102 | #[derive(Debug)] | |
103 | struct LazyBuffer<I: Iterator> { | |
104 | it: I, | |
105 | done: bool, | |
106 | buffer: Vec<I::Item>, | |
107 | } | |
108 | ||
109 | impl<I> LazyBuffer<I> | |
110 | where I: Iterator | |
111 | { | |
112 | pub fn new(it: I) -> LazyBuffer<I> { | |
113 | let mut it = it; | |
114 | let mut buffer = Vec::new(); | |
115 | let done; | |
116 | if let Some(first) = it.next() { | |
117 | buffer.push(first); | |
118 | done = false; | |
119 | } else { | |
120 | done = true; | |
121 | } | |
122 | LazyBuffer { | |
123 | it: it, | |
124 | done: done, | |
125 | buffer: buffer, | |
126 | } | |
127 | } | |
128 | ||
129 | pub fn len(&self) -> usize { | |
130 | self.buffer.len() | |
131 | } | |
132 | ||
133 | pub fn is_done(&self) -> bool { | |
134 | self.done | |
135 | } | |
136 | ||
137 | pub fn get_next(&mut self) -> bool { | |
138 | if self.done { | |
139 | return false; | |
140 | } | |
141 | let next_item = self.it.next(); | |
142 | match next_item { | |
143 | Some(x) => { | |
144 | self.buffer.push(x); | |
145 | true | |
146 | } | |
147 | None => { | |
148 | self.done = true; | |
149 | false | |
150 | } | |
151 | } | |
152 | } | |
153 | } | |
154 | ||
155 | impl<I> Index<usize> for LazyBuffer<I> | |
156 | where I: Iterator, | |
157 | I::Item: Sized | |
158 | { | |
159 | type Output = I::Item; | |
160 | ||
161 | fn index<'b>(&'b self, _index: usize) -> &'b I::Item { | |
162 | self.buffer.index(_index) | |
163 | } | |
164 | } | |
165 |