async_itertools/features/
coalesce.rs

1use core::{
2    pin::Pin,
3    task::{Context, Poll, ready},
4};
5
6use futures_lite::Stream;
7use option_entry::OptionEntry;
8use pin_project_lite::pin_project;
9
10use crate::internal::check::assert_stream;
11
12pin_project! {
13    #[project = Proj]
14    pub struct CoalesceBy<S: Stream, F, C: CountItem<S::Item>> {
15        #[pin]
16        stream: S,
17        last: Option<Option<C::CItem>>,
18        f: F,
19    }
20}
21
22pub trait CoalescePredicate<Item, T> {
23    fn coalesce_pair(&mut self, t: T, item: Item) -> Result<T, (T, T)>;
24}
25
26impl<S: Stream, F: CoalescePredicate<S::Item, C::CItem>, C: CountItem<S::Item>> Stream
27    for CoalesceBy<S, F, C>
28{
29    type Item = C::CItem;
30
31    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
32        let Proj {
33            mut stream,
34            last,
35            f,
36        } = self.project();
37        let current = match last {
38            Some(elt) => elt,
39            None => last.insert(ready!(stream.as_mut().poll_next(cx)).map(C::new)),
40        };
41        let option_entry::Entry::Occupied(mut entry) = current.entry() else {
42            return Poll::Ready(None);
43        };
44        while let Some(next) = ready!(stream.as_mut().poll_next(cx)) {
45            match f.coalesce_pair(entry.remove(), next) {
46                Ok(item) => entry = current.entry().insert_entry(item),
47                Err((last, next)) => {
48                    *current = Some(next);
49                    return Poll::Ready(Some(last));
50                }
51            }
52        }
53        Poll::Ready(Some(entry.remove()))
54    }
55}
56
57pub struct NoCount;
58
59pub trait CountItem<T> {
60    type CItem;
61    fn new(t: T) -> Self::CItem;
62}
63
64impl<T> CountItem<T> for NoCount {
65    type CItem = T;
66    fn new(t: T) -> Self::CItem {
67        t
68    }
69}
70
71pub type Coalesce<S, F> = CoalesceBy<S, F, NoCount>;
72
73impl<F, Item, T> CoalescePredicate<Item, T> for F
74where
75    F: FnMut(T, Item) -> Result<T, (T, T)>,
76{
77    fn coalesce_pair(&mut self, t: T, item: Item) -> Result<T, (T, T)> {
78        self(t, item)
79    }
80}
81
82pub fn coalesce<S: Stream, F>(stream: S, f: F) -> Coalesce<S, F>
83where
84    F: FnMut(S::Item, S::Item) -> Result<S::Item, (S::Item, S::Item)>,
85{
86    assert_stream(CoalesceBy {
87        stream,
88        last: None,
89        f,
90    })
91}
92
93pub type DedupBy<S, Pred> = CoalesceBy<S, DedupPred2CoalescePred<Pred>, NoCount>;
94
95pub struct DedupPred2CoalescePred<DP>(DP);
96
97pub trait DedupPredicate<T> {
98    fn dedup_pair(&mut self, a: &T, b: &T) -> bool;
99}
100
101impl<DP, T> CoalescePredicate<T, T> for DedupPred2CoalescePred<DP>
102where
103    DP: DedupPredicate<T>,
104{
105    fn coalesce_pair(&mut self, t: T, item: T) -> Result<T, (T, T)> {
106        if self.0.dedup_pair(&t, &item) {
107            Ok(t)
108        } else {
109            Err((t, item))
110        }
111    }
112}
113
114pub struct DedupEq;
115
116impl<T: PartialEq> DedupPredicate<T> for DedupEq {
117    fn dedup_pair(&mut self, a: &T, b: &T) -> bool {
118        a == b
119    }
120}
121
122impl<T, F: FnMut(&T, &T) -> bool> DedupPredicate<T> for F {
123    fn dedup_pair(&mut self, a: &T, b: &T) -> bool {
124        self(a, b)
125    }
126}
127
128pub fn dedup_by<S: Stream, Pred>(stream: S, dedup_pred: Pred) -> DedupBy<S, Pred>
129where
130    Pred: FnMut(&S::Item, &S::Item) -> bool,
131{
132    assert_stream(CoalesceBy {
133        stream,
134        last: None,
135        f: DedupPred2CoalescePred(dedup_pred),
136    })
137}
138
139pub type Dedup<S> = DedupBy<S, DedupEq>;
140
141pub fn dedup<S: Stream<Item: PartialEq>>(stream: S) -> Dedup<S> {
142    assert_stream(CoalesceBy {
143        stream,
144        last: None,
145        f: DedupPred2CoalescePred(DedupEq),
146    })
147}
async_itertools/features/coalesce.rs

async_itertools/features/
coalesce.rs
