primitives/utils/

iter.rs

#[macro_export]
/// Create an iterator running multiple iterators in lockstep, panicking if they
/// are not of the same length.
///
/// Length checks are done lazily as the resulting iterator is built.
/// I.e. the `izip_eq_lazy!` iterator yields elements until any subiterator
/// returns `None`.
///
/// **Note:** This macro is slower than `izip_eq!` however it allows to zip iterators of
/// unknown or unbounded size.
macro_rules! izip_eq_lazy {
    // @closure creates a tuple-flattening closure for .map() call. usage:
    // @closure partial_pattern => partial_tuple , rest , of , iterators
    // eg. izip_eq_lazy!( @closure ((a, b), c) => (a, b, c) , dd , ee )
    ( @closure $p:pat => $tup:expr ) => {
        |$p| $tup
    };

    // The "b" identifier is a different identifier on each recursion level thanks to hygiene.
    ( @closure $p:pat => ( $($tup:tt)* ) , $_iter:expr $( , $tail:expr )* ) => {
        $crate::izip_eq_lazy!(@closure ($p, b) => ( $($tup)*, b ) $( , $tail )*)
    };

    // unary
    ($first:expr $(,)*) => {
        std::iter::IntoIterator::into_iter($first)
    };

    // binary
    ($first:expr, $second:expr $(,)*) => {
        itertools::zip_eq(
            std::iter::IntoIterator::into_iter($first),
            $second,
        )
    };

    // n-ary where n > 2
    ( $first:expr $( , $rest:expr )* $(,)* ) => {
        {
            let iter = std::iter::IntoIterator::into_iter($first);
            $(
                let iter = itertools::zip_eq(iter, $rest);
            )*
            std::iter::Iterator::map(
                iter,
                $crate::izip_eq_lazy!(@closure a => (a) $( , $rest )*)
            )
        }
    };
}

#[macro_export]
/// Create an iterator running multiple iterators in lockstep, panicking if they
/// are not the same length.
///
/// Length checks are eagerly performed before the resulting iterator is built.
///
/// **Note:** For performance reasons, this macro is preferable to `izip_eq_lazy!` when the lengths
/// of the iterators are known.
macro_rules! izip_eq {
    (@assert_eq_len $first:expr, $second:expr) => {
        assert!($first.len() == $second.len(), "iterator length mismatch: {} vs {}", $first.len(), $second.len());
    };

    // unary
    ($first:expr $(,)*) => {
        std::iter::IntoIterator::into_iter($first)
    };

    // binary
    ($first:expr, $second:expr $(,)*) => {
        {
            let iter = std::iter::IntoIterator::into_iter($first);
            let second = std::iter::IntoIterator::into_iter($second);
            $crate::izip_eq!(@assert_eq_len iter, second);
            let iter = std::iter::Iterator::zip(iter, second);
            iter
        }
    };

    // n-ary where n > 2
    ( $first:expr $( , $rest:expr )* $(,)* ) => {
        {
            let iter = std::iter::IntoIterator::into_iter($first);
            $(
                let rest = std::iter::IntoIterator::into_iter($rest);
                $crate::izip_eq!(@assert_eq_len iter, rest);
                let iter = std::iter::Iterator::zip(iter, rest);
            )*
            std::iter::Iterator::map(
                iter,
                $crate::izip_eq_lazy!(@closure a => (a) $( , $rest )*)
            )
        }
    };
}

