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
}
};
}
#[macro_export(local_inner_macros)]
macro_rules! iter {
($e:expr; $($rest:tt)*) => { iter!($($rest)*, $e) };
($($rest:tt)*) => { _iter!($($rest)*).unwrap() };
}
#[macro_export(local_inner_macros)]
macro_rules! vector {
($($rest:tt)*) => { iter!($($rest)*).collect::<Vec<_>>() }
}
#[macro_export(local_inner_macros)]
macro_rules! sum {
($($rest:tt)*) => { $crate::sumiter(iter!($($rest)*)) }
}
#[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);
}
}