dfir_pipes/pull/
take_while.rs

1use core::pin::Pin;
2
3use pin_project_lite::pin_project;
4
5use crate::Yes;
6use crate::pull::{Pull, PullStep};
7
8pin_project! {
9    /// Pull combinator that yields items while a predicate returns `true`.
10    #[must_use = "`Pull`s do nothing unless polled"]
11    #[derive(Clone, Debug)]
12    pub struct TakeWhile<Prev, Func> {
13        #[pin]
14        prev: Prev,
15        func: Func,
16    }
17}
18
19impl<Prev, Func> TakeWhile<Prev, Func>
20where
21    Self: Pull,
22{
23    pub(crate) const fn new(prev: Prev, func: Func) -> Self {
24        Self { prev, func }
25    }
26}
27
28impl<Prev, Func> Pull for TakeWhile<Prev, Func>
29where
30    Prev: Pull,
31    Func: FnMut(&Prev::Item) -> bool,
32{
33    type Ctx<'ctx> = Prev::Ctx<'ctx>;
34
35    type Item = Prev::Item;
36    type Meta = Prev::Meta;
37    type CanPend = Prev::CanPend;
38    type CanEnd = Yes;
39
40    fn pull(
41        self: Pin<&mut Self>,
42        ctx: &mut Self::Ctx<'_>,
43    ) -> PullStep<Self::Item, Self::Meta, Self::CanPend, Self::CanEnd> {
44        let this = self.project();
45
46        match this.prev.pull(ctx) {
47            PullStep::Ready(item, meta) => {
48                if (this.func)(&item) {
49                    PullStep::Ready(item, meta)
50                } else {
51                    PullStep::Ended(Yes)
52                }
53            }
54            PullStep::Pending(can_pend) => PullStep::Pending(can_pend),
55            PullStep::Ended(_) => PullStep::Ended(Yes),
56        }
57    }
58
59    fn size_hint(&self) -> (usize, Option<usize>) {
60        let (_, upper) = self.prev.size_hint();
61        (0, upper)
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use core::pin::pin;
68
69    use crate::pull::Pull;
70    use crate::pull::test_utils::TestPull;
71
72    #[test]
73    fn take_while_basic() {
74        let mut pull = TestPull::items_fused(0..5);
75        let mut p = pin!((&mut pull).take_while(|x| *x < 2));
76        let step = p.as_mut().pull(&mut ());
77        assert!(matches!(step, crate::pull::PullStep::Ready(0, _)));
78        let step = p.as_mut().pull(&mut ());
79        assert!(matches!(step, crate::pull::PullStep::Ready(1, _)));
80        let step = p.as_mut().pull(&mut ());
81        assert!(matches!(step, crate::pull::PullStep::Ended(_)));
82    }
83}
dfir_pipes/pull/take_while.rs

dfir_pipes/pull/
take_while.rs
