`iter` blocks

playground.rs

#![feature(arbitrary_self_types, generators, generator_trait, pin)]

use std::{
    marker::Unpin,
    ops::{DerefMut, Deref, Generator, GeneratorState},
    pin::Pin,
};

trait PinIterator {
    type Item;

    fn next(self: Pin<&mut Self>) -> Option<Self::Item>;
}

struct P<T>(T);

impl<T> Iterator for P<T> where T: PinIterator + Unpin {
    type Item = T::Item;

    fn next(&mut self) -> Option<Self::Item> {
        Pin::new(&mut self.0).next()
    }
}

impl<T> PinIterator for P<T> where T: Iterator, P<T>: Unpin {
    type Item = T::Item;

    fn next(self: Pin<&mut Self>) -> Option<Self::Item> {
        Pin::get_mut(self).0.next()
    }
}

impl<P> PinIterator for Pin<P>
where
    P: DerefMut,
    P::Target: PinIterator,
{
    type Item = <<P as Deref>::Target as PinIterator>::Item;

    fn next(self: Pin<&mut Self>) -> Option<Self::Item> {
        Pin::get_mut(self).as_mut().next()
    }
}

// https://docs.rs/pin-utils/0.1.0-alpha.3/src/pin_utils/stack_pin.rs.html#13-24
macro_rules! pin_mut {
    ($($x:ident),*) => { $(
        let mut $x = $x;
        let mut $x = unsafe { Pin::new_unchecked(&mut $x) };
    )* }
}

macro_rules! iter {
    ($($s:stmt);*) => {
        struct GenIter<G>(G);
        impl<G> PinIterator for GenIter<G> where G: Generator<Return = ()> {
            type Item = G::Yield;

            fn next(self: Pin<&mut Self>) -> Option<Self::Item> {
                match unsafe { Pin::get_mut_unchecked(self).0.resume() } {
                    GeneratorState::Yielded(item) => Some(item),
                    GeneratorState::Complete(()) => None,
                }
            }
        }
        GenIter(static move || {
            $($s)*;
            #[allow(unreachable_code)]
            { if false { yield return } }
        })
    };
}

fn foo(maxes: impl Iterator<Item = u32>) -> impl PinIterator<Item = u32> {
    iter! {
        for max in maxes {
            let mut x = 0;
            while x <= max {
                yield x;
                x += 1;
            }
        }
    }
}

fn bar(maxes: impl Iterator<Item = u32>) -> impl PinIterator<Item = u32> {
    iter! {
        let foo = foo(maxes);
        pin_mut!(foo);
        while let Some(i) = foo.as_mut().next() {
            yield i * 2;
        }
    }
}

macro_rules! pin_for {
    (($p:pat in $e:expr) { $($s:stmt);* }) => {
        let iter = $e;
        pin_mut!(iter);
        while let Some($p) = PinIterator::next(iter.as_mut()) {
            $($s);*
        }
    };
}

fn qux(bar: impl PinIterator<Item = u32>) -> impl PinIterator<Item = u32> {
    iter! {
        pin_for!((i in bar) {
            if (i / 2) % 2 == 0 {
                yield i + 1;
            }
        })
    }
}

fn main() {
    let maxes = vec![8, 4];
    
    let b = Box::pinned(bar(maxes.iter().cloned()));
    for i in P(b) {
        println!("{}", i);
    }
    
    println!();
    
    let q = Box::pinned(qux(bar(maxes.into_iter())));
    for i in P(q) {
        println!("{}", i);
    }
}