#![feature(generator_trait, generators)]

use std::ops::Generator;
use std::ops::GeneratorState;

fn test_select() {
        let gen = || {
            yield;
            yield;
            return 1;
        };

        let mut select = gen
            .into_stepfn()
            .map(|n| n * 10)
            .select(Callable::new(|| {
                yield;
                return "Foo";
            }));

        assert!(select.resume().unwrap().is_yield() == true);

        let (ret, mut remaining_callable) = match select.resume() {
            Some(State::Return(Either::Right(result))) => result,
            _ => panic!()
        };

        assert_eq!(ret, "Foo");
        assert_eq!(select.resume().is_none(), true);

        //select called the first generator twice, a next call results in a return!
        assert_eq!(remaining_callable.resume(), Some(State::Return(10)));

}

fn main() {
    test_select()
}

macro_rules! return_from_yield {
    ($g:expr) => {
        unsafe {
            match $g.resume() {
                GeneratorState::Yielded(_) => return Some(State::Yield),
                GeneratorState::Complete(ret) => ret,
            }
        }
    }
}

#[macro_export]
macro_rules! resume_from {
    ($s:expr) => {
        match $s.resume() {
            Some(State::Return(ret)) => ret,
            Some(State::Yield) => return Some(State::Yield),
            None => return None
        }
    }
}

#[derive(PartialEq, Eq, Ord, PartialOrd, Debug, Hash)]
pub enum State<T> {
    Yield,
    Return(T),
}

impl <T> State<T> {
    pub fn is_yield(&self) -> bool {
        match self {
            &State::Yield => true,
            _ => false,
        }
    }

    pub fn is_return(&self) -> bool {
        !self.is_yield()
    }

    pub fn take(self) -> Option<T> {
        match self {
            State::Return(value) => Some(value),
            _ => None
        }
    }
}

pub type Resume<T> = Option<State<T>>;

pub trait StepFn {
    type Return;

    fn resume(&mut self) -> Resume<Self::Return>;

    fn map<F, R>(self, f: F) -> Map<Self, F>
    where
        Self: Sized,
        F: FnMut(Self::Return) -> R,
    {
        Map::new(self, f)
    }

    fn filter<F>(self, predicate: F) -> Filter<Self, F>
    where
        Self:Sized,
        F: Fn(&Self::Return) -> bool
    {
        Filter::new(self, predicate)
    }

    fn select<S: StepFn>(self, other: S) -> Select<Self, S>
    where
        Self: Sized
    {
        Select::new(self, other)
    }
}

pub trait IntoStepFn {
    type StepFn: StepFn;

    fn into_stepfn(self) -> Self::StepFn;
}

impl <G> IntoStepFn for G
where
    G: Generator
{
    type StepFn = Callable<G>;

    fn into_stepfn(self) -> Self::StepFn {
        Callable::new(self)
    }
}
//we are responsible, as soon as we hit a return from the generator, we call .take()
//on the option, so we never call resume() again.
pub struct Callable<G>(Option<G>);

impl <G> Callable<G> {
    pub fn new(generator: G) -> Self {
        Callable(Some(generator))
    }
}

impl <G> StepFn for Callable<G>
where
    G: Generator
{
    type Return = G::Return;

    fn resume(&mut self) -> Resume<Self::Return> {

        let ret = return_from_yield!(self.0.as_mut()?);
        self.0.take();
        Some(State::Return(ret))
    }
}

pub struct Map<S, F> {
    resume: S,
    func: F,
}

impl <S, F> Map<S, F> {

    #[inline]
    fn new(resume: S, func: F) -> Self {
        Map { resume, func }
    }
}

impl <S, F, R> StepFn for Map<S, F>
where
    S: StepFn,
    F: FnMut(S::Return) -> R,
{
    type Return = R;

    #[inline]
    fn resume(&mut self) -> Resume<Self::Return> {
        let r = resume_from!(self.resume);

        Some(State::Return((self.func)(r)))
    }
}

pub struct Filter<S, F> {
    resume: S,
    predicate: F,
}

impl <S, F> Filter<S, F> {
    #[inline]
    fn new(resume: S, predicate: F) -> Self {
        Filter { resume, predicate }
    }
}

impl <S, F> StepFn for Filter<S, F>
where
    S: StepFn,
    F: Fn(&S::Return) -> bool
{
    type Return = S::Return;

    fn resume(&mut self) -> Resume<Self::Return> {
        let r = resume_from!(self.resume);

        if (self.predicate)(&r) {
            return Some(State::Return(r))
        }
        None
    }
}

#[derive(Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub enum Either<L, R> {
    Left(L),
    Right(R),
}

impl <L, R> Either <L, R> {
    pub fn is_right(&self) -> bool {
        match self {
            Either::Right(_) => true,
            _ => false,
        }
    }

    pub fn is_left(&self) -> bool {
        !self.is_right()
    }

    pub fn into_left(self) -> Option<L> {
        match self {
            Either::Left(value) => Some(value),
            _ => None,
        }
    }

    pub fn into_right(self) -> Option<R> {
        match self {
            Either::Right(value) => Some(value),
            _ => None,
        }
    }
}
#[derive(Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct Select<A, B> {
    inner: Option<(A, B)>
}

impl <A, B> Select<A, B> {
    fn new(a: A, b: B) -> Self {
        Select { inner: Some((a, b)) }
    }
}

impl <A: StepFn, B: StepFn> StepFn for Select<A, B>
{
    type Return = Either<(A::Return, B), (B::Return, A)>;

    fn resume(&mut self) -> Resume<Self::Return> {
        let (mut a, mut b) = self.inner.take()?;
        match a.resume() {
            Some(State::Return(ret)) => Some(State::Return(Either::Left((ret, b)))),
            | Some(State::Yield)
            | None => {
                match b.resume() {
                    Some(State::Return(ret)) => Some(State::Return(Either::Right((ret, a)))),
                    None => None,
                    Some(State::Yield) => {
                        self.inner = Some((a, b));
                        Some(State::Yield)
                    }
                }
            }
        }
    }
}