lender/adapters/
peekable.rs

1use alloc::boxed::Box;
2use core::{fmt, ops::ControlFlow};
3
4use crate::{
5    try_trait_v2::{FromResidual, Try},
6    DoubleEndedLender, ExactSizeLender, FusedLender, Lend, Lender, Lending,
7};
8
9#[must_use = "lenders are lazy and do nothing unless consumed"]
10pub struct Peekable<'this, L>
11where
12    L: Lender,
13{
14    lender: Box<L>,
15    peeked: Option<Option<Lend<'this, L>>>,
16}
17impl<'this, L> Peekable<'this, L>
18where
19    L: Lender,
20{
21    pub(crate) fn new(lender: L) -> Peekable<'this, L> {
22        Peekable { lender: Box::new(lender), peeked: None }
23    }
24    pub fn into_inner(self) -> L {
25        *self.lender
26    }
27    pub fn peek(&mut self) -> Option<&'_ Lend<'this, L>> {
28        let lender = &mut self.lender;
29        self.peeked
30            .get_or_insert_with(|| {
31                // SAFETY: The lend is manually guaranteed to be the only one alive
32                unsafe { core::mem::transmute::<Option<Lend<'_, L>>, Option<Lend<'this, L>>>(lender.next()) }
33            })
34            .as_ref()
35    }
36    pub fn peek_mut(&mut self) -> Option<&'_ mut Lend<'this, L>> {
37        let lender = &mut self.lender;
38        self.peeked
39            .get_or_insert_with(|| {
40                // SAFETY: The lend is manually guaranteed to be the only one alive
41                unsafe { core::mem::transmute::<Option<Lend<'_, L>>, Option<Lend<'this, L>>>(lender.next()) }
42            })
43            .as_mut()
44    }
45    pub fn next_if<F>(&mut self, f: F) -> Option<Lend<'_, L>>
46    where
47        F: FnOnce(&Lend<'_, L>) -> bool,
48    {
49        let peeked = unsafe { &mut *(&mut self.peeked as *mut _) };
50        match self.next() {
51            Some(v) if f(&v) => Some(v),
52            v => {
53                // SAFETY: The lend is manually guaranteed to be the only one alive
54                *peeked = Some(unsafe { core::mem::transmute::<Option<Lend<'_, L>>, Option<Lend<'this, L>>>(v) });
55                None
56            }
57        }
58    }
59    pub fn next_if_eq<'a, T>(&'a mut self, t: &T) -> Option<Lend<'a, L>>
60    where
61        T: for<'all> PartialEq<Lend<'all, L>>,
62    {
63        self.next_if(|v| t == v)
64    }
65}
66impl<L> Clone for Peekable<'_, L>
67where
68    L: Lender + Clone,
69{
70    fn clone(&self) -> Self {
71        Peekable { lender: self.lender.clone(), peeked: None }
72    }
73}
74impl<'this, L: fmt::Debug> fmt::Debug for Peekable<'this, L>
75where
76    L: Lender + fmt::Debug,
77    Lend<'this, L>: fmt::Debug,
78{
79    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
80        f.debug_struct("Peekable").field("lender", &self.lender).field("peeked", &self.peeked).finish()
81    }
82}
83impl<'lend, L> Lending<'lend> for Peekable<'_, L>
84where
85    L: Lender,
86{
87    type Lend = Lend<'lend, L>;
88}
89impl<'this, L> Lender for Peekable<'this, L>
90where
91    L: Lender,
92{
93    fn next(&mut self) -> Option<Lend<'_, Self>> {
94        match self.peeked.take() {
95            // SAFETY: The lend is manually guaranteed to be the only one alive
96            Some(peeked) => unsafe { core::mem::transmute::<Option<Lend<'this, Self>>, Option<Lend<'_, Self>>>(peeked) },
97            None => self.lender.next(),
98        }
99    }
100    #[inline]
101    fn count(mut self) -> usize {
102        match self.peeked.take() {
103            Some(None) => 0,
104            Some(Some(_)) => 1 + self.lender.count(),
105            None => self.lender.count(),
106        }
107    }
108    #[inline]
109    fn nth(&mut self, n: usize) -> Option<Lend<'_, Self>> {
110        match self.peeked.take() {
111            Some(None) => None,
112            // SAFETY: The lend is manually guaranteed to be the only one alive
113            Some(v @ Some(_)) if n == 0 => unsafe {
114                core::mem::transmute::<Option<Lend<'this, Self>>, Option<Lend<'_, Self>>>(v)
115            },
116            Some(Some(_)) => self.lender.nth(n - 1),
117            None => self.lender.nth(n),
118        }
119    }
120    #[inline]
121    fn last<'a>(&'a mut self) -> Option<Lend<'a, Self>>
122    where
123        Self: Sized,
124    {
125        let peek_opt = match self.peeked.take() {
126            Some(None) => return None,
127            // SAFETY: 'this: 'call
128            Some(v) => unsafe { core::mem::transmute::<Option<Lend<'this, Self>>, Option<Lend<'a, Self>>>(v) },
129            None => None,
130        };
131        self.lender.last().or(peek_opt)
132    }
133    #[inline]
134    fn size_hint(&self) -> (usize, Option<usize>) {
135        let peek_len = match self.peeked {
136            Some(None) => return (0, Some(0)),
137            Some(Some(_)) => 1,
138            None => 0,
139        };
140        let (l, r) = self.lender.size_hint();
141        (l.saturating_add(peek_len), r.and_then(|r| r.checked_add(peek_len)))
142    }
143    #[inline]
144    fn try_fold<B, F, R>(&mut self, init: B, mut f: F) -> R
145    where
146        Self: Sized,
147        F: FnMut(B, Lend<'_, Self>) -> R,
148        R: Try<Output = B>,
149    {
150        let acc = match self.peeked.take() {
151            Some(None) => return Try::from_output(init),
152            Some(Some(v)) => match f(init, v).branch() {
153                ControlFlow::Break(b) => return FromResidual::from_residual(b),
154                ControlFlow::Continue(a) => a,
155            },
156            None => init,
157        };
158        self.lender.try_fold(acc, f)
159    }
160    #[inline]
161    fn fold<B, F>(mut self, init: B, mut f: F) -> B
162    where
163        Self: Sized,
164        F: FnMut(B, Lend<'_, Self>) -> B,
165    {
166        let acc = match self.peeked.take() {
167            Some(None) => return init,
168            Some(Some(v)) => f(init, v),
169            None => init,
170        };
171        self.lender.fold(acc, f)
172    }
173}
174impl<'this, L: DoubleEndedLender> DoubleEndedLender for Peekable<'this, L> {
175    #[inline]
176    fn next_back(&mut self) -> Option<Lend<'_, Self>> {
177        match self.peeked.as_mut() {
178            // SAFETY: The lend is manually guaranteed to be the only one alive
179            Some(v @ Some(_)) => self
180                .lender
181                .next_back()
182                .or_else(|| unsafe { core::mem::transmute::<Option<Lend<'this, Self>>, Option<Lend<'_, Self>>>(v.take()) }),
183            Some(None) => None,
184            None => self.lender.next_back(),
185        }
186    }
187    #[inline]
188    fn try_rfold<B, F, R>(&mut self, init: B, mut f: F) -> R
189    where
190        Self: Sized,
191        F: FnMut(B, Lend<'_, Self>) -> R,
192        R: Try<Output = B>,
193    {
194        match self.peeked.take() {
195            None => self.lender.try_rfold(init, f),
196            Some(None) => Try::from_output(init),
197            Some(Some(v)) => match self.lender.try_rfold(init, &mut f).branch() {
198                ControlFlow::Continue(acc) => f(acc, v),
199                ControlFlow::Break(r) => {
200                    self.peeked = Some(Some(v));
201                    FromResidual::from_residual(r)
202                }
203            },
204        }
205    }
206    #[inline]
207    fn rfold<B, F>(mut self, init: B, mut f: F) -> B
208    where
209        Self: Sized,
210        F: FnMut(B, Lend<'_, Self>) -> B,
211    {
212        match self.peeked.take() {
213            None => self.lender.rfold(init, f),
214            Some(None) => init,
215            Some(Some(v)) => {
216                let acc = self.lender.rfold(init, &mut f);
217                f(acc, v)
218            }
219        }
220    }
221}
222impl<L: ExactSizeLender> ExactSizeLender for Peekable<'_, L> {}
223
224impl<L: FusedLender> FusedLender for Peekable<'_, L> {}
225
226#[cfg(test)]
227mod test {
228    use super::*;
229
230    struct ArrayLender {
231        array: [i32; 4],
232    }
233
234    impl<'lend> Lending<'lend> for ArrayLender {
235        type Lend = &'lend i32;
236    }
237
238    impl Lender for ArrayLender {
239        fn next(&mut self) -> Option<Lend<'_, Self>> {
240            Some(&self.array[0])
241        }
242    }
243
244    // This test will fail if Peekable stores L instead of Box<L>. In that case,
245    // when Peekable<ArrayLender> is moved, the array inside ArrayLender is
246    // moved, too, but Peekable.peeked will still contain a reference to the
247    // previous location.
248    #[test]
249    fn test_peekable() {
250        let lender = ArrayLender { array: [-1, 1, 2, 3] };
251        let mut peekable = lender.peekable();
252        assert_eq!(**peekable.peek().unwrap(), -1);
253        assert_eq!(peekable.peeked.unwrap().unwrap() as *const _, &peekable.lender.array[0] as *const _);
254        moved_peekable(peekable);
255    }
256
257    fn moved_peekable(peekable: Peekable<ArrayLender>) {
258        let peeked = peekable.peeked.unwrap().unwrap() as *const _;
259        let array = &peekable.lender.array[0] as *const _;
260        assert_eq!(peeked, array, "Peeked element pointer should point to the first element of the array");
261    }
262}
lender/adapters/peekable.rs

lender/adapters/
peekable.rs
