generator_extensions 0.1.1

//
// Copyright (C) 2020 Nathan Sharp.
//
// This file is available under either the terms of the Apache License, Version
// 2.0 or the MIT License, at your discretion.
//

use core::fmt::{self, Debug, Formatter};
use core::marker::PhantomData;
use core::mem;
use core::num::NonZeroUsize;
use core::ops::{Generator, GeneratorState};
use core::pin::Pin;

use crate::NoData;

/// This trait provides extension methods for generators behind references which
/// allow [resumption].
///
/// This trait is blanket implemented for all suitable references to a
/// generator, namely [`Pin`]`<&mut `[`Generator`]`>` and
/// `&mut `[`Generator`]` + `[`Unpin`].
///
/// To take advantage of these extension methods, simply write:
/// ```
/// use generator_extensions::Resumable;
/// ```
/// Or, preferably:
/// ```
/// use generator_extensions::prelude::*;
/// ```
///
/// [`Generator`]: core::ops::Generator
/// [`Pin`]: core::pin::Pin
/// [resumption]: core::ops::Generator::resume
/// [`Unpin`]: core::marker::Unpin
pub trait Resumable<'a, R, G: Generator<R>> {
    /// Reborrow the pinned generator.
    ///
    /// This method serves the same role as [`Pin::as_mut`].
    ///
    /// [`Pin::as_mut`]: core::pin::Pin::as_mut
    #[must_use]
    fn as_mut(&mut self) -> Pin<&mut G>;

    /// [`Resume`] the pinned generator.
    ///
    /// [`Resume`]: core::ops::Generator::resume
    #[must_use]
    fn resume(&mut self, arg: R) -> GeneratorState<G::Yield, G::Return> {
        self.as_mut().resume(arg)
    }

    /// [`Resume`] the pinned generator and panic if it does not [yield].
    ///
    /// [`Resume`]: core::ops::Generator::resume
    /// [yield]: core::ops::GeneratorState::Yielded
    fn expect_yield(&mut self, arg: R) -> G::Yield {
        match self.resume(arg) {
            GeneratorState::Yielded(item) => item,
            GeneratorState::Complete(..) => panic!("generator completed"),
        }
    }

    /// [`Resume`] the pinned generator and panic if it does not [complete].
    ///
    /// [complete]: core::ops::GeneratorState::Complete
    /// [`Resume`]: core::ops::Generator::resume
    fn expect_complete(&mut self, arg: R) -> G::Return {
        match self.resume(arg) {
            GeneratorState::Yielded(..) => panic!("generator yielded"),
            GeneratorState::Complete(item) => item,
        }
    }

    /// Transforms a pinned generator into an [`Iterator`].
    ///
    /// This method is only available on generators which take `()` as the
    /// argument to [`resume`], and which return `()` on completion. To
    /// transform a generator into this form, consider using the [`unify`]
    /// or [`states`] extensions.
    ///
    /// [`unify`]: Resumable::unify
    /// [`Iterator`]: core::iter::Iterator
    /// [`resume`]: core::ops::Generator::resume
    /// [`states`]: Resumable::states
    #[must_use]
    fn iter(self) -> Iter<'a, G>
    where G: Generator<(), Return = ()>;

    /// Produces a generator which will repeatedly return
    /// [`GeneratorState::Complete`]`(())` instead of panicking, provided that
    /// the generator *has not already completed*.
    ///
    /// The generator type must already [return] `()`.
    ///
    /// [`GeneratorState::Complete`]: core::ops::GeneratorState::Complete
    /// [return]: core::ops::Generator::Return
    #[must_use]
    fn fuse(self) -> Fuse<'a, R, G>
    where G: Generator<R, Return = ()>;

    /// Produces a generator which yields its [states] and then completes with
    /// `()`.
    ///
    /// [states]: core::ops::GeneratorState
    #[must_use]
    fn states(self) -> States<'a, R, G>;

    /// Transforms a pinned generator with identical [`Yield`] and [`Return`]
    /// types into a generator which yields these values and then completes
    /// with `()`.
    ///
    /// [`Yield`]: core::ops::Generator::Yield
    /// [`Return`]: core::ops::Generator::Return
    #[must_use]
    fn unify<T>(self) -> Unify<'a, R, T, G>
    where G: Generator<R, Yield = T, Return = T>;
}

impl<'a, R, G: Generator<R>> Resumable<'a, R, G> for Pin<&'a mut G> {
    fn as_mut(&mut self) -> Pin<&mut G> {
        Pin::as_mut(self)
    }

    fn iter(self) -> Iter<'a, G>
    where G: Generator<(), Return = ()> {
        Iter::new(self)
    }

    fn unify<T>(self) -> Unify<'a, R, T, G>
    where G: Generator<R, Yield = T, Return = T> {
        Unify::new(self)
    }

    fn fuse(self) -> Fuse<'a, R, G>
    where G: Generator<R, Return = ()> {
        Fuse::new(self)
    }

    fn states(self) -> States<'a, R, G> {
        States::new(self)
    }
}