#[macro_export]
/// Create a parallel iterator running multiple iterators in lockstep, panicking if they
/// are not the same length.
///
/// This is the rayon parallel version of `izip_eq!`. It uses `into_par_iter()` and performs
/// length checks eagerly before building the parallel iterator.
///
/// **Note:** All input iterators must implement `IntoParallelIterator` and the resulting
/// parallel iterator will process elements in parallel using rayon.
///
/// # Examples
///
/// ```
/// use primitives::par_izip;
/// use rayon::prelude::*;
///
/// let a = vec![1, 2, 3];
/// let b = vec![4, 5, 6];
/// let c = vec![7, 8, 9];
///
/// let results: Vec<_> = par_izip!(&a, &b, &c).map(|(a, b, c)| a + b + c).collect();
///
/// assert_eq!(results, vec![12, 15, 18]);
/// ```
macro_rules! par_izip {
    // Helper to create closure for tuple flattening, reuses izip_eq_lazy's closure builder
    ( @closure $p:pat => $tup:expr ) => {
        |$p| $tup
    };

    ( @closure $p:pat => ( $($tup:tt)* ) , $_iter:expr $( , $tail:expr )* ) => {
        $crate::par_izip!(@closure ($p, b) => ( $($tup)*, b ) $( , $tail )*)
    };

    // Length assertion helper
    (@assert_eq_len $first:expr, $second:expr) => {
        assert!($first.len() == $second.len(), "parallel iterator length mismatch: {} vs {}", $first.len(), $second.len());
    };

    // unary
    ($first:expr $(,)*) => {
        rayon::iter::IntoParallelIterator::into_par_iter($first)
    };

    // binary
    ($first:expr, $second:expr $(,)*) => {
        {
            let iter = rayon::iter::IntoParallelIterator::into_par_iter($first);
            let second = rayon::iter::IntoParallelIterator::into_par_iter($second);
            $crate::par_izip!(@assert_eq_len iter, second);
            let iter = rayon::iter::IndexedParallelIterator::zip(iter, second);
            iter
        }
    };

    // n-ary where n > 2
    ( $first:expr $( , $rest:expr )* $(,)* ) => {
        {
            let iter = rayon::iter::IntoParallelIterator::into_par_iter($first);
            $(
                let rest = rayon::iter::IntoParallelIterator::into_par_iter($rest);
                $crate::par_izip!(@assert_eq_len iter, rest);
                let iter = rayon::iter::IndexedParallelIterator::zip(iter, rest);
            )*
            rayon::iter::ParallelIterator::map(
                iter,
                $crate::par_izip!(@closure a => (a) $( , $rest )*)
            )
        }
    };
}

#[macro_export]
/// An adaptor for `itertools::chain!` macro which creates an ExactSizeIterator.
///
/// **Note:** This macro collects the output of `itertools::chain!` into a vector. This may prevent
/// some compiler optimizations implying iterators.
macro_rules! chain_eq {
    ($first:expr $( , $rest:expr )* $(,)*) => {{
        let it = itertools::chain!($first, $($rest, )*);
        let vec = <std::vec::Vec<_> as std::iter::FromIterator<_>>::from_iter(it);
        <std::vec::Vec<_> as std::iter::IntoIterator>::into_iter(vec)
    }
    };
}

#[macro_export]
/// Create an iterator over hashmaps with the same keys. Accepts a list of keys, and the maps
/// to be iterated over. Panics if any of the keys are not found in the maps.
macro_rules! zip_maps {
    ($keys:expr, $($map:expr),+ $(,)?) => {
        $keys.into_iter().map(move |key| {
            (
                key,
                $(
                    $map.remove(&key).expect(&format!(
                        "Key `{:?}` not found in map `{}`",
                        key,
                        stringify!($map),
                    )),
                )+
            )
        })
    };
}

pub trait IntoExactSizeIterator:
    IntoIterator<IntoIter: ExactSizeIterator<Item = <Self as IntoIterator>::Item>>
{
}

impl<T: IntoIterator<IntoIter = S>, S: ExactSizeIterator<Item = <T as IntoIterator>::Item>>
    IntoExactSizeIterator for T
{
}

