diff --git a/src/stack/generator.rs b/src/stack/generator.rs index 46aa6e5..05e0802 100644 --- a/src/stack/generator.rs +++ b/src/stack/generator.rs @@ -1,10 +1,10 @@ +use std::{future::Future, mem::MaybeUninit, pin::Pin, ptr}; + use crate::{ core::{advance, async_advance, Airlock as _, Next}, - ext::MaybeUninitExt, ops::{Coroutine, GeneratorState}, stack::engine::{Airlock, Co}, }; -use std::{future::Future, mem, pin::Pin, ptr}; /// This data structure holds the transient state of an executing generator. /// @@ -12,14 +12,9 @@ use std::{future::Future, mem, pin::Pin, ptr}; /// `GeneratorState` enum. /// /// [_See the module-level docs for examples._](.) -// Safety: The lifetime of the data is controlled by a `Gen`, which constructs -// it in place, and holds a mutable reference right up until dropping it in -// place. Thus, the data inside is pinned and can never be moved. -pub struct Shelf(mem::MaybeUninit>); - -struct State { +pub struct Shelf { airlock: Airlock, - future: F, + future: MaybeUninit, } impl Shelf { @@ -28,7 +23,13 @@ impl Shelf { /// [_See the module-level docs for examples._](.) #[must_use] pub fn new() -> Self { - Self(mem::MaybeUninit::uninit()) + Self { + airlock: Airlock::default(), + // Safety: The lifetime of the data is controlled by a `Gen`, which constructs + // it in place, and holds a mutable reference right up until dropping it in + // place. Thus, the data inside is pinned and can never be moved. + future: MaybeUninit::uninit(), + } } } @@ -43,7 +44,8 @@ impl Default for Shelf { /// /// [_See the module-level docs for examples._](.) pub struct Gen<'s, Y, R, F: Future> { - state: Pin<&'s mut State>, + airlock: &'s Airlock, + future: Pin<&'s mut F>, } impl<'s, Y, R, F: Future> Gen<'s, Y, R, F> { @@ -89,19 +91,22 @@ impl<'s, Y, R, F: Future> Gen<'s, Y, R, F> { shelf: &'s mut Shelf, producer: impl FnOnce(Co<'s, Y, R>) -> F, ) -> Self { - // Safety: Build the struct in place, by writing each field in place. - let p = &mut *shelf.0.as_mut_ptr() as *mut State; - - let airlock = Airlock::default(); - ptr::write(&mut (*p).airlock, airlock); - - let future = producer(Co::new(&(*p).airlock)); - ptr::write(&mut (*p).future, future); - - // Safety: the state can never be moved again, because we store it inside a - // `Pin` until `Gen::drop`, where the contents are dropped in place. - let state = Pin::new_unchecked(shelf.0.assume_init_get_mut()); - Self { state } + // By splitting the mutable `shelf` into a shared `airlock` and a unique + // pinned `future` reference we ensure the aliasing rules are not violated. + let airlock = &shelf.airlock; + // Safety: Initializes the future in-place using `ptr::write`, which is + // the correct way to initialize a `MaybeUninit` + shelf.future.as_mut_ptr().write(producer(Co::new(airlock))); + // Safety: The `MaybeUninit` is initialized by now, so its safe to create + // a reference to the future itself + // NB: can be replaced by `MaybeUninit::get_mut` once stabilized + let init = &mut *shelf.future.as_mut_ptr(); + // Safety: The `shelf` remains borrowed during the entire lifetime of + // the `Gen`and is hence pinned. + Self { + airlock, + future: Pin::new_unchecked(init), + } } /// Resumes execution of the generator. @@ -115,37 +120,22 @@ impl<'s, Y, R, F: Future> Gen<'s, Y, R, F> { /// /// [_See the module-level docs for examples._](.) pub fn resume_with(&mut self, arg: R) -> GeneratorState { - let (future, airlock) = self.project(); - airlock.replace(Next::Resume(arg)); - advance(future, &airlock) - } - - fn project(&mut self) -> (Pin<&mut F>, &Airlock) { - unsafe { - // Safety: This is a pin projection. `future` is pinned, but never moved. - // `airlock` is never pinned. - let state = self.state.as_mut().get_unchecked_mut(); - - let future = Pin::new_unchecked(&mut state.future); - let airlock = &state.airlock; - (future, airlock) - } + self.airlock.replace(Next::Resume(arg)); + advance(self.future.as_mut(), &self.airlock) } } impl<'s, Y, R, F: Future> Drop for Gen<'s, Y, R, F> { fn drop(&mut self) { - // Safety: `state` is a `MaybeUninit` which is guaranteed to be initialized, - // because the only way to construct a `Gen` is with `Gen::new`, which - // initializes it. + // Safety: `future` itself is a `MaybeUninit`, which is guaranteed to be + // initialized, because the only way to construct a `Gen` is with + // `Gen::new`, which initializes it. // - // Drop `state` in place, by dropping each field in place. Drop `future` first, - // since it likely contains a reference to `airlock` (through the `co` object). - // Since we drop everything in place, the `Pin` invariants are not violated. + // Drop `future` in place first (it likely contains a reference to airlock). + // Since we drop it in place, the `Pin` invariants are not violated. + // The airlock is regularly dropped when the `Shelf` goes out of scope. unsafe { - let state = self.state.as_mut().get_unchecked_mut(); - ptr::drop_in_place(&mut state.future); - ptr::drop_in_place(&mut state.airlock); + ptr::drop_in_place(self.future.as_mut().get_unchecked_mut()); } } } @@ -171,9 +161,8 @@ impl<'s, Y, F: Future> Gen<'s, Y, (), F> { pub fn async_resume( &mut self, ) -> impl Future> + '_ { - let (future, airlock) = self.project(); - airlock.replace(Next::Resume(())); - async_advance(future, airlock) + self.airlock.replace(Next::Resume(())); + async_advance(self.future.as_mut(), self.airlock) } } diff --git a/tests/stack.rs b/tests/stack.rs index 1cccc03..4883d76 100644 --- a/tests/stack.rs +++ b/tests/stack.rs @@ -117,6 +117,23 @@ fn stack_yield_match() { assert_eq!(vec![1, 3, 5, 7, 9], res) } +#[test] +fn stack_yield_closure_no_macro() { + let mut shelf = genawaiter::stack::Shelf::new(); + let gen = unsafe { + Gen::new(&mut shelf, |co| async move { + let mut n = 1; + while n < 10 { + co.yield_(n).await; + n += 2; + } + }) + }; + + let res = gen.into_iter().collect::>(); + assert_eq!(vec![1, 3, 5, 7, 9], res) +} + #[cfg(feature = "proc_macro")] #[test] fn stack_yield_closure() {