/*
 * 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 and vector comprehensions.
//!
//! This crate implements a few macros with generating comprehension
//! expressions.
//!
//! 1. `comprehension!` for generating a sequence of index tuples
//! 2. `map!` for generating a sequence of expressions
//! 3. `vec!` for constructing vectors
//! 4. `sum!` for computing the sum of some values
//! 5. `product!` for computing the product of some values
//!
//! The macro `comprehension!` can be used to generate a sequence of elements using
//! generating sequences and conditional filters.
//!
//! ```nodoc
//! comprehension!(i1 in RANGE1, COND1, ..., ik in RANGEk)
//! ```
//!
//! where RANGE* are iterators (in fact, everything implementing `IntoIterator`)
//! and each COND* is a boolean condition. Each `RANGE` and `COND` term can use
//! the variables declared in preceeding range expressions.
//!
//! The macro `map!` adds an additional expression that computes a value
//! depending on the indices:
//!
//! ```nodoc
//! map!(i1 in RANGE1, COND1, ..., ik in RANGEk, EXPR)
//! map!(EXPR; i1 in RANGE1, COND1, ..., ik in RANGEk)
//! ```
//!
//! # Example
//! The expression $\\{ 5i + j : i \\in \\{0..4\\}, j \\in \\{0..4\\}, i < j \\}$ is
//! equivalent to the following form
//!
//! ```
//! use iter_comprehensions::map;
//! assert_eq!(map!(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]);
//! ```
//!
//! The analogous syntax can be used to create vectors:
//!
//! ```
//! use iter_comprehensions::vec;
//! assert_eq!(vec![i; i in 0..5], vec![0,1,2,3,4]);
//! assert_eq!(vec![(i,j); i in 0..3, j in 0..3, i < j],
//!            vec![(0,1), (0,2), (1,2)]);
//! ```
//!
//! Computing a sum of values:
//!
//! ```
//! use iter_comprehensions::{sum, vec};
//! assert_eq!(sum!(i; i in 1..10, i % 2 == 1), 25);
//! let S = vec![i; i in 1..10];
//! assert_eq!(sum!(i in S, i % 2 == 1, i), 25);
//! ```
//!
//! Computing a product of values:
//!
//! ```
//! use iter_comprehensions::product;
//! assert_eq!(product!(i; i in 1..=5), 120);
//! assert_eq!(product!(i in 1..=5, i), 120);
//! ```

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

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

#[doc(hidden)]
#[macro_export(local_inner_macros)]
macro_rules! _comprehension {
    (($($vars:tt),*); ($($clauses:tt)*); $i:ident in $range:expr) => {
        _comprehension!(($($vars,)* $i); ($($clauses)* ($i in $range));)
    };

    (($($vars:tt),*); ($($clauses:tt)*); $i:ident in $range:expr, $($rest:tt)*) => {
        _comprehension!(($($vars,)* $i); ($($clauses)* ($i in $range)); $($rest)*)
    };

    (($($vars:tt),*); ($($clauses:tt)*); $cond:expr) => {
        _comprehension!(($($vars),*); ($($clauses)* ($cond));)
    };

    (($($vars:tt),*); ($($clauses:tt)*); $cond:expr, $($rest:tt)*) => {
        _comprehension!(($($vars),*); ($($clauses)* ($cond)); $($rest)*)
    };

    (($($vars:tt),*); ($(($($clauses:tt)+))*);) => {
        _comprehension_iter!((); $($($clauses)+,)* ($($vars),*)).unwrap()
    };
}

/// A comprehension for constructing an iterator.
///
/// This macro generates a sequence of elements (i.e. an `Iterator`) using
/// generating sequences and conditional filters. The iterator receives a tuple
/// consisting of the variables in the range expressions.
///
/// # Example
/// ```
/// use iter_comprehensions::comprehension;
/// let mut result = vec![];
/// for (i,j) in comprehension!(i in 0..5, j in i+1..5, (j + i) % 2 == 0) {
///     result.push((i,j));
/// }
/// assert_eq!(result, vec![(0, 2), (0, 4), (1, 3), (2, 4)]);
/// ```
#[macro_export(local_inner_macros)]
macro_rules! comprehension {
    ($($rest:tt)*) => {
        _comprehension!((); (); $($rest)*)
    }
}

#[doc(hidden)]
#[macro_export(local_inner_macros)]
macro_rules! _map {
    (($($vars:tt),*); ($($clauses:tt)*); $i:ident in $range:expr, $($rest:tt)*) => {
        _map!(($($vars,)* $i); ($($clauses)* ($i in $range)); $($rest)*)
    };

    (($($vars:tt),*); ($($clauses:tt)*); $cond:expr, $($rest:tt)*) => {
        _map!(($($vars),*); ($($clauses)* ($cond)); $($rest)*)
    };

    (($var:ident); ($(($($clauses:tt)+))*); $e:expr) => {
        _comprehension_iter!((); $($($clauses)+,)* $var).unwrap().map(|$var| $e)
    };

    (($($vars:tt),*); ($(($($clauses:tt)+))*); $e:expr) => {
        _comprehension_iter!((); $($($clauses)+,)* ($($vars),*)).unwrap().map(|($($vars),*)| $e)
    };
}

/// A map comprehension.
///
/// A sequence of generating expressions and boolean conditions for generating a
/// sequence of values. The two possible macro calls
/// ```norun
/// map!(EXPR; i1 in RANGE1, COND2, ..., ik in RANGEk)
/// map!(i1 in RANGE1, COND2, ..., ik in RANGEk, EXPR)
/// ```
/// is roughly the same as
/// ```norun
/// comprehensions!(i1 in RANGE1, COND2, ..., ik in RANGEk).map(|(i1, i2, ..., ik)| EXPR)
/// ```
///
/// # Example
///
/// ```
/// use iter_comprehensions::map;
/// let mut it = map!(4*i + j; i in 0..5, j in 0..5, i < j);
/// assert_eq!(it.next(), Some(1));
/// assert_eq!(it.next(), Some(2));
/// assert_eq!(it.next(), Some(3));
/// assert_eq!(it.next(), Some(4));
/// assert_eq!(it.next(), Some(6));
/// assert_eq!(it.next(), Some(7));
/// assert_eq!(it.next(), Some(8));
/// assert_eq!(it.next(), Some(11));
/// assert_eq!(it.next(), Some(12));
/// assert_eq!(it.next(), Some(16));
/// assert_eq!(it.next(), None);
/// ```

#[macro_export(local_inner_macros)]
macro_rules! map {
    ($e:expr; $($rest:tt)*) => {
        map!($($rest)*, $e)
    };
    ($($rest:tt)*) => {
        _map!((); (); $($rest)*)
    }
}

/// Extension of `vec!` with comprehensions.
///
/// This is a short form for `map!(...).collect::<Vec<_>>()`.
///
/// For convencience the `vec!` macro also supports the syntax of the
/// `std::vec!`, i.e. `vec![0, 1, 2]` and `vec![x; 5]`.
///
/// # Example
///
/// ```
/// use iter_comprehensions::vec;
/// assert_eq!(vec!(4*i + j; i in 0..5, j in 0..5, i < j),
///            std::vec![1, 2, 3, 4, 6, 7, 8, 11, 12, 16]);
/// ```
#[macro_export(local_inner_macros)]
macro_rules! vec {
    ($e:expr; $i:ident in $range:expr) => {
        map!($e; $i in $range).collect::<Vec<_>>()
    };
    ($e:expr; $i:ident in $range:expr, $($rest:tt)*) => {
        map!($e; $i in $range, $($rest)*).collect::<Vec<_>>()
    };
    ($($rest:tt)*) => {
        std::vec!($($rest)*)
    };
}

/// A sum expression.
///
/// The two possible macro calls
/// ```norun
/// sum!(EXPR; i1 in RANGE1, COND2, ..., ik in RANGEk)
/// sum!(i1 in RANGE1, COND2, ..., ik in RANGEk, EXPR)
/// ```
/// are roughly equivalent to
/// ```norun
/// map!(EXPR; i1 in RANGE1, COND2, ..., ik in RANGEk).sum()
/// map!(i1 in RANGE1, COND2, ..., ik in RANGEk, EXPR).sum()
/// ```
///
/// # 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);
/// assert_eq!(sum!(i; i in 1..=10), 55);
/// ```
///
/// It is also possible to specify the return type of the product in case it
/// cannot be infered automatically.
///
/// ```
/// use iter_comprehensions::sum;
/// assert_eq!(sum!(usize : i; i in 1..=10), 55);
/// ```
#[macro_export(local_inner_macros)]
macro_rules! sum {
    ($t:ty : $e:expr; $($rest:tt)*) => { sum!($t : $($rest)*, $e) };
    ($t:ty : $($rest:tt)*) => { map!($($rest)*).sum::<$t>() };
    ($e:expr; $($rest:tt)*) => { sum!($($rest)*, $e) };
    ($($rest:tt)*) => { $crate::sumiter(map!($($rest)*)) };
}

/// A product expression.
///
/// The two possible macro calls
/// ```norun
/// product!(EXPR; i1 in RANGE1, COND2, ..., ik in RANGEk)
/// product!(i1 in RANGE1, COND2, ..., ik in RANGEk, EXPR)
/// ```
/// are roughly equivalent to
/// ```norun
/// map!(EXPR; i1 in RANGE1, COND2, ..., ik in RANGEk).product()
/// map!(i1 in RANGE1, COND2, ..., ik in RANGEk, EXPR).product()
/// ```
///
/// # 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);
/// assert_eq!(product!(i; i in 1..=5), 120);
/// ```
///
/// It is also possible to specify the return type of the product in case it
/// cannot be infered automatically.
/// ```
/// use iter_comprehensions::product;
/// assert_eq!(product!(usize : i; i in 1..=5), 120);
/// ```
#[macro_export(local_inner_macros)]
macro_rules! product {
    ($t:ty : $e:expr; $($rest:tt)*) => { product!($t : $($rest)*, $e) };
    ($t:ty : $($rest:tt)*) => { map!($($rest)*).product::<$t>() };
    ($e:expr; $($rest:tt)*) => { product!($($rest)*, $e) };
    ($($rest:tt)*) => { $crate::productiter(map!($($rest)*)) };
}

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

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

#[cfg(test)]
mod tests {
    #[test]
    fn test_comprehension() {
        let mut it = comprehension!(k in 0..3);
        assert_eq!(it.next(), Some(0));
        assert_eq!(it.next(), Some(1));
        assert_eq!(it.next(), Some(2));
        assert_eq!(it.next(), None);

        let mut it = comprehension!(i in 0..3, j in 0..3);
        assert_eq!(it.next(), Some((0, 0)));
        assert_eq!(it.next(), Some((0, 1)));
        assert_eq!(it.next(), Some((0, 2)));
        assert_eq!(it.next(), Some((1, 0)));
        assert_eq!(it.next(), Some((1, 1)));
        assert_eq!(it.next(), Some((1, 2)));
        assert_eq!(it.next(), Some((2, 0)));
        assert_eq!(it.next(), Some((2, 1)));
        assert_eq!(it.next(), Some((2, 2)));
        assert_eq!(it.next(), None);

        let mut it = comprehension!(i in 0..3, i % 2 == 0, j in 0..3);
        assert_eq!(it.next(), Some((0, 0)));
        assert_eq!(it.next(), Some((0, 1)));
        assert_eq!(it.next(), Some((0, 2)));
        assert_eq!(it.next(), Some((2, 0)));
        assert_eq!(it.next(), Some((2, 1)));
        assert_eq!(it.next(), Some((2, 2)));
        assert_eq!(it.next(), None);

        let mut it = comprehension!(i in 0..3, i % 2 == 0, j in 0..3, i + j < 4);
        assert_eq!(it.next(), Some((0, 0)));
        assert_eq!(it.next(), Some((0, 1)));
        assert_eq!(it.next(), Some((0, 2)));
        assert_eq!(it.next(), Some((2, 0)));
        assert_eq!(it.next(), Some((2, 1)));
        assert_eq!(it.next(), None);

        assert_eq!(
            comprehension!(i in 0..3, j in 0..3, k in 0..2)
                .map(|(i, j, k)| i + j + k)
                .collect::<Vec<_>>(),
            vec![0, 1, 1, 2, 2, 3, 1, 2, 2, 3, 3, 4, 2, 3, 3, 4, 4, 5]
        );
    }

    #[test]
    fn test_map() {
        let mut it = map!(k in 0..3, 2 * k);
        assert_eq!(it.next(), Some(0));
        assert_eq!(it.next(), Some(2));
        assert_eq!(it.next(), Some(4));
        assert_eq!(it.next(), None);

        assert_eq!(map!(i; i in 0..5).collect::<Vec<_>>(), vec![0, 1, 2, 3, 4]);
        assert_eq!(map!(i; i in 0..5, i % 2 == 0).collect::<Vec<_>>(), vec![0, 2, 4]);
        assert_eq!(
            map!(i*2 + j; i in 0..5, j in 0..2).collect::<Vec<_>>(),
            (0..10).collect::<Vec<_>>()
        );
        assert_eq!(
            map!((i,j); i in 0..3, j in 0..3, i < j).collect::<Vec<_>>(),
            vec![(0, 1), (0, 2), (1, 2)]
        );
        assert_eq!(
            map!((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 vector() {
        assert_eq!(vec!(i; i in 0..5), vec![0, 1, 2, 3, 4]);
    }

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

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

    #[test]
    fn sumvec() {
        let x = vec![1; 10];
        assert_eq!(sum!(x[i]; i in 0..10), 10);
        assert_eq!(sum!(i in 0..10, x[i]), 10);
    }

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

    #[test]
    fn test_mut() {
        let mut x = vec![0, 1, 2];
        sum!({x.push(42 + i + j); x.len()}; i in 0..5, j in 0..3);
    }

}