Skip to main content

braid_core/vendor/rle/
zip.rs

1use std::cmp::Ordering;
2use std::mem::take;
3use crate::vendor::rle::{HasLength, SplitableSpan};
4
5// Also used by intersect.
6#[derive(Clone, Debug)]
7pub(crate) enum Remainder<A, B> {
8    Nothing,
9    SomeA(A),
10    SomeB(B),
11}
12
13impl<A, B> Default for Remainder<A, B> {
14    fn default() -> Self { Remainder::Nothing }
15}
16
17impl<A, B> Remainder<A, B> {
18    pub(crate) fn take_from_iter<AI, BI>(&mut self, ai: &mut AI, bi: &mut BI) -> Option<(A, B)>
19        where AI: Iterator<Item=A>, BI: Iterator<Item=B>
20    {
21        Some(match take(self) {
22            Remainder::Nothing => {
23                // Fetch from both.
24                let a = ai.next()?;
25                let b = bi.next()?;
26                (a, b)
27            }
28            Remainder::SomeA(a) => {
29                let b = bi.next()?;
30                (a, b)
31            }
32            Remainder::SomeB(b) => {
33                let a = ai.next()?;
34                (a, b)
35            }
36        })
37    }
38}
39
40/// A RleZip is a zip iterator over 2 SplitableSpan iterators. Each item it yields contains the
41/// longest readable span from each of A and B.
42///
43/// The iterator ends at the min of A and B.
44#[derive(Clone, Debug)]
45pub struct RleZip<A, B>
46    where A: Iterator, B: Iterator,
47          A::Item: SplitableSpan + HasLength, B::Item: SplitableSpan + HasLength,
48{
49    rem: Remainder<A::Item, B::Item>,
50    a: A,
51    b: B,
52}
53
54impl<A, B> Iterator for RleZip<A, B>
55    where A: Iterator, B: Iterator,
56          A::Item: SplitableSpan + HasLength, B::Item: SplitableSpan + HasLength
57{
58    type Item = (A::Item, B::Item);
59
60    fn next(&mut self) -> Option<Self::Item> {
61        let (mut a, mut b) = self.rem.take_from_iter(&mut self.a, &mut self.b)?;
62
63        let a_len = a.len();
64        let b_len = b.len();
65
66        self.rem = match a_len.cmp(&b_len) {
67            Ordering::Equal => {
68                // Take all of both.
69                Remainder::Nothing
70            }
71            Ordering::Less => {
72                // a < b.
73                let b_rem = b.truncate(a_len);
74                Remainder::SomeB(b_rem)
75            }
76            Ordering::Greater => {
77                // a > b.
78                let a_rem = a.truncate(b_len);
79                Remainder::SomeA(a_rem)
80            }
81        };
82
83        Some((a, b))
84    }
85}
86
87pub fn rle_zip<A, B>(a: A, b: B) -> RleZip<A, B>
88    where A: Iterator, B: Iterator,
89          A::Item: SplitableSpan + HasLength, B::Item: SplitableSpan + HasLength
90{
91    RleZip {
92        rem: Remainder::Nothing,
93        a,
94        b
95    }
96}
97
98pub fn rle_zip3<A, B, C>(a: A, b: B, c: C) -> impl Iterator<Item=(A::Item, B::Item, C::Item)>
99    where A: Iterator, B: Iterator, C: Iterator,
100          A::Item: SplitableSpan + HasLength,
101          B::Item: SplitableSpan + HasLength,
102          C::Item: SplitableSpan + HasLength,
103{
104    rle_zip(rle_zip(a, b), c).map(|((a, b), c)| (a, b, c))
105}
106
107#[cfg(test)]
108mod test {
109    use super::*;
110    use crate::vendor::rle::RleRun;
111
112    fn check_zip(a: &[RleRun<u32>], b: &[RleRun<u32>], expect: &[(RleRun<u32>, RleRun<u32>)]) {
113        assert_eq!(rle_zip(a.iter().copied(), b.iter().copied())
114                       .collect::<Vec<_>>(), expect);
115
116        // And check that if we swap the parameter order we get the same thing.
117        assert_eq!(rle_zip(b.iter().copied(), a.iter().copied())
118                       .map(|(a, b)| (b, a))
119                       .collect::<Vec<_>>(), expect);
120    }
121
122    #[test]
123    fn smoke() {
124        let one = vec![
125            RleRun { val: 1, len: 1 },
126            RleRun { val: 2, len: 4 }
127        ];
128        let two = vec![
129            RleRun { val: 11, len: 4 },
130            RleRun { val: 12, len: 1 }
131        ];
132
133        let expected = vec![
134            (RleRun { val: 1, len: 1 }, RleRun { val: 11, len: 1}),
135            (RleRun { val: 2, len: 3 }, RleRun { val: 11, len: 3}),
136            (RleRun { val: 2, len: 1 }, RleRun { val: 12, len: 1}),
137        ];
138
139        check_zip(&one, &two, &expected);
140    }
141
142    #[test]
143    fn one_is_longer() {
144        let one = vec![
145            RleRun { val: 1, len: 100 },
146        ];
147        let two = vec![
148            RleRun { val: 11, len: 10 },
149        ];
150
151        let expected = vec![
152            (RleRun { val: 1, len: 10 }, RleRun { val: 11, len: 10}),
153        ];
154
155        check_zip(&one, &two, &expected);
156    }
157}