Skip to main content

wonfy_tools/util/iter/
padded_iter.rs

1use std::iter::FusedIterator;
2
3#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4pub enum PaddingSide {
5    Start,
6    End,
7}
8
9#[derive(Debug, Clone)]
10pub struct PaddedIter<I>
11where
12    I: Iterator,
13    I::Item: Clone,
14{
15    inner: I,
16    padding_item: I::Item,
17    padding_remaining: usize,
18    side: PaddingSide,
19    inner_exhausted: bool,
20}
21
22impl<I> PaddedIter<I>
23where
24    I: Iterator,
25    I::Item: Clone,
26{
27    pub fn new(inner: I, padding_item: I::Item, padding_count: usize, side: PaddingSide) -> Self {
28        PaddedIter {
29            inner,
30            padding_item,
31            padding_remaining: padding_count,
32            side,
33            inner_exhausted: false,
34        }
35    }
36}
37
38impl<I> Iterator for PaddedIter<I>
39where
40    I: Iterator,
41    I::Item: Clone,
42{
43    type Item = I::Item;
44
45    fn next(&mut self) -> Option<Self::Item> {
46        if self.side == PaddingSide::Start && self.padding_remaining > 0 {
47            self.padding_remaining -= 1;
48            return Some(self.padding_item.clone());
49        }
50
51        if !self.inner_exhausted {
52            match self.inner.next() {
53                Some(item) => return Some(item),
54                None => {
55                    self.inner_exhausted = true;
56
57                    if self.side == PaddingSide::Start || self.padding_remaining == 0 {
58                        return None;
59                    }
60                }
61            }
62        }
63
64        if self.inner_exhausted && self.side == PaddingSide::End && self.padding_remaining > 0 {
65            self.padding_remaining -= 1;
66            return Some(self.padding_item.clone());
67        }
68
69        None
70    }
71
72    fn size_hint(&self) -> (usize, Option<usize>) {
73        let (inner_low, inner_high) = self.inner.size_hint();
74        // let padding_to_add = if self.inner_exhausted && self.side == PaddingSide::End {
75        //     self.padding_remaining
76        // } else if !self.inner_exhausted && self.side == PaddingSide::Start {
77        //     self.padding_remaining
78        // } else if !self.inner_exhausted && self.side == PaddingSide::End {
79        //     0
80        // } else {
81        //     0
82        // };
83        let padding_to_add = if self.inner_exhausted && self.side == PaddingSide::End {
84            self.padding_remaining
85        } else if !self.inner_exhausted && self.side == PaddingSide::Start {
86            self.padding_remaining
87        } else {
88            0
89        };
90
91        let low = inner_low.saturating_add(padding_to_add);
92        let high = match inner_high {
93            Some(h) => Some(h.saturating_add(self.padding_remaining)),
94            None => None,
95        };
96        (low, high)
97    }
98}
99
100impl<I> DoubleEndedIterator for PaddedIter<I>
101where
102    I: DoubleEndedIterator,
103    I::Item: Clone,
104{
105    fn next_back(&mut self) -> Option<Self::Item> {
106        if self.side == PaddingSide::End && self.padding_remaining > 0 {
107            self.padding_remaining -= 1;
108            return Some(self.padding_item.clone());
109        }
110
111        if !self.inner_exhausted {
112            match self.inner.next_back() {
113                Some(item) => return Some(item),
114                None => {
115                    self.inner_exhausted = true;
116
117                    if self.side == PaddingSide::End || self.padding_remaining == 0 {
118                        return None;
119                    }
120                }
121            }
122        }
123
124        if self.inner_exhausted && self.side == PaddingSide::Start && self.padding_remaining > 0 {
125            self.padding_remaining -= 1;
126            return Some(self.padding_item.clone());
127        }
128
129        None
130    }
131}
132
133impl<I> FusedIterator for PaddedIter<I>
134where
135    I: FusedIterator,
136    I::Item: Clone,
137{
138}
139
140pub trait PadExt: Iterator
141where
142    Self::Item: Clone,
143    Self: Sized,
144{
145    fn pad_start(self, item: Self::Item, count: usize) -> PaddedIter<Self> {
146        PaddedIter::new(self, item, count, PaddingSide::Start)
147    }
148
149    fn pad_end(self, item: Self::Item, count: usize) -> PaddedIter<Self> {
150        PaddedIter::new(self, item, count, PaddingSide::End)
151    }
152}
153
154impl<I> PadExt for I
155where
156    I: Iterator,
157    <I as std::iter::Iterator>::Item: Clone,
158{
159}