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, T>
8where
9    I: Iterator,
10{
11    iter: I,
12    last: Option<T>,
13    f: F,
14}
15
16impl<I: Clone, F: Clone, T: Clone> Clone for CoalesceBy<I, F, T>
17where
18    I: Iterator,
19{
20    clone_fields!(last, iter, f);
21}
22
23impl<I, F, T> fmt::Debug for CoalesceBy<I, F, T>
24where
25    I: Iterator + fmt::Debug,
26    T: fmt::Debug,
27{
28    debug_fmt_fields!(CoalesceBy, iter);
29}
30
31pub trait CoalescePredicate<Item, T> {
32    fn coalesce_pair(&mut self, t: T, item: Item) -> Result<T, (T, T)>;
33}
34
35impl<I, F, T> Iterator for CoalesceBy<I, F, T>
36where
37    I: Iterator,
38    F: CoalescePredicate<I::Item, T>,
39{
40    type Item = T;
41
42    fn next(&mut self) -> Option<Self::Item> {
43        // this fuses the iterator
44        let last = self.last.take()?;
45
46        let self_last = &mut self.last;
47        let self_f = &mut self.f;
48        Some(
49            self.iter
50                .try_fold(last, |last, next| match self_f.coalesce_pair(last, next) {
51                    Ok(joined) => Ok(joined),
52                    Err((last_, next_)) => {
53                        *self_last = Some(next_);
54                        Err(last_)
55                    }
56                })
57                .unwrap_or_else(|x| x),
58        )
59    }
60
61    fn size_hint(&self) -> (usize, Option<usize>) {
62        let (low, hi) = size_hint::add_scalar(self.iter.size_hint(), self.last.is_some() as usize);
63        ((low > 0) as usize, hi)
64    }
65
66    fn fold<Acc, FnAcc>(self, acc: Acc, mut fn_acc: FnAcc) -> Acc
67    where
68        FnAcc: FnMut(Acc, Self::Item) -> Acc,
69    {
70        if let Some(last) = self.last {
71            let mut f = self.f;
72            let (last, acc) = self.iter.fold((last, acc), |(last, acc), elt| {
73                match f.coalesce_pair(last, elt) {
74                    Ok(joined) => (joined, acc),
75                    Err((last_, next_)) => (next_, fn_acc(acc, last_)),
76                }
77            });
78            fn_acc(acc, last)
79        } else {
80            acc
81        }
82    }
83}
84
85impl<I: Iterator, F: CoalescePredicate<I::Item, T>, T> FusedIterator for CoalesceBy<I, F, T> {}
86
87/// An iterator adaptor that may join together adjacent elements.
88///
89/// See [`.coalesce()`](crate::Itertools::coalesce) for more information.
90pub type Coalesce<I, F> = CoalesceBy<I, F, <I as Iterator>::Item>;
91
92impl<F, Item, T> CoalescePredicate<Item, T> for F
93where
94    F: FnMut(T, Item) -> Result<T, (T, T)>,
95{
96    fn coalesce_pair(&mut self, t: T, item: Item) -> Result<T, (T, T)> {
97        self(t, item)
98    }
99}
100
101/// Create a new `Coalesce`.
102pub fn coalesce<I, F>(mut iter: I, f: F) -> Coalesce<I, F>
103where
104    I: Iterator,
105{
106    Coalesce {
107        last: iter.next(),
108        iter,
109        f,
110    }
111}
112
113/// An iterator adaptor that removes repeated duplicates, determining equality using a comparison function.
114///
115/// See [`.dedup_by()`](crate::Itertools::dedup_by) or [`.dedup()`](crate::Itertools::dedup) for more information.
116pub type DedupBy<I, Pred> = CoalesceBy<I, DedupPred2CoalescePred<Pred>, <I as Iterator>::Item>;
117
118#[derive(Clone)]
119pub struct DedupPred2CoalescePred<DP>(DP);
120
121impl<DP> fmt::Debug for DedupPred2CoalescePred<DP> {
122    debug_fmt_fields!(DedupPred2CoalescePred,);
123}
124
125pub trait DedupPredicate<T> {
126    // TODO replace by Fn(&T, &T)->bool once Rust supports it
127    fn dedup_pair(&mut self, a: &T, b: &T) -> bool;
128}
129
130impl<DP, T> CoalescePredicate<T, T> for DedupPred2CoalescePred<DP>
131where
132    DP: DedupPredicate<T>,
133{
134    fn coalesce_pair(&mut self, t: T, item: T) -> Result<T, (T, T)> {
135        if self.0.dedup_pair(&t, &item) {
136            Ok(t)
137        } else {
138            Err((t, item))
139        }
140    }
141}
142
143#[derive(Clone, Debug)]
144pub struct DedupEq;
145
146impl<T: PartialEq> DedupPredicate<T> for DedupEq {
147    fn dedup_pair(&mut self, a: &T, b: &T) -> bool {
148        a == b
149    }
150}
151
152impl<T, F: FnMut(&T, &T) -> bool> DedupPredicate<T> for F {
153    fn dedup_pair(&mut self, a: &T, b: &T) -> bool {
154        self(a, b)
155    }
156}
157
158/// Create a new `DedupBy`.
159pub fn dedup_by<I, Pred>(mut iter: I, dedup_pred: Pred) -> DedupBy<I, Pred>
160where
161    I: Iterator,
162{
163    DedupBy {
164        last: iter.next(),
165        iter,
166        f: DedupPred2CoalescePred(dedup_pred),
167    }
168}
169
170/// An iterator adaptor that removes repeated duplicates.
171///
172/// See [`.dedup()`](crate::Itertools::dedup) for more information.
173pub type Dedup<I> = DedupBy<I, DedupEq>;
174
175/// Create a new `Dedup`.
176pub fn dedup<I>(iter: I) -> Dedup<I>
177where
178    I: Iterator,
179{
180    dedup_by(iter, DedupEq)
181}
182
183/// An iterator adaptor that removes repeated duplicates, while keeping a count of how many
184/// repeated elements were present. This will determine equality using a comparison function.
185///
186/// See [`.dedup_by_with_count()`](crate::Itertools::dedup_by_with_count) or
187/// [`.dedup_with_count()`](crate::Itertools::dedup_with_count) for more information.
188pub type DedupByWithCount<I, Pred> =
189    CoalesceBy<I, DedupPredWithCount2CoalescePred<Pred>, (usize, <I as Iterator>::Item)>;
190
191#[derive(Clone, Debug)]
192pub struct DedupPredWithCount2CoalescePred<DP>(DP);
193
194impl<DP, T> CoalescePredicate<T, (usize, T)> for DedupPredWithCount2CoalescePred<DP>
195where
196    DP: DedupPredicate<T>,
197{
198    fn coalesce_pair(
199        &mut self,
200        (c, t): (usize, T),
201        item: T,
202    ) -> Result<(usize, T), ((usize, T), (usize, T))> {
203        if self.0.dedup_pair(&t, &item) {
204            Ok((c + 1, t))
205        } else {
206            Err(((c, t), (1, item)))
207        }
208    }
209}
210
211/// An iterator adaptor that removes repeated duplicates, while keeping a count of how many
212/// repeated elements were present.
213///
214/// See [`.dedup_with_count()`](crate::Itertools::dedup_with_count) for more information.
215pub type DedupWithCount<I> = DedupByWithCount<I, DedupEq>;
216
217/// Create a new `DedupByWithCount`.
218pub fn dedup_by_with_count<I, Pred>(mut iter: I, dedup_pred: Pred) -> DedupByWithCount<I, Pred>
219where
220    I: Iterator,
221{
222    DedupByWithCount {
223        last: iter.next().map(|v| (1, v)),
224        iter,
225        f: DedupPredWithCount2CoalescePred(dedup_pred),
226    }
227}
228
229/// Create a new `DedupWithCount`.
230pub fn dedup_with_count<I>(iter: I) -> DedupWithCount<I>
231where
232    I: Iterator,
233{
234    dedup_by_with_count(iter, DedupEq)
235}