Simulating generators in Rust

gen.rs

  //! Generators in Rust, simulated at great expense using threads
//!
//! Here's a generator which yields the numbers from one to three:
//!
//!     let gen = generate(|(), out| {
//!         for i in 1..4 {
//!             out.yield_(i);
//!         }
//!     });
//!
//!     assert_eq!(gen.send(()), Some(1));
//!     assert_eq!(gen.send(()), Some(2));
//!     assert_eq!(gen.send(()), Some(3));
//!     assert_eq!(gen.send(()), None);
//!
//! This is implemented by running the closure in a separate thread.

use std::thread::{JoinHandle, spawn};
use std::sync::mpsc::{Receiver, SyncSender, sync_channel};

pub struct Generator<T, U> {
    thread: Option<JoinHandle<()>>,
    tx_in: SyncSender<T>,
    rx_out: Receiver<U>
}

pub struct Yield<T, U> {
    tx_out: SyncSender<U>,
    rx_in: Receiver<T>
}

impl<T, U> Generator<T, U> {
    fn send(&self, t: T) -> Option<U> {
        self.tx_in.send(t).unwrap();
        match self.rx_out.recv() {
            Ok(u) => Some(u),
            Err(_) => {
                None
            }
        }
    }
}

impl<T, U> Drop for Generator<T, U> {
    fn drop(&mut self) {
        // If the generator panicked, propagate that to this thread.
        let thread = self.thread.take().unwrap();
        thread.join().unwrap();
    }
}

impl<T, U> Yield<T, U>
    where T: Send + 'static,
          U: Send + 'static
{
    fn yield_(&self, u: U) -> T {
        self.tx_out.send(u).unwrap();
        self.rx_in.recv().unwrap()
    }
}

pub fn generate<B, T, U>(body: B) -> Generator<T, U>
    where B: FnOnce(T, Yield<T, U>) -> () + Send + 'static,
          T: Send + 'static,
          U: Send + 'static
{
    let (tx_in, rx_in) = sync_channel(0);
    let (tx_out, rx_out) = sync_channel(0);

    Generator {
        thread: Some(spawn(move || {
            let first = rx_in.recv().unwrap();
            body(first, Yield { tx_out: tx_out, rx_in: rx_in });
        })),
        tx_in: tx_in,
        rx_out: rx_out
    }
}

impl<U> Iterator for Generator<(), U>
    where U: Send + 'static
{
    type Item = U;
    fn next(&mut self) -> Option<U> { self.send(()) }
}

#[test]
fn dumb() {
    let gen = generate(|(), out| {
        for i in 1..4 {
            out.yield_(i);
        }
    });

    assert_eq!(gen.send(()), Some(1));
    assert_eq!(gen.send(()), Some(2));
    assert_eq!(gen.send(()), Some(3));
    assert_eq!(gen.send(()), None);
}

#[test]
fn tree_iter() {
    struct Node(Tree, i32, Tree);
    type Tree = Option<Box<Node>>;

    let left = Some(Box::new(Node(Some(Box::new(Node(None, 1, None))),
                                  2,
                                  Some(Box::new(Node(None, 3, None))))));
    let rite = Some(Box::new(Node(Some(Box::new(Node(None, 5, None))),
                                  6,
                                  Some(Box::new(Node(None, 7, None))))));
    let tree = Some(Box::new(Node(left, 4, rite)));

    let gen = generate(move |(), out| {
        fn walk(tree: &Tree, out: &Yield<(), i32>) {
            match tree {
                &None => (),
                &Some(ref node) => {
                    walk(&node.0, out);
                    out.yield_(node.1);
                    walk(&node.2, out);
                }
            }
        }

        walk(&tree, &out)
    });

    assert_eq!(gen.collect::<Vec<i32>>(), [1,2,3,4,5,6,7]);
}