//! # Iterator extensions
//!
//! Zip two iterators, if they have different length, the shortest is extended
//! using a default value.
//!
//! The two iterators must have the same `Item` type.
//!
//! It is similar to `zip_longest` algorithm from the `itertool` python library.
//!
//! # Example
//! ```
//! use zip_fill::ZipFill;
//! let a = [ 1, 2, 3, 4];
//! let b = [ 1, 2];
//! let z: Vec<_> = a.iter().zip_longest(&b, &0).collect();
//! assert_eq!(vec!((&1, &1), (&2, &2), (&3, &0), (&4, &0)), z);
//! let z: Vec<_> = zip_fill::zip_longest_with(&a, &b, &0, |a, b| a + b).collect();
//! assert_eq!(vec!(2, 4, 3, 4), z);
//! ```

#![no_std]

#[cfg(test)]
extern crate std;

#[derive(Clone, Debug)]
pub struct ZipLongest<L, R>
where
    L: Iterator,
    R: Iterator,
{
    a: L,
    b: R,
    fill_value: L::Item,
}

/// Zip two iterators extending the shorter one with the provided `fill_value`.
///
/// # Arguments
///
/// * `a` - first iterator
/// * `b` - second iterator
/// * `fill_value` - default value
pub fn zip_longest<L, R>(a: L, b: R, fill_value: L::Item) -> ZipLongest<L::IntoIter, R::IntoIter>
where
    L::Item: Clone,
    L: IntoIterator,
    R: IntoIterator<Item = L::Item>,
{
    ZipLongest {
        a: a.into_iter(),
        b: b.into_iter(),
        fill_value,
    }
}

impl<L, R> Iterator for ZipLongest<L, R>
where
    L::Item: Clone,
    L: Iterator,
    R: Iterator<Item = L::Item>,
{
    type Item = (L::Item, L::Item);

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        match (self.a.next(), self.b.next()) {
            (Some(l), Some(r)) => Some((l, r)),
            (Some(l), None) => Some((l, self.fill_value.clone())),
            (None, Some(r)) => Some((self.fill_value.clone(), r)),
            _ => None,
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let a_hint = self.a.size_hint();
        let b_hint = self.b.size_hint();

        let lower_bound = core::cmp::max(a_hint.0, b_hint.0);
        let upper_bound = match (a_hint.1, b_hint.1) {
            (Some(a), Some(b)) => Some(core::cmp::max(a, b)),
            _ => None,
        };

        (lower_bound, upper_bound)
    }
}

/// Zip two iterators with the given function extending the shorter one
/// with the provided `fill_value`.
///
/// # Arguments
///
/// * `left` - first iterator
/// * `right` - second iterator
/// * `fill_value` - default value
/// * `f` - function used to zip the two iterators
pub fn zip_longest_with<L, R, T, F>(
    left: L,
    right: R,
    fill_value: L::Item,
    mut f: F,
) -> impl Iterator<Item = T>
where
    L::Item: Clone,
    L: IntoIterator,
    R: IntoIterator<Item = L::Item>,
    F: FnMut(L::Item, L::Item) -> T,
{
    zip_longest(left, right, fill_value).map(move |(l, r)| f(l, r))
}

/// Provide additional methods to `Iterator`
pub trait ZipFill: Iterator {
    /// Zip two iterators extending the shorter one with the provided `fill_value`.
    ///
    /// # Arguments
    ///
    /// * `other` - second iterator
    /// * `fill_value` - default value
    fn zip_longest<R>(self, other: R, fill_value: Self::Item) -> ZipLongest<Self, R::IntoIter>
    where
        Self::Item: Clone,
        R: IntoIterator<Item = Self::Item>,
        Self: Sized,
    {
        zip_longest(self, other, fill_value)
    }
}

impl<T: ?Sized> ZipFill for T where T: Iterator {}

#[cfg(test)]
mod tests {
    use super::*;
    use std::{vec, vec::Vec};

    #[test]
    fn zip_longest_left() {
        let a = zip_longest(1..=4, 6..=7, 0);
        assert_eq!(vec![(1, 6), (2, 7), (3, 0), (4, 0)], a.collect::<Vec<_>>());
    }

