#![feature(generators, generator_trait, type_alias_impl_trait)]
use std::{
    ops::{Generator, GeneratorState},
    pin::Pin,
};

/// A wrapper struct around Generators,
/// providing a safe implementation of the [`Iterator`] trait.
pub struct GenIter<G>(Option<G>);

impl<G: Generator + Unpin> GenIter<G> {
    /// Creates a new `GenIter` instance from a generator.
    /// The returned instance can be iterated over,
    /// consuming the generator.
    #[inline]
    pub fn new(gen: G) -> Self {
        Self(Some(gen))
    }
}

impl<G: Generator + Unpin> Iterator for GenIter<G> {
    type Item = G::Yield;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        Pin::new(self).next()
    }
}

impl<G: Generator> Iterator for Pin<&mut GenIter<G>> {
    type Item = G::Yield;

    fn next(&mut self) -> Option<Self::Item> {
        let this: Pin<&mut GenIter<G>> = self.as_mut();

        // This should be safe.
        // this Iterator implementation is on a Pin<&mut GenIter<G>> where G: Generator.
        // In order to acquire such a Pin<&mut GenIter<G>> if G does *NOT* implement Unpin,
        // the unsafe `new_unchecked` function from the Pin type must be used anyway.
        //
        // Note that if G: Unpin, the Iterator implementation of GenIter<G> itself is used,
        // which just creates a Pin safely, and then delegates to this implementation.
        let gen: Pin<&mut Option<G>> = unsafe { this.map_unchecked_mut(|geniter| &mut geniter.0) };

        let gen: Option<Pin<&mut G>> = Option::as_pin_mut(gen);

        match gen.map(Generator::resume) {
            Some(GeneratorState::Yielded(y)) => Some(y),
            Some(GeneratorState::Complete(_)) => {
                self.set(GenIter(None));
                None
            }
            None => None,
        }
    }
}

macro_rules! try_yield {
    ($option:expr) => {
        match $option {
            Some(x) => yield x,
            None => return,
        }
    };
}

macro_rules! iter {
    ($($b:tt)*) => {
        $crate::GenIter::new(move || { $($b)* })
    };
}

trait AlternateExt: Iterator {
    type Alternate: Iterator<Item = Self::Item>;

    fn alternate(self) -> Self::Alternate;
}

impl <I: DoubleEndedIterator> AlternateExt for I {
    type Alternate = impl Iterator<Item = Self::Item>;

    fn alternate(mut self) -> Self::Alternate {
        iter! {
            loop {
                try_yield!(self.next());
                try_yield!(self.next_back());
            }
        }
    }
}

trait AlternateByExt: Iterator {
    type AlternateBy: Iterator<Item = Self::Item>;

    fn alternate_by(self, count: usize) -> Self::AlternateBy;
}

impl <I: DoubleEndedIterator> AlternateByExt for I {
    type AlternateBy = impl Iterator<Item = Self::Item>;

    fn alternate_by(mut self, count: usize) -> Self::AlternateBy {
        iter! {
            loop {
                for _ in 0..count {
                    try_yield!(self.next());
                }
                for _ in 0..count {
                    try_yield!(self.next_back());
                }
            }
        }
    }
}

fn main() {
    for x in (0..10).alternate_by(3) {
        dbg!(x);
    }
}