impl<'a, R, G: Generator<R> + Unpin> Resumable<'a, R, G> for &'a mut G {
    fn as_mut(&mut self) -> Pin<&mut G> {
        Pin::new(self)
    }

    fn iter(self) -> Iter<'a, G>
    where G: Generator<(), Return = ()> {
        Iter::new(Pin::new(self))
    }

    fn unify<T>(self) -> Unify<'a, R, T, G>
    where G: Generator<R, Yield = T, Return = T> {
        Unify::new(Pin::new(self))
    }

    fn fuse(self) -> Fuse<'a, R, G>
    where G: Generator<R, Return = ()> {
        Fuse::new(Pin::new(self))
    }

    fn states(self) -> States<'a, R, G> {
        States::new(Pin::new(self))
    }
}

/// The return type of [`Resumable::iter`].
pub struct Iter<'a, G: Generator<()>> {
    pin: Pin<&'a mut G>,
}

impl<'a, G: Generator<(), Return = ()>> Iter<'a, G> {
    /// Creates a new `Iter` from a suitable pinned generator.
    ///
    /// Consider using [`Resumable::iter`] instead.
    #[must_use]
    pub fn new(pin: Pin<&'a mut G>) -> Self {
        Self { pin }
    }

    /// Retrieves the original pinned generator.
    #[must_use]
    pub fn into_inner(self) -> Pin<&'a mut G> {
        self.pin
    }
}

impl<'a, G: Generator<(), Return = ()>> Debug for Iter<'a, G> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        f.debug_struct("Iter").field("ptr", &(self.pin.as_ref().get_ref() as *const G)).finish()
    }
}

impl<'a, G: Generator<(), Return = ()>> Iterator for Iter<'a, G> {
    type Item = G::Yield;

    fn next(&mut self) -> Option<Self::Item> {
        match self.pin.as_mut().resume(()) {
            GeneratorState::Yielded(item) => Some(item),
            GeneratorState::Complete(()) => None,
        }
    }
}

#[repr(transparent)]
struct TaggedPin<'a, G> {
    ptr: NonZeroUsize,
    _data: PhantomData<Pin<&'a mut G>>,
}

impl<'a, G> TaggedPin<'a, G> {
    #[must_use]
    fn new(pin: Pin<&'a mut G>) -> Self {
        assert!(mem::align_of::<G>() > 1);

        // Safety:
        //     get_unchecked_mut(): We don't move out of G or expose an unpinned
        //     reference.
        //     NonZeroUSize: References are never null, so the pointer will be nonzero.
        unsafe {
            Self {
                ptr: NonZeroUsize::new_unchecked(pin.get_unchecked_mut() as *mut G as usize),
                _data: PhantomData,
            }
        }
    }

    fn get_ptr(&self) -> *mut G {
        (self.ptr.get() & !1) as *mut G
    }

    fn get_pin(&mut self) -> Pin<&'a mut G> {
        // Safety: The returned reference was previously pinned.
        unsafe { Pin::new_unchecked(&mut *self.get_ptr()) }
    }

    fn get_tag(&self) -> bool {
        (self.ptr.get() & 1) != 0
    }

    fn set_tag(&mut self) {
        // Safety: We are setting a bit, so the result cannot be zero. Additionally,
        //         get_ptr() will clear this bit so the returned pointer will remain
        //         valid to dereference.
        self.ptr = unsafe { NonZeroUsize::new_unchecked(self.ptr.get() | 1) };
    }

    fn into_inner(mut self) -> Pin<&'a mut G> {
        self.get_pin()
    }
}

/// The return type of [`Resumable::states`].
pub struct States<'a, R, G: Generator<R>> {
    pin: TaggedPin<'a, G>,
    _arg: NoData<R>,
}

impl<'a, R, G: Generator<R>> States<'a, R, G> {
    /// Creates a new `States` generator from a pinned generator.
    ///
    /// Consider using [`Resumable::states`] instead.
    #[must_use]
    pub fn new(pin: Pin<&'a mut G>) -> Self {
        Self { pin: TaggedPin::new(pin), _arg: NoData::new() }
    }

    /// Retrieves the original pinned generator.
    #[must_use]
    pub fn into_inner(self) -> Pin<&'a mut G> {
        self.pin.into_inner()
    }
}

impl<'a, R, G: Generator<R>> Debug for States<'a, R, G> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        f.debug_struct("States").field("ptr", &self.pin.get_ptr()).finish()
    }
}

