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::ops::{Generator, GeneratorState};
use core::pin::Pin;

use crate::NoData;

/// This trait represents a value which can be transformed into a [`Generator`].
///
/// Currently, this trait is blanket implemented for all [`Iterator`] types as
/// well as for the [`Gen`] wrapper type.
///
/// To take advantage of these extension methods, simply write:
/// ```
/// use generator_extensions::IntoGenerator;
/// ```
/// Or, preferably:
/// ```
/// use generator_extensions::prelude::*;
/// ```
///
/// # Notes
/// This trait is unfortunately not implemented for compiler-generated types, as
/// those desugared types could theoretically implement multiple distinct
/// generator traits. This does not occur today and is very unlikely to occur in
/// the future, but the Rust trait system is rightly unable to express this
/// fact.
///
/// To work around this, the newtype wrapper [`Gen`] is provided by this crate.
///
/// [`Generator`]: core::ops::Generator
/// [`Iterator`]: core::iter::Iterator
pub trait IntoGenerator {
    /// The argument to [resume][core::ops::Generator::resume] the generator.
    type Resume;
    /// The type of value the generator [yields][core::ops::Generator::Yield].
    type Yield;
    /// The type of value the generator [returns][core::ops::Generator::Return].
    type Return;
    /// Type full type of the generator.
    type Generator: Generator<Self::Resume, Yield = Self::Yield, Return = Self::Return>;

    /// Convert this value into a [`Generator`][core::ops::Generator].
    #[must_use]
    fn into_generator(self) -> Self::Generator;
}

impl<I: Iterator> IntoGenerator for I {
    type Resume = ();
    type Yield = I::Item;
    type Return = ();
    type Generator = GenIter<I>;

    fn into_generator(self) -> Self::Generator {
        GenIter::new(self)
    }
}

/// An [`Iterator`] wrapped to expose a [`Generator`] implementation which
/// yields the iterator's [`Item`] type and then returns `()`.
///
/// [`Generator`]: core::ops::Generator
/// [`Item`]: core::iter::Iterator::Item
/// [`Iterator`]: core::iter::Iterator
#[derive(Clone, Debug)]
#[repr(transparent)]
pub struct GenIter<I: Iterator> {
    iter: I,
}

impl<I: Iterator> GenIter<I> {
    /// Wrap an [`Iterator`] as a [`Generator`].
    ///
    /// Consider using [`IntoGenerator::into_generator`] instead.
    ///
    /// See the [type-level documentation] for more information.
    ///
    /// [`Generator`]: core::ops::Generator
    /// [`Iterator`]: core::iter::Iterator
    /// [type-level documentation]: GenIter
    #[must_use]
    pub fn new(iter: I) -> Self {
        Self { iter }
    }

    /// Retrieves the wrapped [`Iterator`][core::iter::Iterator].
    #[must_use]
    pub fn into_inner(self) -> I {
        self.iter
    }
}

impl<I: Iterator> Generator<()> for GenIter<I> {
    type Yield = I::Item;
    type Return = ();

    fn resume(self: Pin<&mut Self>, _arg: ()) -> GeneratorState<Self::Yield, Self::Return> {
        match unsafe { self.get_unchecked_mut() }.iter.next() {
            Some(item) => GeneratorState::Yielded(item),
            None => GeneratorState::Complete(()),
        }
    }
}

/// A newtype wrapper for built-in generators which provides an implementation
/// of the [`IntoGenerator`] trait.
///
/// See the [`IntoGenerator`] trait documentation for more details.
#[repr(transparent)]
pub struct Gen<R, G: Generator<R>> {
    gen: G,
    _arg: NoData<R>,
}

impl<R, G: Generator<R>> Gen<R, G> {
    /// Wraps a generator.
    #[must_use]
    pub fn new(gen: G) -> Self {
        Self { gen, _arg: NoData::new() }
    }

    /// Projects a pinned `Gen` as its wrapped generator.
    #[allow(clippy::wrong_self_convention)] // Clippy doesn't understand Pin.
    #[must_use]
    pub fn as_pinned(self: Pin<&mut Self>) -> Pin<&mut G> {
        // Safety: `gen` is structurally pinned
        unsafe { self.map_unchecked_mut(|pin| &mut pin.gen) }
    }

    /// Retrieves the wrapped generator.
    #[must_use]
    pub fn into_inner(self) -> G {
        self.gen
    }
}

impl<R, G: Generator<R> + Debug> Debug for Gen<R, G> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        self.gen.fmt(f)
    }
}

impl<R, G: Generator<R>> From<G> for Gen<R, G> {
    fn from(gen: G) -> Self {
        Self::new(gen)
    }
}

impl<R, G: Generator<R>> Generator<R> for Gen<R, G> {
    type Yield = G::Yield;
    type Return = G::Return;

    fn resume(self: Pin<&mut Self>, arg: R) -> GeneratorState<Self::Yield, Self::Return> {
        self.as_pinned().resume(arg)
    }
}

impl<R, G: Generator<R>> IntoGenerator for Gen<R, G> {
    type Resume = R;
    type Yield = G::Yield;
    type Return = G::Return;
    type Generator = G;

    fn into_generator(self) -> Self::Generator {
        self.gen
    }
}