mysten_common/
zip_debug_eq.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::debug_fatal;
5use std::panic::Location;
6
7/// Extension trait providing `zip_debug_eq` on all iterators.
8pub trait ZipDebugEqIteratorExt: Iterator + Sized {
9    /// Like `zip_eq`, but instead of panicking logs `debug_fatal!` if the iterators have
10    /// different lengths. After the diagnostic, behaves like `zip` (returns `None`).
11    #[track_caller]
12    fn zip_debug_eq<J: IntoIterator>(self, other: J) -> ZipDebugEq<Self, J::IntoIter> {
13        ZipDebugEq {
14            a: self,
15            b: other.into_iter(),
16            finished: false,
17            caller: Location::caller(),
18        }
19    }
20}
21
22impl<I: Iterator> ZipDebugEqIteratorExt for I {}
23
24pub struct ZipDebugEq<A, B> {
25    a: A,
26    b: B,
27    finished: bool,
28    caller: &'static Location<'static>,
29}
30
31impl<A: Iterator, B: Iterator> Iterator for ZipDebugEq<A, B> {
32    type Item = (A::Item, B::Item);
33
34    fn next(&mut self) -> Option<Self::Item> {
35        if self.finished {
36            return None;
37        }
38        match (self.a.next(), self.b.next()) {
39            (Some(a), Some(b)) => Some((a, b)),
40            (None, None) => None,
41            (None, Some(_)) => {
42                self.finished = true;
43                debug_fatal!(
44                    "zip_debug_eq: first iterator shorter than second (created at {})",
45                    self.caller
46                );
47                None
48            }
49            (Some(_), None) => {
50                self.finished = true;
51                debug_fatal!(
52                    "zip_debug_eq: second iterator shorter than first (created at {})",
53                    self.caller
54                );
55                None
56            }
57        }
58    }
59
60    fn size_hint(&self) -> (usize, Option<usize>) {
61        if self.finished {
62            return (0, Some(0));
63        }
64        let (a_lower, a_upper) = self.a.size_hint();
65        let (b_lower, b_upper) = self.b.size_hint();
66        let lower = a_lower.min(b_lower);
67        let upper = match (a_upper, b_upper) {
68            (Some(a), Some(b)) => Some(a.min(b)),
69            (Some(a), None) => Some(a),
70            (None, Some(b)) => Some(b),
71            (None, None) => None,
72        };
73        (lower, upper)
74    }
75}
76
77/// Like `itertools::izip!`, but uses `zip_debug_eq` instead of `zip`.
78#[macro_export]
79macro_rules! izip_debug_eq {
80    // Closure helper for tuple flattening (same pattern as itertools::izip!)
81    ( @closure $p:pat => $tup:expr ) => {
82        |$p| $tup
83    };
84    ( @closure $p:pat => ( $($tup:tt)* ) , $_iter:expr $( , $tail:expr )* ) => {
85        $crate::izip_debug_eq!(@closure ($p, b) => ( $($tup)*, b ) $( , $tail )*)
86    };
87
88    // unary
89    ($first:expr $(,)*) => {
90        ::core::iter::IntoIterator::into_iter($first)
91    };
92
93    // binary
94    ($first:expr, $second:expr $(,)*) => {{
95        use $crate::ZipDebugEqIteratorExt as _;
96        $crate::izip_debug_eq!($first).zip_debug_eq($second)
97    }};
98
99    // n-ary where n > 2
100    ( $first:expr $( , $rest:expr )* $(,)* ) => {{
101        use $crate::ZipDebugEqIteratorExt as _;
102        $crate::izip_debug_eq!($first)
103            $(
104                .zip_debug_eq($rest)
105            )*
106            .map(
107                $crate::izip_debug_eq!(@closure a => (a) $( , $rest )*)
108            )
109    }};
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115
116    #[test]
117    fn equal_length_iterators() {
118        let a = vec![1, 2, 3];
119        let b = vec!["a", "b", "c"];
120        let result: Vec<_> = a.into_iter().zip_debug_eq(b).collect();
121        assert_eq!(result, vec![(1, "a"), (2, "b"), (3, "c")]);
122    }
123
124    #[test]
125    fn empty_iterators() {
126        let a: Vec<i32> = vec![];
127        let b: Vec<i32> = vec![];
128        let result: Vec<_> = a.into_iter().zip_debug_eq(b).collect();
129        assert_eq!(result, vec![]);
130    }
131
132    #[test]
133    #[should_panic]
134    fn first_shorter_panics_in_debug() {
135        let a = vec![1, 2];
136        let b = vec!["a", "b", "c"];
137        let _: Vec<_> = a.into_iter().zip_debug_eq(b).collect();
138    }
139
140    #[test]
141    #[should_panic]
142    fn second_shorter_panics_in_debug() {
143        let a = vec![1, 2, 3];
144        let b = vec!["a", "b"];
145        let _: Vec<_> = a.into_iter().zip_debug_eq(b).collect();
146    }
147
148    #[test]
149    fn izip_debug_eq_binary() {
150        let a = vec![1, 2, 3];
151        let b = vec!["a", "b", "c"];
152        let result: Vec<_> = izip_debug_eq!(a, b).collect();
153        assert_eq!(result, vec![(1, "a"), (2, "b"), (3, "c")]);
154    }
155
156    #[test]
157    fn izip_debug_eq_ternary() {
158        let a = vec![1, 2];
159        let b = vec!["a", "b"];
160        let c = vec![10.0, 20.0];
161        let result: Vec<_> = izip_debug_eq!(a, b, c).collect();
162        assert_eq!(result, vec![(1, "a", 10.0), (2, "b", 20.0)]);
163    }
164
165    #[test]
166    #[should_panic]
167    fn izip_debug_eq_mismatch_panics() {
168        let a = vec![1, 2, 3];
169        let b = vec!["a", "b"];
170        let c = vec![10.0, 20.0, 30.0];
171        let _: Vec<_> = izip_debug_eq!(a, b, c).collect();
172    }
173}
mysten_common/zip_debug_eq.rs

mysten_common/
zip_debug_eq.rs
