lender 0.6.2

use core::{num::NonZeroUsize, ops::ControlFlow};

use crate::{
    DoubleEndedLender, ExactSizeLender, FusedLender, Lend, Lender, Lending, try_trait_v2::Try,
};

/// A lender that skips over the first `n` elements of the underlying lender.
///
/// This `struct` is created by the [`skip()`](crate::Lender::skip) method on
/// [`Lender`].
#[derive(Clone, Debug)]
#[must_use = "lenders are lazy and do nothing unless consumed"]
pub struct Skip<L> {
    pub(crate) lender: L,
    pub(crate) n: usize,
}

impl<L> Skip<L> {
    /// Returns the inner lender.
    #[inline(always)]
    pub fn into_inner(self) -> L {
        self.lender
    }

    /// Returns the inner lender and the remaining number of elements to skip.
    #[inline(always)]
    pub fn into_parts(self) -> (L, usize) {
        (self.lender, self.n)
    }
}

impl<L: Lender> Skip<L> {
    #[inline(always)]
    pub(crate) fn new(lender: L, n: usize) -> Skip<L> {
        let _ = L::__check_covariance(crate::CovariantProof::new());
        Skip { lender, n }
    }
}

impl<'lend, L> Lending<'lend> for Skip<L>
where
    L: Lender,
{
    type Lend = Lend<'lend, L>;
}

impl<L> Lender for Skip<L>
where
    L: Lender,
{
    // SAFETY: the lend is that of L
    crate::unsafe_assume_covariance!();
    #[inline]
    fn next(&mut self) -> Option<Lend<'_, Self>> {
        if self.n > 0 {
            self.lender.nth(core::mem::take(&mut self.n))
        } else {
            self.lender.next()
        }
    }

    #[inline]
    fn nth(&mut self, n: usize) -> Option<Lend<'_, Self>> {
        if self.n > 0 {
            let skip = core::mem::take(&mut self.n);
            let n = match skip.checked_add(n) {
                Some(nth) => nth,
                None => {
                    self.lender.nth(skip - 1)?;
                    n
                }
            };
            self.lender.nth(n)
        } else {
            self.lender.nth(n)
        }
    }

    #[inline]
    fn count(mut self) -> usize {
        if self.n > 0 && self.lender.nth(self.n - 1).is_none() {
            return 0;
        }
        self.lender.count()
    }

    #[inline]
    fn last(&mut self) -> Option<Lend<'_, Self>>
    where
        Self: Sized,
    {
        if self.n > 0 {
            self.lender.nth(self.n - 1)?;
            self.n = 0;
        }
        self.lender.last()
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let (lower, upper) = self.lender.size_hint();

        let lower = lower.saturating_sub(self.n);
        let upper = upper.map(|x| x.saturating_sub(self.n));

        (lower, upper)
    }

    #[inline]
    fn try_fold<B, F, R>(&mut self, init: B, f: F) -> R
    where
        Self: Sized,
        F: FnMut(B, Lend<'_, Self>) -> R,
        R: Try<Output = B>,
    {
        let n = self.n;
        self.n = 0;
        if n > 0 && self.lender.nth(n - 1).is_none() {
            return R::from_output(init);
        }
        self.lender.try_fold(init, f)
    }

    #[inline]
    fn fold<B, F>(mut self, init: B, f: F) -> B
    where
        Self: Sized,
        F: FnMut(B, Lend<'_, Self>) -> B,
    {
        if self.n > 0 && self.lender.nth(self.n - 1).is_none() {
            return init;
        }
        self.lender.fold(init, f)
    }

    #[inline]
    fn advance_by(&mut self, mut n: usize) -> Result<(), NonZeroUsize> {
        let skip_inner = self.n;
        let skip_and_advance = skip_inner.saturating_add(n);

        let remainder = match self.lender.advance_by(skip_and_advance) {
            Ok(()) => 0,
            Err(n) => n.get(),
        };
        let advanced_inner = skip_and_advance - remainder;
        n -= advanced_inner.saturating_sub(skip_inner);
        self.n = self.n.saturating_sub(advanced_inner);

        if remainder == 0 && n > 0 {
            n = match self.lender.advance_by(n) {
                Ok(()) => 0,
                Err(n) => n.get(),
            }
        }

        NonZeroUsize::new(n).map_or(Ok(()), Err)
    }
}

impl<L> ExactSizeLender for Skip<L> where L: ExactSizeLender {}

impl<L> DoubleEndedLender for Skip<L>
where
    L: DoubleEndedLender + ExactSizeLender,
{
    #[inline]
    fn next_back(&mut self) -> Option<Lend<'_, Self>> {
        if self.len() > 0 {
            self.lender.next_back()
        } else {
            None
        }
    }

    #[inline]
    fn nth_back(&mut self, n: usize) -> Option<Lend<'_, Self>> {
        let len = self.len();
        if len > n {
            self.lender.nth_back(n)
        } else {
            if len > 0 {
                self.lender.nth_back(len - 1);
            }
            None
        }
    }

    #[inline]
    fn try_rfold<B, F, R>(&mut self, init: B, mut f: F) -> R
    where
        Self: Sized,
        F: FnMut(B, Lend<'_, Self>) -> R,
        R: Try<Output = B>,
    {
        let mut len = self.len();
        if len == 0 {
            R::from_output(init)
        } else {
            match self.lender.try_rfold(init, move |acc, x| {
                len -= 1;
                let r = f(acc, x);
                if len == 0 {
                    ControlFlow::Break(r)
                } else {
                    match r.branch() {
                        ControlFlow::Continue(r) => ControlFlow::Continue(r),
                        ControlFlow::Break(r) => ControlFlow::Break(R::from_residual(r)),
                    }
                }
            }) {
                ControlFlow::Continue(r) => R::from_output(r),
                ControlFlow::Break(r) => r,
            }
        }
    }

    #[inline]
    fn advance_back_by(&mut self, n: usize) -> Result<(), NonZeroUsize> {
        let min = core::cmp::min(self.len(), n);
        let rem = self.lender.advance_back_by(min);
        assert!(rem.is_ok(), "ExactSizeLender contract violation");
        NonZeroUsize::new(n - min).map_or(Ok(()), Err)
    }
}

impl<L> FusedLender for Skip<L> where L: FusedLender {}