/// A trait for iterators that take exactly N elements, panicking if the iterator
/// is shorter than N.
pub trait TakeExact: Iterator {
    /// Takes exactly `n` elements from the iterator, panicking if the iterator
    /// is shorter than `n`. Returns an iterator that implements `ExactSizeIterator`.
    fn take_exact(self, n: usize) -> TakeExactIter<Self>
    where
        Self: Sized,
    {
        TakeExactIter {
            iter: self,
            remaining: n,
        }
    }
}

impl<I: Iterator> TakeExact for I {}

/// An iterator that takes exactly N elements from the underlying iterator.
#[derive(Clone, Debug)]
pub struct TakeExactIter<I> {
    iter: I,
    remaining: usize,
}

impl<I: Iterator> Iterator for TakeExactIter<I> {
    type Item = I::Item;

    fn next(&mut self) -> Option<Self::Item> {
        if self.remaining == 0 {
            return None;
        }
        self.remaining -= 1;
        match self.iter.next() {
            Some(item) => Some(item),
            None => panic!("iterator shorter than expected length"),
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.remaining, Some(self.remaining))
    }
}

impl<I: Iterator> ExactSizeIterator for TakeExactIter<I> {}

#[cfg(test)]
mod tests {
    use rayon::prelude::*;

    use crate::{izip_eq, izip_eq_lazy, par_izip, utils::TakeExact};

    #[test]
    fn test_izip_eq() {
        let a = [1, 2, 3];
        let b = [4, 5, 6];
        let c = [7, 8, 9];

        {
            let mut results = [0, 0, 0];
            for (r, aa, bb, cc) in izip_eq_lazy!(&mut results, &a, &b, &c) {
                *r = aa + bb + cc;
            }

            assert_eq!(results, [1 + 4 + 7, 2 + 5 + 8, 3 + 6 + 9]);
        }
        {
            let mut results = [0, 0, 0];
            for (r, aa, bb) in izip_eq_lazy!(&mut results, &a, &b) {
                *r = aa + bb;
            }

            assert_eq!(results, [1 + 4, 2 + 5, 3 + 6]);
        }
        {
            let mut results = [0, 0, 0];
            for (r, aa) in izip_eq_lazy!(&mut results, &a) {
                *r = *aa;
            }

            assert_eq!(results, [1, 2, 3]);
        }
        {
            let mut result = 0;
            for aa in izip_eq_lazy!(&a) {
                result += *aa;
            }

            assert_eq!(result, 1 + 2 + 3);
        }
        {
            let mut results = [0, 0, 0];
            for (r, aa, bb, cc) in izip_eq!(&mut results, &a, &b, &c) {
                *r = aa + bb + cc;
            }

            assert_eq!(results, [1 + 4 + 7, 2 + 5 + 8, 3 + 6 + 9]);
        }
        {
            let mut results = [0, 0, 0];
            for (r, aa, bb) in izip_eq!(&mut results, &a, &b) {
                *r = aa + bb;
            }

            assert_eq!(results, [1 + 4, 2 + 5, 3 + 6]);
        }
        {
            let mut results = [0, 0, 0];
            for (r, aa) in izip_eq!(&mut results, &a) {
                *r = *aa;
            }

            assert_eq!(results, [1, 2, 3]);
        }
        {
            let mut result = 0;
            for aa in izip_eq!(&a) {
                result += *aa;
            }

            assert_eq!(result, 1 + 2 + 3);
        }
    }

    #[test]
    #[should_panic(expected = "itertools: .zip_eq() reached end of one iterator before the other")]
    fn test_izip_eq_lazy_panic() {
        let a = [1, 2, 3];
        let b = [4, 5];
        let c = [7, 8, 9];

        let mut results = [0, 0, 0];
        for (r, aa, bb, cc) in izip_eq_lazy!(&mut results, &a, &b, &c) {
            *r = aa + bb + cc;
        }
        unreachable!()
    }

