/*
 * Copyright (c) 2019 Frank Fischer <frank-fischer@shadow-soft.de>
 *
 * This program is free software: you can redistribute it and/or
 * modify it under the terms of the GNU General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see  <http://www.gnu.org/licenses/>
 */

//! Compact iterator comprehensions.
//!
//! The macro `iter!` can be used to generate a sequence of elements using
//! generating sequences and conditional filters. There are two general forms
//!
//! ```nodoc
//! iter!(EXPR; VAR in RANGE, COND, ...)
//! iter!(VAR in RANGE, COND, ..., EXPR)
//! ```
//!
//! where `EXPR` the generator expression, `VAR` is a variable `RANGE` a some
//! sequence, e.g. `0..n`. `COND` is an arbitrary boolean expression. Each
//! `RANGE` and `COND` term can use the variables declared in preceeding range
//! expressions.
//!
//! # Example
//! The expression $\\{ 5i + j : i \\in \\{0..4\\}, j \\in \\{0..4\\}, i < j \\}$ is
//! equivalent to the following form
//!
//! ```
//! use iter_comprehensions::iter;
//! assert_eq!(iter!(4*i + j; i in 0..5, j in 0..5, i < j).collect::<Vec<_>>(),
//!            vec![1, 2, 3, 4, 6, 7, 8, 11, 12, 16]);
//! ```

use std::iter::{Product, Sum};

#[doc(hidden)]
#[macro_export(local_inner_macros)]
macro_rules! _iter {
    ($e:expr) => { Some(std::iter::once($e)) };
    ($i:ident in $range:expr, $($rest:tt)+ ) => {
        Some($range.filter_map(move |$i| _iter!($($rest)*)).flatten())
    };
    ($cond:expr, $($rest:tt)*) => {
        if $cond {
            _iter!($($rest)*)
        } else {
            None
        }
    };
}

/// An iterator comprehension.
///
/// A sequence of generating expressions and boolean conditions for generating a
/// sequence of values.
///
/// # Example
///
/// ```
/// use iter_comprehensions::iter;
/// assert_eq!(iter!(4*i + j; i in 0..5, j in 0..5, i < j).collect::<Vec<_>>(),
///            vec![1, 2, 3, 4, 6, 7, 8, 11, 12, 16]);
/// ```
#[macro_export(local_inner_macros)]
macro_rules! iter {
    ($e:expr; $($rest:tt)*) => { iter!($($rest)*, $e) };
    ($($rest:tt)*) => { _iter!($($rest)*).unwrap() };
}

/// A iterator comprehensions generating a vector.
///
/// This is short form form for `iter!(...).collect::<Vec<_>>()`.
///
/// # Example
///
/// ```
/// use iter_comprehensions::vector;
/// assert_eq!(vector!(4*i + j; i in 0..5, j in 0..5, i < j),
///            vec![1, 2, 3, 4, 6, 7, 8, 11, 12, 16]);
/// ```
#[macro_export(local_inner_macros)]
macro_rules! vector {
    ($($rest:tt)*) => { iter!($($rest)*).collect::<Vec<_>>() }
}

/// A sum expression.
///
/// The macro takes as first argument a type.
///
/// # Example
///
/// The following expression corresponds to the mathematical expression
/// $\\sum_{i=1}^{10} i$.
///
/// ```
/// use iter_comprehensions::sum;
/// assert_eq!(sum!(i in 1..=10, i), 55);
/// ```
#[macro_export(local_inner_macros)]
macro_rules! sum {
    ($($rest:tt)*) => { $crate::sumiter(iter!($($rest)*)) }
}

/// A product expression.
///
/// The macro takes as first argument a type.
///
/// # Example
///
/// The following expression corresponds to the mathematical expression
/// $\\prod_{i=1}^5 i$.
///
/// ```
/// use iter_comprehensions::product;
/// assert_eq!(product!(i in 1..=5, i), 120);
/// ```
#[macro_export(local_inner_macros)]
macro_rules! product {
    ($($rest:tt)*) => { $crate::productiter(iter!($($rest)*)) }
}

pub fn sumiter<T, I>(it: I) -> T
where
    T: Sum<T>,
    I: Iterator<Item = T>,
{
    it.sum()
}

pub fn productiter<T, I>(it: I) -> T
where
    T: Product<T>,
    I: Iterator<Item = T>,
{
    it.product()
}

#[cfg(test)]
mod tests {
    #[test]
    fn iter_set_form() {
        assert_eq!(iter!(i; i in 0..5).collect::<Vec<_>>(), vec![0, 1, 2, 3, 4]);
        assert_eq!(iter!(i; i in 0..5, i % 2 == 0).collect::<Vec<_>>(), vec![0, 2, 4]);
        assert_eq!(
            iter!(i*2 + j; i in 0..5, j in 0..2).collect::<Vec<_>>(),
            (0..10).collect::<Vec<_>>()
        );
        assert_eq!(
            iter!((i,j); i in 0..3, j in 0..3, i < j).collect::<Vec<_>>(),
            vec![(0, 1), (0, 2), (1, 2)]
        );
        assert_eq!(
            iter!((i,j); i in 0..4, i % 2 == 0, j in 0..4, i < j).collect::<Vec<_>>(),
            vec![(0, 1), (0, 2), (0, 3), (2, 3)]
        );
    }

    #[test]
    fn iter() {
        assert_eq!(iter!(i in 0..5, i).collect::<Vec<_>>(), vec![0, 1, 2, 3, 4]);
        assert_eq!(iter!(i in 0..5, i % 2 == 0, i).collect::<Vec<_>>(), vec![0, 2, 4]);
        assert_eq!(
            iter!(i in 0..5, j in 0..2, i*2 + j).collect::<Vec<_>>(),
            (0..10).collect::<Vec<_>>()
        );
        assert_eq!(
            iter!(i in 0..3, j in 0..3, i < j, (i,j)).collect::<Vec<_>>(),
            vec![(0, 1), (0, 2), (1, 2)]
        );
        assert_eq!(
            iter!(i in 0..4, i % 2 == 0, j in 0..4, i < j, (i,j)).collect::<Vec<_>>(),
            vec![(0, 1), (0, 2), (0, 3), (2, 3)]
        );
    }

    #[test]
    fn vector() {
        assert_eq!(vector!(i; i in 0..5), vec![0, 1, 2, 3, 4]);
    }

    #[test]
    fn sum() {
        assert_eq!(sum!(i in 0..5, i), 10);
    }

    #[test]
    fn product() {
        assert_eq!(product!(i in 1..5, i), 24);
    }
}