use core::marker::PhantomData;
use core::{ptr, mem};
use core::cell::Cell;
use stack;
use debug;
use arch::{self, StackPointer};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum State {
Runnable,
Unavailable
}
#[derive(Debug)]
pub struct Generator<'a, Input: 'a, Output: 'a, Stack: stack::Stack> {
state: State,
stack: NoDrop<Stack>,
stack_id: NoDrop<debug::StackId>,
stack_ptr: arch::StackPointer,
phantom: PhantomData<(&'a (), *mut Input, *const Output)>
}
#[allow(unions_with_drop_fields)]
union NoDrop<T> {
inner: T
}
impl<T: ::core::fmt::Debug> ::core::fmt::Debug for NoDrop<T> {
fn fmt(&self, w: &mut ::core::fmt::Formatter) -> ::core::fmt::Result {
unsafe {
self.inner.fmt(w)
}
}
}
impl<'a, Input, Output, Stack> Generator<'a, Input, Output, Stack>
where Input: 'a, Output: 'a, Stack: stack::Stack {
pub fn new<F>(stack: Stack, f: F) -> Generator<'a, Input, Output, Stack>
where Stack: stack::GuardedStack + 'static,
F: FnOnce(&Yielder<Input, Output>, Input) + 'a {
unsafe { Generator::unsafe_new(stack, f) }
}
pub unsafe fn unsafe_new<F>(stack: Stack, f: F) -> Generator<'a, Input, Output, Stack>
where F: FnOnce(&Yielder<Input, Output>, Input) + 'a {
unsafe extern "C" fn generator_wrapper<Input, Output, Stack, F>(env: usize, stack_ptr: StackPointer) -> !
where Stack: stack::Stack, F: FnOnce(&Yielder<Input, Output>, Input) {
let f = ptr::read(env as *const F);
let (data, stack_ptr) = arch::swap(0, stack_ptr, None);
let input = ptr::read(data as *const Input);
let yielder = Yielder::new(stack_ptr);
f(&yielder, input);
loop { yielder.suspend_bare(None); }
}
let stack_id = debug::StackId::register(&stack);
let stack_ptr = arch::init(&stack, generator_wrapper::<Input, Output, Stack, F>);
let stack_ptr = arch::swap(&f as *const F as usize, stack_ptr, Some(&stack)).1;
mem::forget(f);
Generator {
state: State::Runnable,
stack: NoDrop { inner: stack },
stack_id: NoDrop { inner: stack_id },
stack_ptr: stack_ptr,
phantom: PhantomData
}
}
#[inline]
pub fn resume(&mut self, input: Input) -> Option<Output> {
match self.state {
State::Runnable => {
self.state = State::Unavailable;
let val = unsafe {
let (data_out, stack_ptr) = arch::swap(&input as *const Input as usize, self.stack_ptr, Some(&self.stack.inner));
self.stack_ptr = stack_ptr;
mem::forget(input);
ptr::read(data_out as *const Option<Output>)
};
if val.is_some() { self.state = State::Runnable }
val
}
State::Unavailable => None
}
}
#[inline]
pub fn state(&self) -> State { self.state }
pub fn unwrap(self) -> Stack {
match self.state {
State::Runnable => {
mem::forget(self);
panic!("Argh! Bastard! Don't touch that!")
}
State::Unavailable => unsafe { self.unsafe_unwrap() }
}
}
pub unsafe fn unsafe_unwrap(mut self) -> Stack {
ptr::drop_in_place(&mut self.stack_id.inner);
let stack = ptr::read(&mut self.stack.inner);
mem::forget(self);
stack
}
}
impl<'a, Input, Output, Stack> Drop for Generator<'a, Input, Output, Stack>
where Input: 'a, Output: 'a, Stack: stack::Stack {
fn drop(&mut self) {
unsafe {
ptr::drop_in_place(&mut self.stack_id.inner);
match self.state {
State::Runnable => panic!("dropped unfinished Generator"),
State::Unavailable => ptr::drop_in_place(&mut self.stack.inner)
}
}
}
}
#[derive(Debug)]
pub struct Yielder<Input, Output> {
stack_ptr: Cell<StackPointer>,
phantom: PhantomData<(*const Input, *mut Output)>
}
impl<Input, Output> Yielder<Input, Output> {
fn new(stack_ptr: StackPointer) -> Yielder<Input, Output> {
Yielder {
stack_ptr: Cell::new(stack_ptr),
phantom: PhantomData
}
}
#[inline(always)]
fn suspend_bare(&self, val: Option<Output>) -> Input {
unsafe {
let (data, stack_ptr) = arch::swap(&val as *const Option<Output> as usize, self.stack_ptr.get(), None);
self.stack_ptr.set(stack_ptr);
mem::forget(val);
ptr::read(data as *const Input)
}
}
#[inline(always)]
pub fn suspend(&self, item: Output) -> Input {
self.suspend_bare(Some(item))
}
}
impl<'a, Output, Stack> Iterator for Generator<'a, (), Output, Stack>
where Output: 'a, Stack: stack::Stack {
type Item = Output;
fn next(&mut self) -> Option<Self::Item> { self.resume(()) }
}