incremental_map/
symmetric_fold.rs

1#[cfg(test)]
2use test_log::test;
3
4use std::cmp::Ordering;
5use std::collections::{
6    btree_map::{IntoIter, Keys},
7    BTreeMap,
8};
9use std::iter::Peekable;
10use std::ops::Deref;
11use std::rc::Rc;
12
13// Adapted from itertools.
14// For [1, 2, 3].merge([2, 4])
15// you should get [1, 2, 3, 4].
16struct MergeOnce<I, J>
17where
18    I: Iterator,
19    J: Iterator<Item = I::Item>,
20{
21    a: Peekable<I>,
22    b: Peekable<J>,
23    fused: Option<bool>,
24}
25
26impl<I, J> MergeOnce<I, J>
27where
28    I: Iterator,
29    J: Iterator<Item = I::Item>,
30{
31    fn new(a: I, b: J) -> Self {
32        Self {
33            a: a.peekable(),
34            b: b.peekable(),
35            fused: None,
36        }
37    }
38}
39
40impl<I, J> Iterator for MergeOnce<I, J>
41where
42    I: Iterator,
43    J: Iterator<Item = I::Item>,
44    I::Item: PartialOrd,
45{
46    type Item = I::Item;
47    fn next(&mut self) -> Option<Self::Item> {
48        let (less_than, both) = match self.fused {
49            Some(lt) => (lt, false),
50            None => match (self.a.peek(), self.b.peek()) {
51                (Some(a), Some(b)) => (a <= b, a == b),
52                (Some(_), None) => {
53                    self.fused = Some(true);
54                    (true, false)
55                }
56                (None, Some(_)) => {
57                    self.fused = Some(false);
58                    (false, false)
59                }
60                (None, None) => return None,
61            },
62        };
63
64        if less_than {
65            if both {
66                drop(self.b.next());
67            }
68            self.a.next()
69        } else {
70            if both {
71                drop(self.a.next());
72            }
73            self.b.next()
74        }
75    }
76}
77
78// Same but with a custom comparator
79pub(crate) struct MergeOnceWith<I, J, F>
80where
81    I: Iterator,
82    J: Iterator,
83    F: Fn(&I::Item, &J::Item) -> Ordering,
84{
85    a: Peekable<I>,
86    b: Peekable<J>,
87    fcmp: F,
88    fused: Option<bool>,
89}
90
91impl<I: Iterator, J: Iterator, FCmp: Fn(&I::Item, &J::Item) -> Ordering> MergeOnceWith<I, J, FCmp> {
92    pub(crate) fn new(a: I, b: J, fcmp: FCmp) -> Self {
93        Self {
94            a: a.peekable(),
95            b: b.peekable(),
96            fcmp,
97            fused: None,
98        }
99    }
100}
101
102impl<I, J, FCmp> Iterator for MergeOnceWith<I, J, FCmp>
103where
104    I: Iterator,
105    J: Iterator,
106    FCmp: Fn(&I::Item, &J::Item) -> Ordering,
107{
108    type Item = MergeElement<I::Item, J::Item>;
109    fn next(&mut self) -> Option<Self::Item> {
110        let ordering: Ordering = match self.fused {
111            Some(true) => Ordering::Less,
112            Some(false) => Ordering::Greater,
113            None => match (self.a.peek(), self.b.peek()) {
114                (Some(a), Some(b)) => (self.fcmp)(a, b),
115                (Some(_), None) => {
116                    self.fused = Some(true);
117                    Ordering::Less
118                }
119                (None, Some(_)) => {
120                    self.fused = Some(false);
121                    Ordering::Greater
122                }
123                (None, None) => return None,
124            },
125        };
126
127        match ordering {
128            Ordering::Equal => self
129                .a
130                .next()
131                .zip(self.b.next())
132                .map(|(a, b)| MergeElement::Both(a, b)),
133            Ordering::Less => self.a.next().map(MergeElement::Left),
134            Ordering::Greater => self.b.next().map(MergeElement::Right),
135        }
136    }
137}
138
139pub(crate) struct SymmetricDiffOwned<K, V> {
140    self_: Peekable<IntoIter<K, V>>,
141    other: Peekable<IntoIter<K, V>>,
142    fused: Option<bool>, // keys: MergeOnce<Keys<'a, K, V>, Keys<'a, K, V>>,
143}
144
145impl<K, V> Iterator for SymmetricDiffOwned<K, V>
146where
147    K: Ord,
148    V: PartialEq,
149{
150    type Item = DiffElement<(K, V)>;
151    fn next(&mut self) -> Option<Self::Item> {
152        let less_than = loop {
153            match self.fused {
154                Some(lt) => break lt,
155                None => match (self.self_.peek(), self.other.peek()) {
156                    (Some((ka, va)), Some((kb, vb))) => match ka.cmp(kb) {
157                        Ordering::Less => break true,
158                        Ordering::Greater => break false,
159                        Ordering::Equal => {
160                            let unequal = va != vb;
161                            let (sk, sv) = self.self_.next()?;
162                            let (ok, ov) = self.other.next()?;
163                            if unequal {
164                                return Some(DiffElement::Unequal((sk, sv), (ok, ov)));
165                            } else {
166                                continue;
167                            }
168                        }
169                    },
170                    (Some(_), None) => {
171                        self.fused = Some(true);
172                        break true;
173                    }
174                    (None, Some(_)) => {
175                        self.fused = Some(false);
176                        break false;
177                    }
178                    (None, None) => return None,
179                },
180            }
181        };
182        if less_than {
183            self.self_.next().map(|(k, v)| DiffElement::Left((k, v)))
184        } else {
185            self.other.next().map(|(k, v)| DiffElement::Right((k, v)))
186        }
187    }
188}
189
190pub(crate) struct SymmetricDiff<'a, K: 'a, V: 'a> {
191    self_: &'a BTreeMap<K, V>,
192    other: &'a BTreeMap<K, V>,
193    keys: MergeOnce<Keys<'a, K, V>, Keys<'a, K, V>>,
194}
195impl<'a, K: 'a, V: 'a> Iterator for SymmetricDiff<'a, K, V>
196where
197    K: Ord,
198    V: PartialEq,
199{
200    type Item = (&'a K, DiffElement<&'a V>);
201    fn next(&mut self) -> Option<Self::Item> {
202        let mut key;
203        let elem = loop {
204            key = self.keys.next()?;
205            let s = self.self_.get(key);
206            let o = self.other.get(key);
207            match (s, o) {
208                (Some(a), Some(b)) if a != b => break DiffElement::Unequal(a, b),
209                (Some(_), Some(_)) => continue,
210                (Some(a), _) => break DiffElement::Left(a),
211                (_, Some(b)) => break DiffElement::Right(b),
212                _ => return None,
213            }
214        };
215        Some((key, elem))
216    }
217}
218
219#[derive(Debug, PartialEq, Eq, Clone, Copy)]
220pub enum MergeElement<L, R> {
221    /// A key in `incr_merge` was present only in the left map
222    Left(L),
223    /// A key in `incr_merge` was present only in the right map
224    Right(R),
225    /// A key in `incr_merge` was present in both maps
226    Both(L, R),
227}
228
229impl<L, R> MergeElement<L, R> {
230    pub fn left(&self) -> Option<&L> {
231        match self {
232            Self::Left(l) | Self::Both(l, _) => Some(l),
233            _ => None,
234        }
235    }
236    pub fn into_left(self) -> Option<L> {
237        match self {
238            Self::Left(l) | Self::Both(l, _) => Some(l),
239            _ => None,
240        }
241    }
242    pub fn right(&self) -> Option<&R> {
243        match self {
244            Self::Right(r) | Self::Both(_, r) => Some(r),
245            _ => None,
246        }
247    }
248    pub fn into_right(self) -> Option<R> {
249        match self {
250            Self::Right(r) | Self::Both(_, r) => Some(r),
251            _ => None,
252        }
253    }
254}
255
256impl<L, R> MergeElement<&L, &R>
257where
258    L: Clone,
259    R: Clone,
260{
261    pub fn cloned(&self) -> MergeElement<L, R> {
262        match *self {
263            MergeElement::Both(a, b) => MergeElement::Both(a.clone(), b.clone()),
264            MergeElement::Left(a) => MergeElement::Left(a.clone()),
265            MergeElement::Right(b) => MergeElement::Right(b.clone()),
266        }
267    }
268}
269
270#[derive(Debug, PartialEq, Eq)]
271pub enum DiffElement<V> {
272    Unequal(V, V),
273    Left(V),
274    Right(V),
275}
276
277impl<T> DiffElement<T> {
278    pub fn new_data(self) -> Option<T> {
279        match self {
280            DiffElement::Left(_) => None,
281            DiffElement::Right(r) | DiffElement::Unequal(_, r) => Some(r),
282        }
283    }
284}
285
286#[test]
287fn test_merge_once() {
288    let i = [1i32, 2, 3][..].iter();
289    let j = [2i32, 4][..].iter();
290    let v: Vec<_> = MergeOnce::new(i, j).cloned().collect();
291    assert_eq!(v, vec![1, 2, 3, 4]);
292}
293
294pub(crate) trait SymmetricDiffMap<'a, K: 'a, V: 'a> {
295    type Iter: Iterator<Item = (&'a K, DiffElement<&'a V>)>;
296
297    fn symmetric_diff(&'a self, other: &'a Self) -> Self::Iter;
298
299    // Could be useful, I guess
300    #[allow(unused)]
301    fn symmetric_fold_with_inverse<R, FAdd, FRemove>(
302        &'a self,
303        other: &'a Self,
304        init: R,
305        mut add: FAdd,
306        mut remove: FRemove,
307    ) -> R
308    where
309        FAdd: FnMut(R, &K, &V) -> R,
310        FRemove: FnMut(R, &K, &V) -> R,
311        K: Ord,
312        V: PartialEq,
313    {
314        self.symmetric_diff(other)
315            .fold(init, |mut acc, (key, elem)| match elem {
316                DiffElement::Unequal(left, right) => {
317                    acc = add(acc, key, right);
318                    remove(acc, key, left)
319                }
320                DiffElement::Left(left) => remove(acc, key, left),
321                DiffElement::Right(right) => add(acc, key, right),
322            })
323    }
324}
325
326// Could be useful, I guess
327#[allow(unused)]
328pub(crate) trait SymmetricDiffMapOwned<K, V> {
329    type Iter: Iterator<Item = DiffElement<(K, V)>>;
330
331    fn symmetric_diff_owned(self, other: Self) -> Self::Iter;
332
333    fn symmetric_fold_owned<R, FAdd, FRemove>(
334        self,
335        other: Self,
336        init: R,
337        mut add: FAdd,
338        mut remove: FRemove,
339    ) -> R
340    where
341        FAdd: FnMut(R, K, V) -> R,
342        FRemove: FnMut(R, K, V) -> R,
343        K: Ord,
344        V: PartialEq,
345        Self: Sized,
346    {
347        self.symmetric_diff_owned(other)
348            .fold(init, |mut acc, elem| match elem {
349                DiffElement::Unequal(left, right) => {
350                    acc = add(acc, right.0, right.1);
351                    remove(acc, left.0, left.1)
352                }
353                DiffElement::Left(left) => remove(acc, left.0, left.1),
354                DiffElement::Right(right) => add(acc, right.0, right.1),
355            })
356    }
357}
358
359impl<'a, K: Ord + 'a, V: PartialEq + 'a> SymmetricDiffMap<'a, K, V> for BTreeMap<K, V> {
360    type Iter = SymmetricDiff<'a, K, V>;
361    fn symmetric_diff(&'a self, other: &'a Self) -> Self::Iter {
362        SymmetricDiff {
363            self_: self,
364            other,
365            keys: MergeOnce::new(self.keys(), other.keys()),
366        }
367    }
368}
369
370impl<K: Ord, V: PartialEq> SymmetricDiffMapOwned<K, V> for BTreeMap<K, V> {
371    type Iter = SymmetricDiffOwned<K, V>;
372    fn symmetric_diff_owned(self, other: Self) -> SymmetricDiffOwned<K, V> {
373        SymmetricDiffOwned {
374            self_: self.into_iter().peekable(),
375            other: other.into_iter().peekable(),
376            fused: None,
377        }
378    }
379}
380
381/// A trait implemented by `BTreeMap` and `im_rc::OrdMap`.
382pub trait GenericMap<K, V> {
383    fn remove(&mut self, key: &K) -> Option<V>;
384    fn insert(&mut self, key: K, value: V) -> Option<V>;
385}
386
387impl<K: Ord, V> GenericMap<K, V> for BTreeMap<K, V> {
388    #[inline]
389    fn remove(&mut self, key: &K) -> Option<V> {
390        BTreeMap::remove(self, key)
391    }
392
393    #[inline]
394    fn insert(&mut self, key: K, value: V) -> Option<V> {
395        BTreeMap::insert(self, key, value)
396    }
397}
398
399/// A trait implemented by `BTreeMap`, `im_rc::OrdMap` and `Rc<BTreeMap>`.
400///
401/// You frequently want to clone incrementals a lot. Making it so
402pub trait MutableMap<K, V> {
403    type UnderlyingMap: GenericMap<K, V>;
404    /// For Rc<BTreeMap>, this is BTreeMap.
405    fn make_mut(&mut self) -> &mut Self::UnderlyingMap;
406}
407
408pub trait SymmetricMapMap<K, V>: MutableMap<K, V> {
409    type OutputMap<V2: PartialEq + Clone>: SymmetricMapMap<K, V2>;
410
411    fn filter_map_collect<V2: PartialEq + Clone>(
412        &self,
413        f: &mut impl FnMut(&K, &V) -> Option<V2>,
414    ) -> Self::OutputMap<V2>;
415}
416
417pub trait SymmetricFoldMap<K, V> {
418    fn symmetric_fold<R>(
419        &self,
420        other: &Self,
421        init: R,
422        f: impl FnMut(R, (&K, DiffElement<&V>)) -> R,
423    ) -> R;
424    fn len(&self) -> usize;
425    fn is_empty(&self) -> bool {
426        self.len() == 0
427    }
428    /// Basically `self.iter().fold(init, f)`.
429    fn nonincremental_fold<R>(&self, init: R, f: impl FnMut(R, (&K, &V)) -> R) -> R;
430}
431
432impl<K: Ord + Clone, V: Clone> MutableMap<K, V> for Rc<BTreeMap<K, V>> {
433    type UnderlyingMap = BTreeMap<K, V>;
434    #[inline]
435    fn make_mut(&mut self) -> &mut Self::UnderlyingMap {
436        Rc::make_mut(self)
437    }
438}
439
440impl<K: Ord + Clone, V: PartialEq + Clone> SymmetricMapMap<K, V> for Rc<BTreeMap<K, V>> {
441    type OutputMap<V2: PartialEq + Clone> = Rc<BTreeMap<K, V2>>;
442    #[inline]
443    fn filter_map_collect<V2: PartialEq + Clone>(
444        &self,
445        f: &mut impl FnMut(&K, &V) -> Option<V2>,
446    ) -> Self::OutputMap<V2> {
447        Rc::new(self.deref().filter_map_collect(f))
448    }
449}
450
451impl<K: Ord, V: PartialEq> SymmetricFoldMap<K, V> for Rc<BTreeMap<K, V>> {
452    fn symmetric_fold<R>(
453        &self,
454        other: &Self,
455        init: R,
456        f: impl FnMut(R, (&K, DiffElement<&V>)) -> R,
457    ) -> R {
458        let self_target = self.deref();
459        let other_target = other.deref();
460        self_target.symmetric_diff(other_target).fold(init, f)
461    }
462    #[inline]
463    fn len(&self) -> usize {
464        self.deref().len()
465    }
466    fn nonincremental_fold<R>(&self, init: R, f: impl FnMut(R, (&K, &V)) -> R) -> R {
467        self.deref().nonincremental_fold(init, f)
468    }
469}
470
471impl<K: Ord, V> MutableMap<K, V> for BTreeMap<K, V> {
472    type UnderlyingMap = Self;
473    #[inline]
474    fn make_mut(&mut self) -> &mut Self::UnderlyingMap {
475        self
476    }
477}
478
479impl<K: Ord + Clone, V> SymmetricMapMap<K, V> for BTreeMap<K, V> {
480    type OutputMap<V2: PartialEq + Clone> = BTreeMap<K, V2>;
481    fn filter_map_collect<V2: PartialEq + Clone>(
482        &self,
483        f: &mut impl FnMut(&K, &V) -> Option<V2>,
484    ) -> Self::OutputMap<V2> {
485        self.iter()
486            .filter_map(|(k, v)| f(k, v).map(|v2| (k.clone(), v2)))
487            .collect()
488    }
489}
490impl<K: Ord, V: PartialEq> SymmetricFoldMap<K, V> for BTreeMap<K, V> {
491    fn symmetric_fold<R>(
492        &self,
493        other: &Self,
494        init: R,
495        f: impl FnMut(R, (&K, DiffElement<&V>)) -> R,
496    ) -> R {
497        self.symmetric_diff(other).fold(init, f)
498    }
499    #[inline]
500    fn len(&self) -> usize {
501        self.len()
502    }
503    #[inline]
504    fn nonincremental_fold<R>(&self, init: R, f: impl FnMut(R, (&K, &V)) -> R) -> R {
505        self.iter().fold(init, f)
506    }
507}