    #[test]
    #[should_panic(expected = "itertools: .zip_eq() reached end of one iterator before the other")]
    fn test_izip_eq_lazy_panic_2() {
        let a = [1, 2, 3];
        let b = [4, 5, 6];
        let c = [7, 8, 9];

        let mut results = [0, 0];
        for (r, aa, bb, cc) in izip_eq_lazy!(&mut results, &a, &b, &c) {
            *r = aa + bb + cc;
        }
        unreachable!()
    }

    #[test]
    #[should_panic(expected = "iterator length mismatch: 3 vs 2")]
    fn test_izip_eq_eager_panic() {
        let a = [1, 2, 3];
        let b = [4, 5];
        let c = [7, 8, 9];

        let mut results = [0, 0, 0];
        for (r, aa, bb, cc) in izip_eq!(&mut results, &a, &b, &c) {
            *r = aa + bb + cc;
        }
        unreachable!()
    }

    #[test]
    #[should_panic(expected = "iterator length mismatch: 2 vs 3")]
    fn test_izip_eq_eager_panic_2() {
        let a = [1, 2, 3];
        let b = [4, 5, 6];
        let c = [7, 8, 9];

        let mut results = [0, 0];
        for (r, aa, bb, cc) in izip_eq!(&mut results, &a, &b, &c) {
            *r = aa + bb + cc;
        }
        unreachable!()
    }

    #[test]
    #[should_panic(expected = "iterator length mismatch: 2 vs 3")]
    fn test_izip_eq_eager_panic_3() {
        let a = [1, 2, 3];

        let mut results = [0, 0];
        for (r, aa) in izip_eq!(&mut results, &a) {
            *r = *aa;
        }
        unreachable!()
    }

    #[test]
    #[should_panic(expected = "iterator length mismatch: 3 vs 2")]
    fn test_izip_eq_eager_panic_4() {
        let a = [1, 2];

        let mut results = [0, 0, 0];
        for (r, aa) in izip_eq!(&mut results, &a) {
            *r = *aa;
        }
        unreachable!()
    }

    #[test]
    fn test_zip_maps() {
        let keys = vec![1, 2, 3];
        let mut map1: std::collections::HashMap<_, _> = [(1, "a"), (2, "b"), (3, "c")].into();
        let mut map2: std::collections::HashMap<_, _> = [(1, "x"), (2, "y"), (3, "z")].into();

        let result: Vec<_> = zip_maps!(keys, map1, map2).collect();

        assert_eq!(result, vec![(1, "a", "x"), (2, "b", "y"), (3, "c", "z"),]);
    }

    #[test]
    #[should_panic(expected = "Key `4` not found in map `map1`")]
    fn test_zip_maps_missing_key_panic() {
        let keys = vec![1, 2, 3, 4];
        let mut map1: std::collections::HashMap<_, _> = [(1, "a"), (2, "b"), (3, "c")].into();
        let mut map2: std::collections::HashMap<_, _> = [(1, "x"), (2, "y"), (3, "z")].into();

        let _result: Vec<_> = zip_maps!(keys, map1, map2).collect();
    }

    #[test]
    fn test_take_exact() {
        let v = vec![1, 2, 3, 4, 5];
        let mut iter = v.into_iter().take_exact(3);

        assert_eq!(iter.len(), 3);
        assert_eq!(iter.next(), Some(1));
        assert_eq!(iter.len(), 2);
        assert_eq!(iter.next(), Some(2));
        assert_eq!(iter.len(), 1);
        assert_eq!(iter.next(), Some(3));
        assert_eq!(iter.len(), 0);
        assert_eq!(iter.next(), None);
    }

    #[test]
    #[should_panic(expected = "iterator shorter than expected length")]
    fn test_take_exact_panic() {
        let v = vec![1, 2];
        let mut iter = v.into_iter().take_exact(3);

        assert_eq!(iter.next(), Some(1));
        assert_eq!(iter.next(), Some(2));
        iter.next(); // This should panic
    }

    #[test]
    fn test_take_exact_collect() {
        let v = vec![1, 2, 3, 4, 5];
        let result: Vec<_> = v.into_iter().take_exact(3).collect();
        assert_eq!(result, vec![1, 2, 3]);
    }

