Skip to main content

itertools/
pad_tail.rs

1use std::iter::{Fuse, FusedIterator};
2
3/// An iterator adaptor that pads a sequence to a minimum length by filling
4/// missing elements using a closure.
5///
6/// Iterator element type is `I::Item`.
7///
8/// See [`.pad_using()`](crate::Itertools::pad_using) for more information.
9#[derive(Clone)]
10#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
11pub struct PadUsing<I, F> {
12    iter: Fuse<I>,
13    elements_from_next: usize,
14    elements_from_next_back: usize,
15    elements_required: usize,
16    filler: F,
17}
18
19impl<I, F> std::fmt::Debug for PadUsing<I, F>
20where
21    I: std::fmt::Debug,
22{
23    debug_fmt_fields!(
24        PadUsing,
25        iter,
26        elements_from_next,
27        elements_from_next_back,
28        elements_required
29    );
30}
31
32/// Create a new `PadUsing` iterator.
33pub fn pad_using<I, F>(iter: I, elements_required: usize, filler: F) -> PadUsing<I, F>
34where
35    I: Iterator,
36    F: FnMut(usize) -> I::Item,
37{
38    PadUsing {
39        iter: iter.fuse(),
40        elements_from_next: 0,
41        elements_from_next_back: 0,
42        elements_required,
43        filler,
44    }
45}
46
47impl<I, F> Iterator for PadUsing<I, F>
48where
49    I: Iterator,
50    F: FnMut(usize) -> I::Item,
51{
52    type Item = I::Item;
53
54    #[inline]
55    fn next(&mut self) -> Option<Self::Item> {
56        let total_consumed = self.elements_from_next + self.elements_from_next_back;
57
58        if total_consumed >= self.elements_required {
59            self.iter.next()
60        } else if let Some(e) = self.iter.next() {
61            self.elements_from_next += 1;
62            Some(e)
63        } else {
64            let e = (self.filler)(self.elements_from_next);
65            self.elements_from_next += 1;
66            Some(e)
67        }
68    }
69
70    fn size_hint(&self) -> (usize, Option<usize>) {
71        let total_consumed = self.elements_from_next + self.elements_from_next_back;
72
73        if total_consumed >= self.elements_required {
74            return self.iter.size_hint();
75        }
76
77        let elements_remaining = self.elements_required - total_consumed;
78        let (low, high) = self.iter.size_hint();
79
80        let lower_bound = low.max(elements_remaining);
81        let upper_bound = high.map(|h| h.max(elements_remaining));
82
83        (lower_bound, upper_bound)
84    }
85}
86
87impl<I, F> DoubleEndedIterator for PadUsing<I, F>
88where
89    I: DoubleEndedIterator + ExactSizeIterator,
90    F: FnMut(usize) -> I::Item,
91{
92    fn next_back(&mut self) -> Option<Self::Item> {
93        let total_consumed = self.elements_from_next + self.elements_from_next_back;
94
95        if total_consumed >= self.elements_required {
96            return self.iter.next_back();
97        }
98
99        let elements_remaining = self.elements_required - total_consumed;
100        self.elements_from_next_back += 1;
101
102        if self.iter.len() < elements_remaining {
103            Some((self.filler)(
104                self.elements_required - self.elements_from_next_back,
105            ))
106        } else {
107            let item = self.iter.next_back();
108            debug_assert!(item.is_some()); // If this triggers, we should not have incremented elements_from_next_back, and the input iterator mistakenly reported that it would be able to produce at least elements_remaining items.
109            item
110        }
111    }
112}
113
114impl<I, F> ExactSizeIterator for PadUsing<I, F>
115where
116    I: ExactSizeIterator,
117    F: FnMut(usize) -> I::Item,
118{
119}
120
121impl<I, F> FusedIterator for PadUsing<I, F>
122where
123    I: FusedIterator,
124    F: FnMut(usize) -> I::Item,
125{
126}