blob: 03867cbf0c4d50efbfe15feb48141df6da729a89 [file] [log] [blame]
Joel Galensonb593e252021-06-21 13:15:57 -07001use std::iter::{Fuse, FusedIterator};
Jakub Kotura425e552020-12-21 17:28:15 +01002use crate::size_hint;
3
4/// An iterator adaptor that pads a sequence to a minimum length by filling
5/// missing elements using a closure.
6///
7/// Iterator element type is `I::Item`.
8///
Joel Galensonb593e252021-06-21 13:15:57 -07009/// See [`.pad_using()`](crate::Itertools::pad_using) for more information.
Jakub Kotura425e552020-12-21 17:28:15 +010010#[derive(Clone)]
11#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
12pub struct PadUsing<I, F> {
13 iter: Fuse<I>,
14 min: usize,
15 pos: usize,
16 filler: F,
17}
18
19/// Create a new **PadUsing** iterator.
20pub fn pad_using<I, F>(iter: I, min: usize, filler: F) -> PadUsing<I, F>
21 where I: Iterator,
22 F: FnMut(usize) -> I::Item
23{
24 PadUsing {
25 iter: iter.fuse(),
26 min,
27 pos: 0,
28 filler,
29 }
30}
31
32impl<I, F> Iterator for PadUsing<I, F>
33 where I: Iterator,
34 F: FnMut(usize) -> I::Item
35{
36 type Item = I::Item;
37
38 #[inline]
Joel Galenson6f798712021-04-01 17:03:06 -070039 fn next(&mut self) -> Option<Self::Item> {
Jakub Kotura425e552020-12-21 17:28:15 +010040 match self.iter.next() {
41 None => {
42 if self.pos < self.min {
43 let e = Some((self.filler)(self.pos));
44 self.pos += 1;
45 e
46 } else {
47 None
48 }
49 },
50 e => {
51 self.pos += 1;
52 e
53 }
54 }
55 }
56
57 fn size_hint(&self) -> (usize, Option<usize>) {
58 let tail = self.min.saturating_sub(self.pos);
59 size_hint::max(self.iter.size_hint(), (tail, Some(tail)))
60 }
61}
62
63impl<I, F> DoubleEndedIterator for PadUsing<I, F>
64 where I: DoubleEndedIterator + ExactSizeIterator,
65 F: FnMut(usize) -> I::Item
66{
Joel Galenson6f798712021-04-01 17:03:06 -070067 fn next_back(&mut self) -> Option<Self::Item> {
Jakub Kotura425e552020-12-21 17:28:15 +010068 if self.min == 0 {
69 self.iter.next_back()
70 } else if self.iter.len() >= self.min {
71 self.min -= 1;
72 self.iter.next_back()
73 } else {
74 self.min -= 1;
75 Some((self.filler)(self.min))
76 }
77 }
78}
79
80impl<I, F> ExactSizeIterator for PadUsing<I, F>
81 where I: ExactSizeIterator,
82 F: FnMut(usize) -> I::Item
83{}
Joel Galensonb593e252021-06-21 13:15:57 -070084
85
86impl<I, F> FusedIterator for PadUsing<I, F>
87 where I: FusedIterator,
88 F: FnMut(usize) -> I::Item
89{}