    #[test]
    fn test_take_exact_size_hint() {
        let v = vec![1, 2, 3, 4, 5];
        let iter = v.into_iter().take_exact(3);
        assert_eq!(iter.size_hint(), (3, Some(3)));
    }

    #[test]
    fn test_par_izip_unary() {
        let a = vec![1, 2, 3];
        let result: Vec<_> = par_izip!(&a).map(|x| x * 2).collect();
        assert_eq!(result, vec![2, 4, 6]);
    }

    #[test]
    fn test_par_izip_binary() {
        let a = vec![1, 2, 3];
        let b = vec![4, 5, 6];
        let result: Vec<_> = par_izip!(&a, &b).map(|(x, y)| x + y).collect();
        assert_eq!(result, vec![5, 7, 9]);
    }

    #[test]
    fn test_par_izip_ternary() {
        let a = vec![1, 2, 3];
        let b = vec![4, 5, 6];
        let c = vec![7, 8, 9];
        let result: Vec<_> = par_izip!(&a, &b, &c).map(|(x, y, z)| x + y + z).collect();
        assert_eq!(result, vec![12, 15, 18]);
    }

    #[test]
    fn test_par_izip_quaternary() {
        let a = vec![1, 2, 3];
        let b = vec![4, 5, 6];
        let c = vec![7, 8, 9];
        let d = vec![10, 11, 12];
        let result: Vec<_> = par_izip!(&a, &b, &c, &d)
            .map(|(w, x, y, z)| w + x + y + z)
            .collect();
        assert_eq!(result, vec![22, 26, 30]);
    }

    #[test]
    fn test_par_izip_with_mutation() {
        let a = vec![1, 2, 3];
        let b = vec![4, 5, 6];
        let mut results = vec![0, 0, 0];

        par_izip!(&mut results, &a, &b).for_each(|(r, a, b)| {
            *r = a + b;
        });

        assert_eq!(results, vec![5, 7, 9]);
    }

    #[test]
    #[should_panic(expected = "parallel iterator length mismatch: 3 vs 2")]
    fn test_par_izip_length_mismatch_binary() {
        let a = vec![1, 2, 3];
        let b = vec![4, 5];
        let _: Vec<_> = par_izip!(&a, &b).collect();
    }

    #[test]
    #[should_panic(expected = "parallel iterator length mismatch: 3 vs 2")]
    fn test_par_izip_length_mismatch_ternary_first() {
        let a = vec![1, 2, 3];
        let b = vec![4, 5];
        let c = vec![7, 8, 9];
        let _: Vec<_> = par_izip!(&a, &b, &c).collect();
    }

    #[test]
    #[should_panic(expected = "parallel iterator length mismatch: 3 vs 2")]
    fn test_par_izip_length_mismatch_ternary_second() {
        let a = vec![1, 2, 3];
        let b = vec![4, 5, 6];
        let c = vec![7, 8];
        let _: Vec<_> = par_izip!(&a, &b, &c).collect();
    }

    #[test]
    fn test_par_izip_empty() {
        let a: Vec<i32> = vec![];
        let b: Vec<i32> = vec![];
        let result: Vec<_> = par_izip!(&a, &b).collect();
        assert_eq!(result, vec![]);
    }

    #[test]
    fn test_par_izip_large() {
        // Test with larger dataset to ensure parallel execution works
        let a: Vec<_> = (0..1000).collect();
        let b: Vec<_> = (1000..2000).collect();
        let c: Vec<_> = (2000..3000).collect();

        let result: Vec<_> = par_izip!(&a, &b, &c).map(|(x, y, z)| x + y + z).collect();

        assert_eq!(result.len(), 1000);
        assert_eq!(result[0], 1000 + 2000);
        assert_eq!(result[999], 999 + 1999 + 2999);
    }
}