product.rs

#![feature(macro_rules)]

/// A trait for (x,y,z) ++ w => (x,y,z,w)
trait AppendTuple<T, U> {
    fn append(self, x: T) -> U;
}

macro_rules! impl_append_tuple(
    () => (
        impl<T> AppendTuple<T, (T, )> for () {
            fn append(self, x: T) -> (T, ) {
                (x, )
            }
        }
    );

    ($A:ident $(,$B:ident)*) => (
        impl_append_tuple!($($B),*)
        #[allow(uppercase_variables)]
        impl<$A, $($B,)* T> AppendTuple<T, ($A, $($B,)* T)> for ($A, $($B),*) {
            fn append(self, x: T) -> ($A, $($B,)* T) {
                let ($A, $($B),*) = self;
                ($A, $($B,)* x)
            }
        }
    );
)

impl_append_tuple!(A, B, C, D, E, F, G, H)

pub macro_rules! product(
    ($I:expr) => (
        ($I)
    );
    ($I:expr, $J:expr $(, $K:expr)*) => (
        {
            let it = product($I, $J);
            //($I).flat_map(|x| Repeat::new(x).zip($J))
            $(
                let it = product(it, $K).map(|(t, x)| t.append(x));
            )*
            it
        }
    );
)


pub macro_rules! icompr(
    ($r:expr for $x:pat in $J:expr if $pred:expr) => (
        ($J).filter_map(|$x| if $pred { Some($r) } else { None })
    );
    ($r:expr for $x:pat in $J:expr) => (
        ($J).filter_map(|$x| Some($r))
    );
)

#[deriving(Clone)]
pub struct Product<A, I, J> {
    a: I,
    a_cur: Option<A>,
    b: J,
    b_orig: J,
}

impl<A: Clone, I: Iterator<A>, B, J: Clone + Iterator<B>>
Iterator<(A, B)> for Product<A, I, J>
{
    fn next(&mut self) -> Option<(A, B)>
    {
        let elt_b = match self.b.next() {
            None => {
                self.b = self.b_orig.clone();
                match self.b.next() {
                    None => return None,
                    Some(x) => {
                        self.a_cur = self.a.next();
                        x
                    }
                }
            }
            Some(x) => x
        };
        match self.a_cur {
            None => None,
            Some(ref a) => {
                Some((a.clone(), elt_b))
            }
        }
    }

    fn size_hint(&self) -> (uint, Option<uint>)
    {
        let (a, ah) = self.a.size_hint();
        let (b, bh) = self.b.size_hint();
        let (bo, boh) = self.b_orig.size_hint();

        // Compute a * bo + b for both lower and upper bound
        let low = a.checked_mul(&bo)
                    .and_then(|x| x.checked_add(&b))
                    .unwrap_or(std::uint::MAX);
        let high = ah.and_then(|x| boh.and_then(|y| x.checked_mul(&y)))
                     .and_then(|x| bh.and_then(|y| x.checked_add(&y)));
        (low, high)
    }
}

pub fn product<A: Clone, B,
             I: Iterator<A>, J: Clone + Iterator<B>>
             (i: I, j: J) -> Product<A, I, J>
{
    let mut i = i;
    Product{a_cur: i.next(), a: i, b: j.clone(), b_orig: j}
}


/// Clonable iterator adaptor to map elementwise
/// from `Iterator<A>` to `Iterator<B>`
pub struct StaticMap<A, B, I> {
    map: fn(A) -> B,
    iter: I,
}

impl<A, B, I> StaticMap<A, B, I> 
{
    #[inline]
    fn new(iter: I, map: fn(A) -> B) -> StaticMap<A, B, I> {
        StaticMap{iter: iter, map: map}
    }
}

impl<A, B, I: Iterator<A>> Iterator<B> for StaticMap<A, B, I>
{
    #[inline]
    fn next(&mut self) -> Option<B> {
        self.iter.next().map(|a| (self.map)(a))
    }

    fn size_hint(&self) -> (uint, Option<uint>) {
        self.iter.size_hint()
    }
}

impl<A, B, I: Clone> Clone for StaticMap<A, B, I>
{
    fn clone(&self) -> StaticMap<A, B, I> {
        StaticMap::new(self.iter.clone(), self.map)
    }
}

fn main() {
    let v = ["a", "b", "c"];

    for (a, b) in product!(v.iter(), range(-3, 0i)) {
        print!("({}, {}) ", a, b);
    }
    println!("");

    for a in icompr!(x*x for x in range(0i, 10) if x % 2 == 1) {
        print!("{}, ", a);
    }
    println!("");

    for elt in product!(range(0, 2i), range(0, 3i), range(0, 3i)) {
        print!("{}, ", elt);
    }
    println!("");

    for elt in product!(range(0, 2i), range(0, 2i), range(0, 2i), range(0, 2i), range(0,2i)) {
        print!("{}, ", elt);
    }
    println!("");

    /*
    */

    for elt in icompr!(x + y + z for (x, y, z) in
                       product!(range(0, 2i), range(0, 3i), range(0, 3i))) {
        print!("{}, ", elt);
    }
    println!("");

    println!("{}", ().append(2i));
    println!("{}", ('a', 3u, &[1i], 2i8, b'z').append(2i));
}