Skip to main content

itertools/adaptors/
coalesce.rs

1use std::fmt;
2use std::iter::FusedIterator;
3
4use crate::size_hint;
5
6#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
7pub struct CoalesceBy<I, F, C>
8where
9    I: Iterator,
10    C: CountItem<I::Item>,
11{
12    iter: I,
13    /// `last` is `None` while no item have been taken out of `iter` (at definition).
14    /// Then `last` will be `Some(Some(item))` until `iter` is exhausted,
15    /// in which case `last` will be `Some(None)`.
16    last: Option<Option<C::CItem>>,
17    f: F,
18}
19
20impl<I, F, C> Clone for CoalesceBy<I, F, C>
21where
22    I: Clone + Iterator,
23    F: Clone,
24    C: CountItem<I::Item>,
25    C::CItem: Clone,
26{
27    clone_fields!(last, iter, f);
28}
29
30impl<I, F, C> fmt::Debug for CoalesceBy<I, F, C>
31where
32    I: Iterator + fmt::Debug,
33    C: CountItem<I::Item>,
34    C::CItem: fmt::Debug,
35{
36    debug_fmt_fields!(CoalesceBy, iter, last);
37}
38
39pub trait CoalescePredicate<Item, T> {
40    fn coalesce_pair(&mut self, t: T, item: Item) -> Result<T, (T, T)>;
41}
42
43impl<I, F, C> Iterator for CoalesceBy<I, F, C>
44where
45    I: Iterator,
46    F: CoalescePredicate<I::Item, C::CItem>,
47    C: CountItem<I::Item>,
48{
49    type Item = C::CItem;
50
51    fn next(&mut self) -> Option<Self::Item> {
52        let Self { iter, last, f } = self;
53        // this fuses the iterator
54        let init = match last {
55            Some(elt) => elt.take(),
56            None => {
57                *last = Some(None);
58                iter.next().map(C::new)
59            }
60        }?;
61
62        Some(
63            iter.try_fold(init, |accum, next| match f.coalesce_pair(accum, next) {
64                Ok(joined) => Ok(joined),
65                Err((last_, next_)) => {
66                    *last = Some(Some(next_));
67                    Err(last_)
68                }
69            })
70            .unwrap_or_else(|x| x),
71        )
72    }
73
74    fn size_hint(&self) -> (usize, Option<usize>) {
75        let (low, hi) = size_hint::add_scalar(
76            self.iter.size_hint(),
77            usize::from(matches!(self.last, Some(Some(_)))),
78        );
79        (usize::from(low > 0), hi)
80    }
81
82    fn fold<Acc, FnAcc>(self, acc: Acc, mut fn_acc: FnAcc) -> Acc
83    where
84        FnAcc: FnMut(Acc, Self::Item) -> Acc,
85    {
86        let Self {
87            mut iter,
88            last,
89            mut f,
90        } = self;
91        if let Some(last) = last.unwrap_or_else(|| iter.next().map(C::new)) {
92            let (last, acc) = iter.fold((last, acc), |(last, acc), elt| {
93                match f.coalesce_pair(last, elt) {
94                    Ok(joined) => (joined, acc),
95                    Err((last_, next_)) => (next_, fn_acc(acc, last_)),
96                }
97            });
98            fn_acc(acc, last)
99        } else {
100            acc
101        }
102    }
103}
104
105impl<I, F, C> FusedIterator for CoalesceBy<I, F, C>
106where
107    I: Iterator,
108    F: CoalescePredicate<I::Item, C::CItem>,
109    C: CountItem<I::Item>,
110{
111}
112
113#[derive(Debug)]
114pub struct NoCount;
115
116#[derive(Debug)]
117pub struct WithCount;
118
119pub trait CountItem<T> {
120    type CItem;
121    fn new(t: T) -> Self::CItem;
122}
123
124impl<T> CountItem<T> for NoCount {
125    type CItem = T;
126    #[inline(always)]
127    fn new(t: T) -> T {
128        t
129    }
130}
131
132impl<T> CountItem<T> for WithCount {
133    type CItem = (usize, T);
134    #[inline(always)]
135    fn new(t: T) -> (usize, T) {
136        (1, t)
137    }
138}
139
140/// An iterator adaptor that may join together adjacent elements.
141///
142/// See [`.coalesce()`](crate::Itertools::coalesce) for more information.
143pub type Coalesce<I, F> = CoalesceBy<I, F, NoCount>;
144
145impl<F, Item, T> CoalescePredicate<Item, T> for F
146where
147    F: FnMut(T, Item) -> Result<T, (T, T)>,
148{
149    fn coalesce_pair(&mut self, t: T, item: Item) -> Result<T, (T, T)> {
150        self(t, item)
151    }
152}
153
154/// Create a new `Coalesce`.
155pub fn coalesce<I, F>(iter: I, f: F) -> Coalesce<I, F>
156where
157    I: Iterator,
158{
159    Coalesce {
160        last: None,
161        iter,
162        f,
163    }
164}
165
166/// An iterator adaptor that removes repeated duplicates, determining equality using a comparison function.
167///
168/// See [`.dedup_by()`](crate::Itertools::dedup_by) or [`.dedup()`](crate::Itertools::dedup) for more information.
169pub type DedupBy<I, Pred> = CoalesceBy<I, DedupPred2CoalescePred<Pred>, NoCount>;
170
171#[derive(Clone)]
172pub struct DedupPred2CoalescePred<DP>(DP);
173
174impl<DP> fmt::Debug for DedupPred2CoalescePred<DP> {
175    debug_fmt_fields!(DedupPred2CoalescePred,);
176}
177
178pub trait DedupPredicate<T> {
179    // TODO replace by Fn(&T, &T)->bool once Rust supports it
180    fn dedup_pair(&mut self, a: &T, b: &T) -> bool;
181}
182
183impl<DP, T> CoalescePredicate<T, T> for DedupPred2CoalescePred<DP>
184where
185    DP: DedupPredicate<T>,
186{
187    fn coalesce_pair(&mut self, t: T, item: T) -> Result<T, (T, T)> {
188        if self.0.dedup_pair(&t, &item) {
189            Ok(t)
190        } else {
191            Err((t, item))
192        }
193    }
194}
195
196#[derive(Clone, Debug)]
197pub struct DedupEq;
198
199impl<T: PartialEq> DedupPredicate<T> for DedupEq {
200    fn dedup_pair(&mut self, a: &T, b: &T) -> bool {
201        a == b
202    }
203}
204
205impl<T, F: FnMut(&T, &T) -> bool> DedupPredicate<T> for F {
206    fn dedup_pair(&mut self, a: &T, b: &T) -> bool {
207        self(a, b)
208    }
209}
210
211/// Create a new `DedupBy`.
212pub fn dedup_by<I, Pred>(iter: I, dedup_pred: Pred) -> DedupBy<I, Pred>
213where
214    I: Iterator,
215{
216    DedupBy {
217        last: None,
218        iter,
219        f: DedupPred2CoalescePred(dedup_pred),
220    }
221}
222
223/// An iterator adaptor that removes repeated duplicates.
224///
225/// See [`.dedup()`](crate::Itertools::dedup) for more information.
226pub type Dedup<I> = DedupBy<I, DedupEq>;
227
228/// Create a new `Dedup`.
229pub fn dedup<I>(iter: I) -> Dedup<I>
230where
231    I: Iterator,
232{
233    dedup_by(iter, DedupEq)
234}
235
236/// An iterator adaptor that removes repeated duplicates, while keeping a count of how many
237/// repeated elements were present. This will determine equality using a comparison function.
238///
239/// See [`.dedup_by_with_count()`](crate::Itertools::dedup_by_with_count) or
240/// [`.dedup_with_count()`](crate::Itertools::dedup_with_count) for more information.
241pub type DedupByWithCount<I, Pred> =
242    CoalesceBy<I, DedupPredWithCount2CoalescePred<Pred>, WithCount>;
243
244#[derive(Clone, Debug)]
245pub struct DedupPredWithCount2CoalescePred<DP>(DP);
246
247impl<DP, T> CoalescePredicate<T, (usize, T)> for DedupPredWithCount2CoalescePred<DP>
248where
249    DP: DedupPredicate<T>,
250{
251    fn coalesce_pair(
252        &mut self,
253        (c, t): (usize, T),
254        item: T,
255    ) -> Result<(usize, T), ((usize, T), (usize, T))> {
256        if self.0.dedup_pair(&t, &item) {
257            Ok((c + 1, t))
258        } else {
259            Err(((c, t), (1, item)))
260        }
261    }
262}
263
264/// An iterator adaptor that removes repeated duplicates, while keeping a count of how many
265/// repeated elements were present.
266///
267/// See [`.dedup_with_count()`](crate::Itertools::dedup_with_count) for more information.
268pub type DedupWithCount<I> = DedupByWithCount<I, DedupEq>;
269
270/// Create a new `DedupByWithCount`.
271pub fn dedup_by_with_count<I, Pred>(iter: I, dedup_pred: Pred) -> DedupByWithCount<I, Pred>
272where
273    I: Iterator,
274{
275    DedupByWithCount {
276        last: None,
277        iter,
278        f: DedupPredWithCount2CoalescePred(dedup_pred),
279    }
280}
281
282/// Create a new `DedupWithCount`.
283pub fn dedup_with_count<I>(iter: I) -> DedupWithCount<I>
284where
285    I: Iterator,
286{
287    dedup_by_with_count(iter, DedupEq)
288}