    #[test]
    fn zip_longest_right() {
        let a = zip_longest(['a', 'b'].iter(), ['a', 'b', 'c', 'd'].iter(), &'z');
        assert_eq!(
            vec![(&'a', &'a'), (&'b', &'b'), (&'z', &'c'), (&'z', &'d')],
            a.collect::<Vec<_>>()
        );
    }

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

    #[test]
    fn zip_longest_with_right() {
        let a = zip_longest_with(
            &[true, false],
            &[false, true, true, false],
            &true,
            |&x, &y| x && y,
        );
        assert_eq!(vec![false, false, true, false], a.collect::<Vec<_>>());
    }

    #[test]
    fn zip_longest_two_iters_copy() {
        // test zip_longest with iterator of different type
        let a = (1..6).map(|x| x * 2);
        let b = (1..7).filter(|x| x % 2 == 1);
        let res = zip_longest(a, b, 0);
        assert_eq!(
            vec![(2, 1), (4, 3), (6, 5), (8, 0), (10, 0)],
            res.collect::<Vec<_>>()
        );
    }

    #[allow(clippy::map_identity)]
    #[test]
    fn zip_longest_two_iters_non_copy() {
        // test zip_longest with iterator of different type non Copy
        let a = vec![vec![0], vec![1], vec![2]];
        let a = a.iter().map(|x| x);
        let b = vec![vec![5], vec![4]];
        let b = b.iter().filter(|_| true);
        let c = vec![10];
        let res = zip_longest(a, b, &c).map(|(x, y)| (x[0], y[0]));
        assert_eq!(vec![(0, 5), (1, 4), (2, 10)], res.collect::<Vec<_>>());
    }

    #[test]
    fn zip_longest_struct_left() {
        let mut a = zip_longest(&[1, 2, 3], &[1, 2], &0);
        assert_eq!(Some((&1, &1)), a.next());
        assert_eq!(Some((&2, &2)), a.next());
        assert_eq!(Some((&3, &0)), a.next());
        assert_eq!(None, a.next());
    }

    #[test]
    fn zip_longest_struct_collect() {
        let v1 = vec![1, 2, 3];
        let v2 = vec![1, 2];
        let a = zip_longest(&v1, &v2, &0);
        assert_eq!(vec![(&1, &1), (&2, &2), (&3, &0)], a.collect::<Vec<_>>());
    }

    #[test]
    fn zip_longest_trait() {
        let a: Vec<_> = (1..4).zip_longest(1..3, 0).collect();
        assert_eq!(vec![(1, 1), (2, 2), (3, 0)], a);
    }

    #[test]
    fn zip_longest_struct_right() {
        let mut a = zip_longest(&[true, false], &[false, true, false], &true);
        assert_eq!(Some((&true, &false)), a.next());
        assert_eq!(Some((&false, &true)), a.next());
        assert_eq!(Some((&true, &false)), a.next());
        assert_eq!(None, a.next());
    }

    #[test]
    fn zip_longest_with_new_left() {
        let mut a = zip_longest_with(&[1, 2, 3], &[1, 2], &0, |x, y| x * y);
        assert_eq!(Some(1), a.next());
        assert_eq!(Some(4), a.next());
        assert_eq!(Some(0), a.next());
        assert_eq!(None, a.next());
    }

    #[test]
    fn zip_longest_with_new_right() {
        let mut a = zip_longest_with(&[true, false], &[false, true, false], &true, |&x, &y| {
            x || y
        });
        assert_eq!(Some(true), a.next());
        assert_eq!(Some(true), a.next());
        assert_eq!(Some(true), a.next());
        assert_eq!(None, a.next());
    }

    #[test]
    fn zip_longest_with_collect() {
        let v1 = vec![1, 2, 3];
        let v2 = vec![1, 2];
        let a = zip_longest_with(&v1, &v2, &0, core::ops::Mul::mul);
        assert_eq!(vec![1, 4, 0], a.collect::<Vec<_>>());
    }

    #[test]
    fn size_hint() {
        let v1 = vec![1, 2, 3];
        let v2 = vec![1, 2];
        let a = zip_longest(&v1, &v2, &0);
        assert_eq!((3, Some(3)), a.size_hint());
        assert_eq!(v1.iter().size_hint(), a.size_hint());
    }
}