itertools/
tuple_impl.rs

1//! Some iterator that produces tuples
2
3use std::iter::Cycle;
4use std::iter::Fuse;
5use std::iter::FusedIterator;
6
7use crate::size_hint;
8
9// `HomogeneousTuple` is a public facade for `TupleCollect`, allowing
10// tuple-related methods to be used by clients in generic contexts, while
11// hiding the implementation details of `TupleCollect`.
12// See https://github.com/rust-itertools/itertools/issues/387
13
14/// Implemented for homogeneous tuples of size up to 12.
15pub trait HomogeneousTuple: TupleCollect {}
16
17impl<T: TupleCollect> HomogeneousTuple for T {}
18
19/// An iterator over a incomplete tuple.
20///
21/// See [`.tuples()`](crate::Itertools::tuples) and
22/// [`Tuples::into_buffer()`].
23#[derive(Clone, Debug)]
24pub struct TupleBuffer<T>
25where
26    T: HomogeneousTuple,
27{
28    cur: usize,
29    buf: T::Buffer,
30}
31
32impl<T> TupleBuffer<T>
33where
34    T: HomogeneousTuple,
35{
36    fn new(buf: T::Buffer) -> Self {
37        Self { cur: 0, buf }
38    }
39}
40
41impl<T> Iterator for TupleBuffer<T>
42where
43    T: HomogeneousTuple,
44{
45    type Item = T::Item;
46
47    fn next(&mut self) -> Option<Self::Item> {
48        let s = self.buf.as_mut();
49        if let Some(ref mut item) = s.get_mut(self.cur) {
50            self.cur += 1;
51            item.take()
52        } else {
53            None
54        }
55    }
56
57    fn size_hint(&self) -> (usize, Option<usize>) {
58        let buffer = &self.buf.as_ref()[self.cur..];
59        let len = if buffer.is_empty() {
60            0
61        } else {
62            buffer
63                .iter()
64                .position(|x| x.is_none())
65                .unwrap_or(buffer.len())
66        };
67        (len, Some(len))
68    }
69}
70
71impl<T> ExactSizeIterator for TupleBuffer<T> where T: HomogeneousTuple {}
72
73/// An iterator that groups the items in tuples of a specific size.
74///
75/// See [`.tuples()`](crate::Itertools::tuples) for more information.
76#[derive(Clone, Debug)]
77#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
78pub struct Tuples<I, T>
79where
80    I: Iterator<Item = T::Item>,
81    T: HomogeneousTuple,
82{
83    iter: Fuse<I>,
84    buf: T::Buffer,
85}
86
87/// Create a new tuples iterator.
88pub fn tuples<I, T>(iter: I) -> Tuples<I, T>
89where
90    I: Iterator<Item = T::Item>,
91    T: HomogeneousTuple,
92{
93    Tuples {
94        iter: iter.fuse(),
95        buf: Default::default(),
96    }
97}
98
99impl<I, T> Iterator for Tuples<I, T>
100where
101    I: Iterator<Item = T::Item>,
102    T: HomogeneousTuple,
103{
104    type Item = T;
105
106    fn next(&mut self) -> Option<Self::Item> {
107        T::collect_from_iter(&mut self.iter, &mut self.buf)
108    }
109
110    fn size_hint(&self) -> (usize, Option<usize>) {
111        // The number of elts we've drawn from the underlying iterator, but have
112        // not yet produced as a tuple.
113        let buffered = T::buffer_len(&self.buf);
114        // To that, we must add the size estimates of the underlying iterator.
115        let (unbuffered_lo, unbuffered_hi) = self.iter.size_hint();
116        // The total low estimate is the sum of the already-buffered elements,
117        // plus the low estimate of remaining unbuffered elements, divided by
118        // the tuple size.
119        let total_lo = add_then_div(unbuffered_lo, buffered, T::num_items()).unwrap_or(usize::MAX);
120        // And likewise for the total high estimate, but using the high estimate
121        // of the remaining unbuffered elements.
122        let total_hi = unbuffered_hi.and_then(|hi| add_then_div(hi, buffered, T::num_items()));
123        (total_lo, total_hi)
124    }
125}
126
127/// `(n + a) / d` avoiding overflow when possible, returns `None` if it overflows.
128fn add_then_div(n: usize, a: usize, d: usize) -> Option<usize> {
129    debug_assert_ne!(d, 0);
130    (n / d).checked_add(a / d)?.checked_add((n % d + a % d) / d)
131}
132
133impl<I, T> ExactSizeIterator for Tuples<I, T>
134where
135    I: ExactSizeIterator<Item = T::Item>,
136    T: HomogeneousTuple,
137{
138}
139
140impl<I, T> Tuples<I, T>
141where
142    I: Iterator<Item = T::Item>,
143    T: HomogeneousTuple,
144{
145    /// Return a buffer with the produced items that was not enough to be grouped in a tuple.
146    ///
147    /// ```
148    /// use itertools::Itertools;
149    ///
150    /// let mut iter = (0..5).tuples();
151    /// assert_eq!(Some((0, 1, 2)), iter.next());
152    /// assert_eq!(None, iter.next());
153    /// itertools::assert_equal(vec![3, 4], iter.into_buffer());
154    /// ```
155    pub fn into_buffer(self) -> TupleBuffer<T> {
156        TupleBuffer::new(self.buf)
157    }
158}
159
160/// An iterator over all contiguous windows that produces tuples of a specific size.
161///
162/// See [`.tuple_windows()`](crate::Itertools::tuple_windows) for more
163/// information.
164#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
165#[derive(Clone, Debug)]
166pub struct TupleWindows<I, T>
167where
168    I: Iterator<Item = T::Item>,
169    T: HomogeneousTuple,
170{
171    iter: I,
172    last: Option<T>,
173}
174
175/// Create a new tuple windows iterator.
176pub fn tuple_windows<I, T>(iter: I) -> TupleWindows<I, T>
177where
178    I: Iterator<Item = T::Item>,
179    T: HomogeneousTuple,
180    T::Item: Clone,
181{
182    TupleWindows { last: None, iter }
183}
184
185impl<I, T> Iterator for TupleWindows<I, T>
186where
187    I: Iterator<Item = T::Item>,
188    T: HomogeneousTuple + Clone,
189    T::Item: Clone,
190{
191    type Item = T;
192
193    fn next(&mut self) -> Option<Self::Item> {
194        if T::num_items() == 1 {
195            return T::collect_from_iter_no_buf(&mut self.iter);
196        }
197        if let Some(new) = self.iter.next() {
198            if let Some(ref mut last) = self.last {
199                last.left_shift_push(new);
200                Some(last.clone())
201            } else {
202                use std::iter::once;
203                let iter = once(new).chain(&mut self.iter);
204                self.last = T::collect_from_iter_no_buf(iter);
205                self.last.clone()
206            }
207        } else {
208            None
209        }
210    }
211
212    fn size_hint(&self) -> (usize, Option<usize>) {
213        let mut sh = self.iter.size_hint();
214        // Adjust the size hint at the beginning
215        // OR when `num_items == 1` (but it does not change the size hint).
216        if self.last.is_none() {
217            sh = size_hint::sub_scalar(sh, T::num_items() - 1);
218        }
219        sh
220    }
221}
222
223impl<I, T> ExactSizeIterator for TupleWindows<I, T>
224where
225    I: ExactSizeIterator<Item = T::Item>,
226    T: HomogeneousTuple + Clone,
227    T::Item: Clone,
228{
229}
230
231impl<I, T> FusedIterator for TupleWindows<I, T>
232where
233    I: FusedIterator<Item = T::Item>,
234    T: HomogeneousTuple + Clone,
235    T::Item: Clone,
236{
237}
238
239/// An iterator over all windows, wrapping back to the first elements when the
240/// window would otherwise exceed the length of the iterator, producing tuples
241/// of a specific size.
242///
243/// See [`.circular_tuple_windows()`](crate::Itertools::circular_tuple_windows) for more
244/// information.
245#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
246#[derive(Debug, Clone)]
247pub struct CircularTupleWindows<I, T>
248where
249    I: Iterator<Item = T::Item> + Clone,
250    T: TupleCollect + Clone,
251{
252    iter: TupleWindows<Cycle<I>, T>,
253    len: usize,
254}
255
256pub fn circular_tuple_windows<I, T>(iter: I) -> CircularTupleWindows<I, T>
257where
258    I: Iterator<Item = T::Item> + Clone + ExactSizeIterator,
259    T: TupleCollect + Clone,
260    T::Item: Clone,
261{
262    let len = iter.len();
263    let iter = tuple_windows(iter.cycle());
264
265    CircularTupleWindows { iter, len }
266}
267
268impl<I, T> Iterator for CircularTupleWindows<I, T>
269where
270    I: Iterator<Item = T::Item> + Clone,
271    T: TupleCollect + Clone,
272    T::Item: Clone,
273{
274    type Item = T;
275
276    fn next(&mut self) -> Option<Self::Item> {
277        if self.len != 0 {
278            self.len -= 1;
279            self.iter.next()
280        } else {
281            None
282        }
283    }
284
285    fn size_hint(&self) -> (usize, Option<usize>) {
286        (self.len, Some(self.len))
287    }
288}
289
290impl<I, T> ExactSizeIterator for CircularTupleWindows<I, T>
291where
292    I: Iterator<Item = T::Item> + Clone,
293    T: TupleCollect + Clone,
294    T::Item: Clone,
295{
296}
297
298impl<I, T> FusedIterator for CircularTupleWindows<I, T>
299where
300    I: Iterator<Item = T::Item> + Clone,
301    T: TupleCollect + Clone,
302    T::Item: Clone,
303{
304}
305
306pub trait TupleCollect: Sized {
307    type Item;
308    type Buffer: Default + AsRef<[Option<Self::Item>]> + AsMut<[Option<Self::Item>]>;
309
310    fn buffer_len(buf: &Self::Buffer) -> usize {
311        let s = buf.as_ref();
312        s.iter().position(Option::is_none).unwrap_or(s.len())
313    }
314
315    fn collect_from_iter<I>(iter: I, buf: &mut Self::Buffer) -> Option<Self>
316    where
317        I: IntoIterator<Item = Self::Item>;
318
319    fn collect_from_iter_no_buf<I>(iter: I) -> Option<Self>
320    where
321        I: IntoIterator<Item = Self::Item>;
322
323    fn num_items() -> usize;
324
325    fn left_shift_push(&mut self, item: Self::Item);
326}
327
328macro_rules! rev_for_each_ident{
329    ($m:ident, ) => {};
330    ($m:ident, $i0:ident, $($i:ident,)*) => {
331        rev_for_each_ident!($m, $($i,)*);
332        $m!($i0);
333    };
334}
335
336macro_rules! impl_tuple_collect {
337    ($dummy:ident,) => {}; // stop
338    ($dummy:ident, $($Y:ident,)*) => (
339        impl_tuple_collect!($($Y,)*);
340        impl<A> TupleCollect for ($(ignore_ident!($Y, A),)*) {
341            type Item = A;
342            type Buffer = [Option<A>; count_ident!($($Y)*) - 1];
343
344            #[allow(unused_assignments, unused_mut)]
345            fn collect_from_iter<I>(iter: I, buf: &mut Self::Buffer) -> Option<Self>
346                where I: IntoIterator<Item = A>
347            {
348                let mut iter = iter.into_iter();
349                $(
350                    let mut $Y = None;
351                )*
352
353                loop {
354                    $(
355                        $Y = iter.next();
356                        if $Y.is_none() {
357                            break
358                        }
359                    )*
360                    return Some(($($Y.unwrap()),*,))
361                }
362
363                let mut i = 0;
364                let mut s = buf.as_mut();
365                $(
366                    if i < s.len() {
367                        s[i] = $Y;
368                        i += 1;
369                    }
370                )*
371                return None;
372            }
373
374            fn collect_from_iter_no_buf<I>(iter: I) -> Option<Self>
375                where I: IntoIterator<Item = A>
376            {
377                let mut iter = iter.into_iter();
378
379                Some(($(
380                    { let $Y = iter.next()?; $Y },
381                )*))
382            }
383
384            fn num_items() -> usize {
385                count_ident!($($Y)*)
386            }
387
388            fn left_shift_push(&mut self, mut item: A) {
389                use std::mem::replace;
390
391                let &mut ($(ref mut $Y),*,) = self;
392                macro_rules! replace_item{($i:ident) => {
393                    item = replace($i, item);
394                }}
395                rev_for_each_ident!(replace_item, $($Y,)*);
396                drop(item);
397            }
398        }
399    )
400}
401impl_tuple_collect!(dummy, a, b, c, d, e, f, g, h, i, j, k, l,);