asynciter 0.1.0

use std::{fmt::Debug, future::Future};

use crate::AsyncIterator;

#[must_use = "iterators are lazy and do nothing unless consumed"]
pub struct Peekable<I>
where
    I: AsyncIterator,
{
    inner: I,
    /// Remember a peeked value, even if it was None.
    peeked: Option<Option<I::Item>>,
}

impl<I: AsyncIterator + Debug> Debug for Peekable<I>
where
    I::Item: Debug,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Peekable")
            .field("peeked", &self.peeked)
            .field("inner", &self.inner)
            .finish()
    }
}

impl<I: AsyncIterator + Clone> std::clone::Clone for Peekable<I>
where
    I::Item: Clone,
{
    fn clone(&self) -> Self {
        Self {
            peeked: self.peeked.clone(),
            inner: self.inner.clone(),
        }
    }

    fn clone_from(&mut self, source: &Self) {
        self.inner = source.inner.clone();
        self.peeked = source.peeked.clone();
    }
}

impl<I> Peekable<I>
where
    I: AsyncIterator,
{
    #[inline]
    pub fn new(inner: I) -> Peekable<I> {
        Self {
            inner,
            peeked: None,
        }
    }

    #[inline]
    pub async fn peek(&mut self) -> Option<&I::Item> {
        if self.peeked.is_none() {
            self.peeked = Some(self.inner.next().await);
        }

        // SAFETY: a `None` variant for `self` would have been replaced by a `Some`
        // variant in the code above.
        unsafe { self.peeked.as_mut().unwrap_unchecked() }.as_ref()
    }

    #[inline]
    pub async fn peek_mut(&mut self) -> Option<&mut I::Item> {
        if self.peeked.is_none() {
            self.peeked = Some(self.inner.next().await);
        }

        // SAFETY: a `None` variant for `self` would have been replaced by a `Some`
        // variant in the code above.
        unsafe { self.peeked.as_mut().unwrap_unchecked() }.as_mut()
    }

    pub async fn next_if(&mut self, func: impl FnOnce(&I::Item) -> bool) -> Option<I::Item> {
        match self.next().await {
            Some(matched) if func(&matched) => Some(matched),
            other => {
                // Since we called `self.next()`, we consumed `self.peeked`.
                assert!(self.peeked.is_none());
                self.peeked = Some(other);
                None
            }
        }
    }

    pub async fn anext_if<F>(&mut self, func: impl FnOnce(&I::Item) -> F) -> Option<I::Item>
    where
        F: Future<Output = bool>,
    {
        match self.next().await {
            Some(matched) if func(&matched).await => Some(matched),
            other => {
                // Since we called `self.next()`, we consumed `self.peeked`.
                assert!(self.peeked.is_none());
                self.peeked = Some(other);
                None
            }
        }
    }

    pub async fn next_if_eq<T>(&mut self, expected: &T) -> Option<I::Item>
    where
        T: ?Sized,
        I::Item: PartialEq<T>,
    {
        self.next_if(|next| next == expected).await
    }
}

impl<I> AsyncIterator for Peekable<I>
where
    I: AsyncIterator,
{
    type Item = <I as AsyncIterator>::Item;

    async fn next(&mut self) -> Option<Self::Item> {
        match self.peeked.take() {
            Some(v) => v,
            None => self.inner.next().await,
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let peek_len = match self.peeked {
            Some(None) => return (0, Some(0)),
            Some(Some(_)) => 1,
            None => 0,
        };
        let (lo, hi) = self.inner.size_hint();
        let lo = lo.saturating_add(peek_len);
        let hi = match hi {
            Some(x) => x.checked_add(peek_len),
            None => None,
        };
        (lo, hi)
    }
}