impl<'a, R, G: Generator<R>> Generator<R> for States<'a, R, G> {
    type Yield = GeneratorState<G::Yield, G::Return>;
    type Return = ();

    fn resume(mut self: Pin<&mut Self>, arg: R) -> GeneratorState<Self::Yield, Self::Return> {
        if self.pin.get_tag() {
            GeneratorState::Complete(())
        } else {
            match self.pin.get_pin().resume(arg) {
                yielded @ GeneratorState::Yielded(..) => GeneratorState::Yielded(yielded),
                complete @ GeneratorState::Complete(..) => {
                    self.pin.set_tag();
                    GeneratorState::Yielded(complete)
                }
            }
        }
    }
}

/// The return type of [`Resumable::unify`].
pub struct Unify<'a, R, T, G: Generator<R, Yield = T, Return = T>> {
    pin: TaggedPin<'a, G>,
    _arg: NoData<R>,
}

impl<'a, R, T, G: Generator<R, Yield = T, Return = T>> Unify<'a, R, T, G> {
    /// Creates a new `Unify` generator from a pinned generator.
    ///
    /// Consider using [`Resumable::unify`] instead.
    #[must_use]
    pub fn new(pin: Pin<&'a mut G>) -> Self {
        Self { pin: TaggedPin::new(pin), _arg: NoData::new() }
    }

    /// Retrieves the original pinned generator.
    #[must_use]
    pub fn into_inner(self) -> Pin<&'a mut G> {
        self.pin.into_inner()
    }
}

impl<'a, R, T, G: Generator<R, Yield = T, Return = T>> Debug for Unify<'a, R, T, G> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        f.debug_struct("Unify").field("ptr", &self.pin.get_ptr()).finish()
    }
}

impl<'a, R, T, G: Generator<R, Yield = T, Return = T>> Generator<R> for Unify<'a, R, T, G> {
    type Yield = T;
    type Return = ();

    fn resume(mut self: Pin<&mut Self>, arg: R) -> GeneratorState<Self::Yield, Self::Return> {
        if self.pin.get_tag() {
            GeneratorState::Complete(())
        } else {
            match self.pin.get_pin().resume(arg) {
                GeneratorState::Yielded(item) => GeneratorState::Yielded(item),
                GeneratorState::Complete(item) => {
                    self.pin.set_tag();
                    GeneratorState::Yielded(item)
                }
            }
        }
    }
}

/// The return type of [`Resumable::fuse`].
pub struct Fuse<'a, R, G: Generator<R>> {
    pin: TaggedPin<'a, G>,
    _arg: NoData<R>,
}

impl<'a, R, G: Generator<R, Return = ()>> Fuse<'a, R, G> {
    /// Creates a new `Fuse` generator from a pinned generator.
    ///
    /// Consider using [`Resumable::fuse`] instead.
    #[must_use]
    pub fn new(pin: Pin<&'a mut G>) -> Self {
        Self { pin: TaggedPin::new(pin), _arg: NoData::new() }
    }

    /// Returns `true` if the pinned generator has completed.
    ///
    /// # Notes
    /// This method may return `false` even if the generator is complete if
    /// [`new`] was called with a generator which was already complete.
    ///
    /// [`new`]: Fuse::new
    #[must_use]
    pub fn is_complete(&self) -> bool {
        self.pin.get_tag()
    }

    /// Forces the `Fuse` to repeatedly return
    /// [`GeneratorState::Complete`]`(())`, even if the underlying pinned
    /// generator is not yet complete.
    ///
    /// The underlying pinned generator is neither made to complete nor dropped.
    ///
    /// [`GeneratorState::Complete`]: core::ops::GeneratorState::Complete
    pub fn stop(&mut self) {
        self.pin.set_tag()
    }

    /// Retrieves the original pinned generator.
    #[must_use]
    pub fn into_inner(self) -> Pin<&'a mut G> {
        self.pin.into_inner()
    }
}

impl<'a, R, G: Generator<R, Return = ()>> Debug for Fuse<'a, R, G> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        f.debug_struct("Fuse")
            .field("fuse", &self.pin.get_tag())
            .field("ptr", &self.pin.get_ptr())
            .finish()
    }
}

impl<'a, R, G: Generator<R, Return = ()>> Generator<R> for Fuse<'a, R, G> {
    type Yield = G::Yield;
    type Return = ();

    fn resume(mut self: Pin<&mut Self>, arg: R) -> GeneratorState<Self::Yield, Self::Return> {
        if self.pin.get_tag() {
            GeneratorState::Complete(())
        } else {
            match self.pin.get_pin().resume(arg) {
                yielded @ GeneratorState::Yielded(..) => yielded,
                complete @ GeneratorState::Complete(..) => {
                    self.pin.set_tag();
                    complete
                }
            }
        }
    }
}