p3_util/
zip_eq.rs

1/// An iterator which iterates two other iterators of the same length simultaneously.
2///
3/// Equality of the lengths of `a` abd `b` are at construction time.
4#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
5pub struct ZipEq<A, B> {
6    a: A,
7    b: B,
8}
9
10/// Zips two iterators but **panics** if they are not of the same length.
11///
12/// Similar to `itertools::zip_eq`, but we check the lengths at construction time.
13pub fn zip_eq<A, AIter, B, BIter, Error>(
14    a: A,
15    b: B,
16    err: Error,
17) -> Result<ZipEq<A::IntoIter, B::IntoIter>, Error>
18where
19    A: IntoIterator<IntoIter = AIter>,
20    AIter: ExactSizeIterator,
21    B: IntoIterator<IntoIter = BIter>,
22    BIter: ExactSizeIterator,
23{
24    let a_iter = a.into_iter();
25    let b_iter = b.into_iter();
26    match a_iter.len() == b_iter.len() {
27        true => Ok(ZipEq {
28            a: a_iter,
29            b: b_iter,
30        }),
31        false => Err(err),
32    }
33}
34
35impl<A, B> Iterator for ZipEq<A, B>
36where
37    A: ExactSizeIterator, // We need to use ExactSizeIterator here otherwise the size_hint() methods could differ.
38    B: ExactSizeIterator,
39{
40    type Item = (A::Item, B::Item);
41
42    fn next(&mut self) -> Option<Self::Item> {
43        match (self.a.next(), self.b.next()) {
44            (Some(a), Some(b)) => Some((a, b)),
45            (None, None) => None,
46            _ => unreachable!("The iterators must have the same length."),
47        }
48    }
49
50    fn size_hint(&self) -> (usize, Option<usize>) {
51        // self.a.size_hint() = self.b.size_hint() as a and b are ExactSizeIterators
52        // and we checked that they are the same length at construction time.
53        debug_assert_eq!(self.a.size_hint(), self.b.size_hint());
54        self.a.size_hint()
55    }
56}
57
58impl<A, B> ExactSizeIterator for ZipEq<A, B>
59where
60    A: ExactSizeIterator,
61    B: ExactSizeIterator,
62{
63}
64
65#[cfg(test)]
66mod tests {
67    use alloc::vec;
68    use alloc::vec::Vec;
69
70    use super::*;
71
72    #[test]
73    fn test_zip_eq_success() {
74        let a = [1, 2, 3];
75        let b = ['a', 'b', 'c'];
76
77        // Expect zip_eq to succeed since both slices are length 3.
78        let zipped = zip_eq(a, b, "length mismatch").unwrap();
79
80        let result: Vec<_> = zipped.collect();
81
82        // Expect tuples zipped together positionally.
83        assert_eq!(result, vec![(1, 'a'), (2, 'b'), (3, 'c')]);
84    }
85
86    #[test]
87    fn test_zip_eq_length_mismatch() {
88        let a = [1, 2];
89        let b = ['x', 'y', 'z'];
90
91        // Use pattern matching instead of .unwrap_err()
92        match zip_eq(a, b, "oops") {
93            Err(e) => assert_eq!(e, "oops"),
94            Ok(_) => panic!("expected error due to mismatched lengths"),
95        }
96    }
97
98    #[test]
99    fn test_zip_eq_empty_iterators() {
100        let a: [i32; 0] = [];
101        let b: [char; 0] = [];
102
103        // Zipping two empty iterators should succeed and produce an empty iterator.
104        let zipped = zip_eq(a, b, "mismatch").unwrap();
105
106        let result: Vec<_> = zipped.collect();
107
108        // The result should be an empty vector.
109        assert!(result.is_empty());
110    }
111
112    #[test]
113    fn test_zip_eq_size_hint() {
114        let a = [10, 20];
115        let b = [100, 200];
116
117        let zipped = zip_eq(a, b, "bad").unwrap();
118
119        // Size hint should reflect the number of items remaining.
120        assert_eq!(zipped.size_hint(), (2, Some(2)));
121    }
122
123    #[test]
124    fn test_zip_eq_unreachable_case() {
125        let a = [1, 2];
126        let b = [3, 4];
127
128        let mut zipped = zip_eq(a, b, "fail").unwrap();
129
130        // Manually advance past the last element
131        assert_eq!(zipped.next(), Some((1, 3)));
132        assert_eq!(zipped.next(), Some((2, 4)));
133        assert_eq!(zipped.next(), None);
134    }
135}
p3_util/zip_eq.rs

p3_util/
zip_eq.rs
