mysten_common/
zip_debug_eq.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::debug_fatal_at;
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_at!(
44                    &format!("{}:{}", self.caller.file(), self.caller.line()),
45                    "zip_debug_eq: first iterator shorter than second (created at {})",
46                    self.caller
47                );
48                None
49            }
50            (Some(_), None) => {
51                self.finished = true;
52                debug_fatal_at!(
53                    &format!("{}:{}", self.caller.file(), self.caller.line()),
54                    "zip_debug_eq: second iterator shorter than first (created at {})",
55                    self.caller
56                );
57                None
58            }
59        }
60    }
61
62    fn size_hint(&self) -> (usize, Option<usize>) {
63        if self.finished {
64            return (0, Some(0));
65        }
66        let (a_lower, a_upper) = self.a.size_hint();
67        let (b_lower, b_upper) = self.b.size_hint();
68        let lower = a_lower.min(b_lower);
69        let upper = match (a_upper, b_upper) {
70            (Some(a), Some(b)) => Some(a.min(b)),
71            (Some(a), None) => Some(a),
72            (None, Some(b)) => Some(b),
73            (None, None) => None,
74        };
75        (lower, upper)
76    }
77}
78
79/// Like `itertools::izip!`, but uses `zip_debug_eq` instead of `zip`.
80#[macro_export]
81macro_rules! izip_debug_eq {
82    // Closure helper for tuple flattening (same pattern as itertools::izip!)
83    ( @closure $p:pat => $tup:expr ) => {
84        |$p| $tup
85    };
86    ( @closure $p:pat => ( $($tup:tt)* ) , $_iter:expr $( , $tail:expr )* ) => {
87        $crate::izip_debug_eq!(@closure ($p, b) => ( $($tup)*, b ) $( , $tail )*)
88    };
89
90    // unary
91    ($first:expr $(,)*) => {
92        ::core::iter::IntoIterator::into_iter($first)
93    };
94
95    // binary
96    ($first:expr, $second:expr $(,)*) => {{
97        use $crate::ZipDebugEqIteratorExt as _;
98        $crate::izip_debug_eq!($first).zip_debug_eq($second)
99    }};
100
101    // n-ary where n > 2
102    ( $first:expr $( , $rest:expr )* $(,)* ) => {{
103        use $crate::ZipDebugEqIteratorExt as _;
104        $crate::izip_debug_eq!($first)
105            $(
106                .zip_debug_eq($rest)
107            )*
108            .map(
109                $crate::izip_debug_eq!(@closure a => (a) $( , $rest )*)
110            )
111    }};
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117
118    #[test]
119    fn equal_length_iterators() {
120        let a = vec![1, 2, 3];
121        let b = vec!["a", "b", "c"];
122        let result: Vec<_> = a.into_iter().zip_debug_eq(b).collect();
123        assert_eq!(result, vec![(1, "a"), (2, "b"), (3, "c")]);
124    }
125
126    #[test]
127    fn empty_iterators() {
128        let a: Vec<i32> = vec![];
129        let b: Vec<i32> = vec![];
130        let result: Vec<_> = a.into_iter().zip_debug_eq(b).collect();
131        assert_eq!(result, vec![]);
132    }
133
134    #[test]
135    #[should_panic]
136    fn first_shorter_panics_in_debug() {
137        let a = vec![1, 2];
138        let b = vec!["a", "b", "c"];
139        let _: Vec<_> = a.into_iter().zip_debug_eq(b).collect();
140    }
141
142    #[test]
143    #[should_panic]
144    fn second_shorter_panics_in_debug() {
145        let a = vec![1, 2, 3];
146        let b = vec!["a", "b"];
147        let _: Vec<_> = a.into_iter().zip_debug_eq(b).collect();
148    }
149
150    #[test]
151    fn izip_debug_eq_binary() {
152        let a = vec![1, 2, 3];
153        let b = vec!["a", "b", "c"];
154        let result: Vec<_> = izip_debug_eq!(a, b).collect();
155        assert_eq!(result, vec![(1, "a"), (2, "b"), (3, "c")]);
156    }
157
158    #[test]
159    fn izip_debug_eq_ternary() {
160        let a = vec![1, 2];
161        let b = vec!["a", "b"];
162        let c = vec![10.0, 20.0];
163        let result: Vec<_> = izip_debug_eq!(a, b, c).collect();
164        assert_eq!(result, vec![(1, "a", 10.0), (2, "b", 20.0)]);
165    }
166
167    #[test]
168    #[should_panic]
169    fn izip_debug_eq_mismatch_panics() {
170        let a = vec![1, 2, 3];
171        let b = vec!["a", "b"];
172        let c = vec![10.0, 20.0, 30.0];
173        let _: Vec<_> = izip_debug_eq!(a, b, c).collect();
174    }
175}
mysten_common/zip_debug_eq.rs

mysten_common/
zip_debug_eq.rs
