Reimplement Generator to pass both input and output values around.

Author: whitequark

README.md

@@ -13,7 +13,7 @@ without relying on kernel services. It can be used in hosted environments
 (using `std`) as well as on bare metal (using `core`).
 
 It provides the following safe abstractions:
-  * an implementation of internal iterators, also known as generators,
+  * an implementation of generators,
     [Generator](https://nathan7.github.io/libfringe/fringe/generator/struct.Generator.html).
 
 It also provides the necessary low-level building blocks:

benches/generator.rs

@@ -10,9 +10,9 @@ use fringe::{OsStack, Generator};
 #[bench]
 fn generate(b: &mut test::Bencher) {
   let stack = OsStack::new(0).unwrap();
-  let mut gen = Generator::new(stack, move |yielder| {
-    for i in 1.. { yielder.generate(i) }
+  let mut identity = Generator::new(stack, move |yielder, mut input| {
+    loop { input = yielder.generate(input) }
   });
 
-  b.iter(|| test::black_box(gen.next()));
+  b.iter(|| test::black_box(identity.resume(test::black_box(0))));
 }

src/generator.rs

@@ -9,11 +9,10 @@
 //! afterwards.
 
 use core::marker::PhantomData;
-use core::iter::Iterator;
 use core::{ptr, mem};
 
 use stack;
-use context;
+use context::Context;
 
 #[derive(Debug, Clone, Copy)]
 pub enum State {
@@ -24,18 +23,20 @@ pub enum State {
   Unavailable
 }
 
-/// Generator wraps a function and allows suspending its execution more than
-/// once, return a value each time.
+/// Generator wraps a function and allows suspending its execution more than once, returning
+/// a value each time.
 ///
-/// It implements the Iterator trait. The first time `next()` is called,
-/// the function is called as `f(yielder)`; every time `next()` is called afterwards,
-/// the function is resumed. In either case, it runs until it suspends its execution
-/// through `yielder.generate(val)`), in which case `next()` returns `Some(val)`, or
-/// returns, in which case `next()` returns `None`. `next()` will return `None`
-/// every time after that.
+/// The first time `resume(input0)` is called, the function is called as `f(yielder, input0)`.
+/// It runs until it suspends its execution through `yielder.generate(output0)`, after which
+/// `resume(input0)` returns `output0`. The function can be resumed again using `resume(input1)`,
+/// after which `yielder.generate(output0)` returns `input1`, and so on. Once the function returns,
+/// the `resume()` call will return `None`, and it will return `None` every time it is called
+/// after that.
 ///
-/// After the generator function returns, it is safe to reclaim the generator
-/// stack using `unwrap()`.
+/// If the generator function panics, the panic is propagated through the `resume()` call as usual.
+///
+/// After the generator function returns or panics, it is safe to reclaim the generator stack
+/// using `unwrap()`.
 ///
 /// `state()` can be used to determine whether the generator function has returned;
 /// the state is `State::Runnable` after creation and suspension, and `State::Unavailable`
@@ -47,28 +48,29 @@ pub enum State {
 /// use fringe::{OsStack, Generator};
 ///
 /// let stack = OsStack::new(0).unwrap();
-/// let mut gen = Generator::new(stack, move |yielder| {
-///   for i in 1..4 {
-///     yielder.generate(i);
+/// let mut add_one = Generator::new(stack, move |yielder, mut input| {
+///   loop {
+///     if input == 0 { break }
+///     input = yielder.generate(input + 1)
 ///   }
 /// });
-/// println!("{:?}", gen.next()); // prints Some(1)
-/// println!("{:?}", gen.next()); // prints Some(2)
-/// println!("{:?}", gen.next()); // prints Some(3)
-/// println!("{:?}", gen.next()); // prints None
+/// println!("{:?}", add_one.resume(2)); // prints Some(3)
+/// println!("{:?}", add_one.resume(3)); // prints Some(4)
+/// println!("{:?}", add_one.resume(0)); // prints None
 /// ```
 #[derive(Debug)]
-pub struct Generator<Item: Send, Stack: stack::Stack> {
+pub struct Generator<Input: Send, Output: Send, Stack: stack::Stack> {
   state:   State,
-  context: context::Context<Stack>,
-  phantom: PhantomData<*const Item>
+  context: Context<Stack>,
+  phantom: (PhantomData<*const Input>, PhantomData<*const Output>)
 }
 
-impl<Item, Stack> Generator<Item, Stack>
-    where Item: Send, Stack: stack::Stack {
+impl<Input, Output, Stack> Generator<Input, Output, Stack>
+    where Input: Send, Output: Send, Stack: stack::Stack {
   /// Creates a new generator.
-  pub fn new<F>(stack: Stack, f: F) -> Generator<Item, Stack>
-      where Stack: stack::GuardedStack, F: FnOnce(&mut Yielder<Item, Stack>) + Send {
+  pub fn new<F>(stack: Stack, f: F) -> Generator<Input, Output, Stack>
+      where Stack: stack::GuardedStack,
+            F: FnOnce(&mut Yielder<Input, Output, Stack>, Input) + Send {
     unsafe { Generator::unsafe_new(stack, f) }
   }
 
@@ -77,36 +79,70 @@ impl<Item, Stack> Generator<Item, Stack>
   /// This function is unsafe because the generator function can easily violate
   /// memory safety by overflowing the stack. It is useful in environments where
   /// guarded stacks do not exist, e.g. in absence of an MMU.
-  pub unsafe fn unsafe_new<F>(stack: Stack, f: F) -> Generator<Item, Stack>
-      where F: FnOnce(&mut Yielder<Item, Stack>) + Send {
-    unsafe extern "C" fn generator_wrapper<Item, Stack, F>(info: usize) -> !
-        where Item: Send, Stack: stack::Stack, F: FnOnce(&mut Yielder<Item, Stack>) {
+  pub unsafe fn unsafe_new<F>(stack: Stack, f: F) -> Generator<Input, Output, Stack>
+      where F: FnOnce(&mut Yielder<Input, Output, Stack>, Input) + Send {
+    unsafe extern "C" fn generator_wrapper<Input, Output, Stack, F>(env: usize) -> !
+        where Input: Send, Output: Send, Stack: stack::Stack,
+              F: FnOnce(&mut Yielder<Input, Output, Stack>, Input) {
       // Retrieve our environment from the callee and return control to it.
-      let (mut yielder, f) = ptr::read(info as *mut (Yielder<Item, Stack>, F));
-      let new_context = context::Context::swap(yielder.context, yielder.context, 0);
-      // See Yielder::return_.
-      yielder.context = new_context as *mut context::Context<Stack>;
+      let (mut yielder, f) = ptr::read(env as *mut (Yielder<Input, Output, Stack>, F));
+      let data = Context::swap(yielder.context, yielder.context, 0);
+      // See the second half of Yielder::generate_bare.
+      let (new_context, input) = ptr::read(data as *mut (*mut Context<Stack>, Input));
+      yielder.context = new_context as *mut Context<Stack>;
       // Run the body of the generator.
-      f(&mut yielder);
+      f(&mut yielder, input);
       // Past this point, the generator has dropped everything it has held.
-      loop { yielder.return_(None) }
+      loop { yielder.generate_bare(None); }
     }
 
     let mut generator = Generator {
       state:   State::Runnable,
-      context: context::Context::new(stack, generator_wrapper::<Item, Stack, F>),
-      phantom: PhantomData
+      context: Context::new(stack, generator_wrapper::<Input, Output, Stack, F>),
+      phantom: (PhantomData, PhantomData)
     };
 
     // Transfer environment to the callee.
-    let mut data = (Yielder::new(&mut generator.context), f);
-    context::Context::swap(&mut generator.context, &generator.context,
-                           &mut data as *mut (Yielder<Item, Stack>, F) as usize);
-    mem::forget(data);
+    let mut env = (Yielder::new(&mut generator.context), f);
+    Context::swap(&mut generator.context, &generator.context,
+                  &mut env as *mut (Yielder<Input, Output, Stack>, F) as usize);
+    mem::forget(env);
 
     generator
   }
 
+  /// Resumes the generator and return the next value it yields.
+  /// If the generator function has returned, returns `None`.
+  #[inline]
+  pub fn resume(&mut self, input: Input) -> Option<Output> {
+    match self.state {
+      State::Runnable => {
+        // Set the state to Unavailable. Since we have exclusive access to the generator,
+        // the only case where this matters is the generator function panics, after which
+        // it must not be invocable again.
+        self.state = State::Unavailable;
+
+        // Switch to the generator function, and retrieve the yielded value.
+        let val = unsafe {
+          let mut data_in = (&mut self.context as *mut Context<Stack>, input);
+          let data_out =
+            ptr::read(Context::swap(&mut self.context, &self.context,
+                                    &mut data_in as *mut (*mut Context<Stack>, Input)  as usize)
+                      as *mut Option<Output>);
+          mem::forget(data_in);
+          data_out
+        };
+
+        // Unless the generator function has returned, it can be switched to again, so
+        // set the state to Runnable.
+        if val.is_some() { self.state = State::Runnable }
+
+        val
+      }
+      State::Unavailable => None
+    }
+  }
+
   /// Returns the state of the generator.
   #[inline]
   pub fn state(&self) -> State { self.state }
@@ -125,72 +161,41 @@ impl<Item, Stack> Generator<Item, Stack>
 /// Yielder is an interface provided to every generator through which it
 /// returns a value.
 #[derive(Debug)]
-pub struct Yielder<Item: Send, Stack: stack::Stack> {
-  context: *mut context::Context<Stack>,
-  phantom: PhantomData<Item>
+pub struct Yielder<Input: Send, Output: Send, Stack: stack::Stack> {
+  context: *mut Context<Stack>,
+  phantom: (PhantomData<*const Input>, PhantomData<*const Output>)
 }
 
-impl<Item, Stack> Yielder<Item, Stack>
-    where Item: Send, Stack: stack::Stack {
-  fn new(context: *mut context::Context<Stack>) -> Yielder<Item, Stack> {
+impl<Input, Output, Stack> Yielder<Input, Output, Stack>
+    where Input: Send, Output: Send, Stack: stack::Stack {
+  fn new(context: *mut Context<Stack>) -> Yielder<Input, Output, Stack> {
     Yielder {
       context: context,
-      phantom: PhantomData
+      phantom: (PhantomData, PhantomData)
     }
   }
 
   #[inline(always)]
-  fn return_(&mut self, mut val: Option<Item>) {
+  fn generate_bare(&mut self, mut val: Option<Output>) -> Input {
     unsafe {
-      let new_context = context::Context::swap(self.context, self.context,
-                                               &mut val as *mut Option<Item> as usize);
+      let data = Context::swap(self.context, self.context,
+                               &mut val as *mut Option<Output> as usize);
+      let (new_context, input) = ptr::read(data as *mut (*mut Context<Stack>, Input));
       // The generator can be moved (and with it, the context).
       // This changes the address of the context.
       // Thus, we update it after each swap.
-      self.context = new_context as *mut context::Context<Stack>;
-      // However, between this point and the next time we enter return_
+      self.context = new_context;
+      // However, between this point and the next time we enter generate_bare
       // the generator cannot be moved, as a &mut Generator is necessary
       // to resume the generator function.
+      input
     }
   }
 
-  /// Suspends the generator and returns `Some(item)` from the `next()`
+  /// Suspends the generator and returns `Some(item)` from the `resume()`
   /// invocation that resumed the generator.
   #[inline(always)]
-  pub fn generate(&mut self, item: Item) {
-    self.return_(Some(item))
-  }
-}
-
-impl<Item, Stack> Iterator for Generator<Item, Stack>
-    where Item: Send, Stack: stack::Stack {
-  type Item = Item;
-
-  /// Resumes the generator and return the next value it yields.
-  /// If the generator function has returned, returns `None`.
-  #[inline]
-  fn next(&mut self) -> Option<Self::Item> {
-    match self.state {
-      State::Runnable => {
-        // Set the state to Unavailable. Since we have exclusive access to the generator,
-        // the only case where this matters is the generator function panics, after which
-        // it must not be invocable again.
-        self.state = State::Unavailable;
-
-        // Switch to the generator function.
-        let new_context = &mut self.context as *mut context::Context<Stack> as usize;
-        let val = unsafe {
-          ptr::read(context::Context::swap(&mut self.context, &self.context, new_context)
-                    as *mut Option<Item>)
-        };
-
-        // Unless the generator function has returned, it can be switched to again, so
-        // set the state to Runnable.
-        if val.is_some() { self.state = State::Runnable }
-
-        val
-      }
-      State::Unavailable => None
-    }
+  pub fn generate(&mut self, item: Output) -> Input {
+    self.generate_bare(Some(item))
   }
 }

src/lib.rs

@@ -13,7 +13,7 @@
 //!
 //! It provides the following safe abstractions:
 //!
-//!   * an implementation of internal iterators, also known as generators,
+//!   * an implementation of generators,
 //!     [Generator](generator/struct.Generator.html).
 //!
 //! It also provides the necessary low-level building blocks:

tests/fpe.rs

@@ -20,10 +20,10 @@ extern {
 #[ignore]
 fn fpe() {
   let stack = OsStack::new(0).unwrap();
-  let mut gen = Generator::new(stack, move |yielder| {
+  let mut gen = Generator::new(stack, move |yielder, ()| {
     yielder.generate(1.0 / black_box(0.0));
   });
 
   unsafe { feenableexcept(FE_DIVBYZERO); }
-  println!("{:?}", gen.next());
+  println!("{:?}", gen.resume(()));
 }

tests/generator.rs

@@ -6,57 +6,56 @@ extern crate fringe;
 use fringe::OsStack;
 use fringe::generator::Generator;
 
-#[test]
-fn generator() {
+fn new_add_one() -> Generator<i32, i32, OsStack> {
   let stack = OsStack::new(0).unwrap();
-  let mut gen = Generator::new(stack, move |yielder| {
-    for i in 1..4 {
-      yielder.generate(i);
+  Generator::new(stack, move |yielder, mut input| {
+    loop {
+      if input == 0 { break }
+      input = yielder.generate(input + 1)
     }
-  });
-  assert_eq!(gen.next(), Some(1));
-  assert_eq!(gen.next(), Some(2));
-  assert_eq!(gen.next(), Some(3));
-  assert_eq!(gen.next(), None);
+  })
+}
+
+#[test]
+fn generator() {
+  let mut add_one = new_add_one();
+  assert_eq!(add_one.resume(1), Some(2));
+  assert_eq!(add_one.resume(2), Some(3));
+  assert_eq!(add_one.resume(0), None);
 }
 
 #[test]
 fn move_after_new() {
-  let stack = OsStack::new(0).unwrap();
-  let mut gen = Generator::new(stack, move |yielder| {
-    for i in 1..4 {
-      yielder.generate(i);
-    }
-  });
-  assert_eq!(gen.next(), Some(1));
+  let mut add_one = new_add_one();
+  assert_eq!(add_one.resume(1), Some(2));
 
   #[inline(never)]
-  fn rest(mut gen: Generator<u32, OsStack>) {
-    assert_eq!(gen.next(), Some(2));
-    assert_eq!(gen.next(), Some(3));
-    assert_eq!(gen.next(), None);
+  fn run_moved(mut add_one: Generator<i32, i32, OsStack>) {
+    assert_eq!(add_one.resume(2), Some(3));
+    assert_eq!(add_one.resume(3), Some(4));
+    assert_eq!(add_one.resume(0), None);
   }
-  rest(gen);
+  run_moved(add_one);
 }
 
 #[test]
 #[should_panic]
 fn panic_safety() {
   struct Wrapper {
-    gen: Generator<u32, OsStack>
+    gen: Generator<(), (), OsStack>
   }
 
   impl Drop for Wrapper {
     fn drop(&mut self) {
-      self.gen.next();
+      self.gen.resume(());
     }
   }
 
   let stack = OsStack::new(4 << 20).unwrap();
-  let gen = Generator::new(stack, move |_yielder| {
+  let gen = Generator::new(stack, move |_yielder, ()| {
     panic!("foo")
   });
 
   let mut wrapper = Wrapper { gen: gen };
-  wrapper.gen.next();
+  wrapper.gen.resume(());
 }