1use std::cmp::Ordering;
2use std::mem::take;
3use crate::vendor::rle::{HasLength, SplitableSpan};
4
5#[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 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#[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 Remainder::Nothing
70 }
71 Ordering::Less => {
72 let b_rem = b.truncate(a_len);
74 Remainder::SomeB(b_rem)
75 }
76 Ordering::Greater => {
77 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 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}