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]
/// 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 crate::{izip_eq, izip_eq_lazy, 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)));
    }
}