iter_comprehensions/

lib.rs

/*
 * MIT License
 *
 * Copyright (c) 2019-2024 Frank Fischer <frank-fischer@shadow-soft.de>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */

//! 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.
//!
//! ```no_test
//! 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 preceding range expressions.
//!
//! The macro `map!` adds an additional expression that computes a value
//! depending on the indices:
//!
//! ```no_test
//! 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, \\ldots, 4\\}, j \\in \\{0,
//! \\ldots, 4\\}, i < j \\}$ is equivalent to the following form
//!
//! ```
//! use iter_comprehensions::map;
//! assert_eq!(map!(5*i + j; i in 0..5, j in 0..5, i < j).collect::<Vec<_>>(),
//!            vec![1, 2, 3, 4, 7, 8, 9, 13, 14, 19]);
//! ```
//!
//! 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);
//! ```
//!
//! It is also possible to use tuple destructuring binds in
//! RANGE expressions:
//! ```
//! use iter_comprehensions::sum;
//! assert_eq!(sum!(i*j; (i,j) in [(1,2), (3,4), (5,6)]), 44);
//! assert_eq!(sum!((i,j) in [(1,2), (3,4), (5,6)], i*j), 44);
//! ```


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:tt movein $range:expr, $($rest:tt)+ ) => {
        Some($range
             .into_iter()
             .zip(std::iter::repeat($r))
             .filter_map(move |($i, $r)| _comprehension_iter!(($i, $r); $($rest)*)).flatten())
    };
    ($r:tt; $i:tt 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:tt in $range:expr) => {
        _comprehension!(($($vars,)* $i); ($($clauses)* ($i in $range));)
    };

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

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

    (($($vars:tt),*); ($($clauses:tt)*); move $i:tt in $range:expr, $($rest:tt)*) => {
        _comprehension!(($($vars,)* $i); ($($clauses)* ($i movein $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:tt 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:tt); ($(($($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
/// ```no_test
/// map!(EXPR; i1 in RANGE1, COND2, ..., ik in RANGEk)
/// map!(i1 in RANGE1, COND2, ..., ik in RANGEk, EXPR)
/// ```
/// is roughly the same as
/// ```no_test
/// 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 convenience 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:tt in $range:expr) => {
        map!($e; $i in $range).collect::<Vec<_>>()
    };
    ($e:expr; $i:tt 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
/// ```no_test
/// sum!(EXPR; i1 in RANGE1, COND2, ..., ik in RANGEk)
/// sum!(i1 in RANGE1, COND2, ..., ik in RANGEk, EXPR)
/// ```
/// are roughly equivalent to
/// ```no_test
/// 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 sum in case it
/// cannot be inferred 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
/// ```no_test
/// product!(EXPR; i1 in RANGE1, COND2, ..., ik in RANGEk)
/// product!(i1 in RANGE1, COND2, ..., ik in RANGEk, EXPR)
/// ```
/// are roughly equivalent to
/// ```no_test
/// 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 inferred 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];
        assert_eq!(
            18 * 19 / 2 - 6,
            sum!({x.push(42 + i + j); x.len()}; i in 0..5, j in 0..3)
        );
        assert_eq!(
            x,
            vec![0, 1, 2, 42, 43, 44, 43, 44, 45, 44, 45, 46, 45, 46, 47, 46, 47, 48]
        )
    }

    #[test]
    fn test_tuple() {
        let pairs = [(1,2), (3,4), (5,6)];
        let mut it = map!(format!("{}+{}", i, j); (i, j) in pairs);
        assert_eq!(it.next(), Some("1+2".to_string()));
        assert_eq!(it.next(), Some("3+4".to_string()));
        assert_eq!(it.next(), Some("5+6".to_string()));
        assert_eq!(it.next(), None);
    }

    #[test]
    fn test_tuple2() {
        let pairs = [(1,2), (3,4), (5,6)];
        let mut it = map!((i, j) in pairs, format!("{}+{}", i, j));
        assert_eq!(it.next(), Some("1+2".to_string()));
        assert_eq!(it.next(), Some("3+4".to_string()));
        assert_eq!(it.next(), Some("5+6".to_string()));
        assert_eq!(it.next(), None);
    }

    #[test]
    fn test_sum() {
        let pairs = [(1,2), (3,4), (5,6)];
        assert_eq!(44, sum!(i * j; (i,j) in pairs));
        assert_eq!(44, sum!((i,j) in pairs, i * j));
    }
}