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

mysten_common/
zip_debug_eq.rs
