Python list comprehensions in Rust

list_comprehensions.rs

#![feature(macro_rules)]

/// MIT license etc. etc.
///
/// Experimenting with Python-like list comprehensions.
///
/// An attempt to explore the possibility space for ergonomic improvements
/// in Rust that can come post v1.0. Notice that there are not type declerations. 
/// Aside from the "c!" macro invocation, Rust allows for an exact copy of the
/// Python comprehension syntax.
///
/// This should probably work on slices. Something to look forward to.


use std::num::{abs, pow}; // Used in unit tests only.


macro_rules! c(
    ($m:expr for $el:ident in $iter:expr) => (
        {
            let (lower, upper) = $iter.size_hint();
            let mut result = match upper {
                Some(bound) => Vec::with_capacity(bound),
                None => Vec::with_capacity(lower),
            };

            for $el in $iter { 
                let $el = $el.clone();
                result.push($m);
            }
            result
        }
    );

    ($m:expr for $el:ident in $iter:expr if $cond:expr) => (
        {
            let mut result = Vec::new();
            for $el in $iter { 
                let $el = $el.clone();
                if $cond {
                    result.push($m);
                }
            }
            result
        }
    );
    ($m:expr for $el:ident in $iter:expr 
     $(for $inner_el:ident in $inner_iter:expr)* 
     $(if $cond:expr)*) => (

        {
            let mut result = Vec::new();
            for $el in $iter {
                let $el = $el.clone();
                result.push_all_move(c!($($m for $inner_el in $inner_iter)+ $(if $cond)*));
            }
            result
        }
    );
)

#[test]
fn simple_comprehension() {
    let squares = c![ x * x for x in range(0, 10i) ];
    assert_eq!(squares, vec![0i, 1, 4, 9, 16, 25, 36, 49, 64, 81]);
}

#[test]
fn vector_to_iter_comprehension() {
    let vec = vec![-4i, -2, 0, 2, 4];
    let output: Vec<int> = c![x*2 for x in vec.iter()];
    assert_eq!(output, vec![-8i, -4, 0, 4, 8]);
}

#[test]
fn filter_comprehension() {
    let vec = vec![-4i, -2, 0, 2, 4];
    let output = c![x for x in vec.iter() if x >= 0i];
    assert_eq!(output, vec![0i, 2, 4]);
}

#[test]
fn apply_function_comprehension() {
    let vec = vec![-4i, -2, 0, 2, 4];
    let output = c![abs(x) for x in vec.iter()];
    assert_eq!(output, vec![4i, 2, 0, 2, 4]);
}

#[test]
fn call_method_comprehension() {
    let fresh_fruit = vec!["  banana", " loganberry ", "passion fruit  "];
    let output = c![fruit.trim_chars(' ') for fruit in fresh_fruit.iter()];
    assert_eq!(output, vec!["banana", "loganberry", "passion fruit"]);
}


#[test]
fn change_type_comprehension() {
    let output = c![(x, pow(x,2)) for x in range(0i,6)];
    assert_eq!(output, vec![(0i, 0i), (1, 1), (2, 4), (3, 9), (4, 16), (5, 25)]);
}

#[test]
fn combined_comprehension() {
    let a = vec![1i,2,3];
    let b = vec![3i,1,4];
    let not_equal = c![(x, y) for x in a.iter() for y in b.iter() if x != y];
    assert_eq!(not_equal, vec![(1i, 3i), (1, 4), (2, 3), (2, 1), (2, 4), (3, 1), (3, 4)]);
    println!("{}", a);
}

 #[test]
 fn nested_comprehensions() {
     let matrix = vec![vec![1i,2,3,4],vec![5,6,7,8],vec![9,10,11,12]];
     let output = c![ c![ row[i] for row in matrix.iter()] for i in range(0u,4)];
     assert_eq!(output, vec![vec![1i,5,9], vec![2i,6,10], vec![3i,7,11], vec![4i,8,12]]);
 }