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

p3_util/
zip_eq.rs
