gat-lending-iterator 0.1.8

use crate::LendingIterator;

/// A lending iterator for stepping lending iterators by a custom amount.
///
/// This `struct` is created by the [`step_by`] method on [`LendingIterator`]. See
/// its documentation for more.
///
/// [`LendingIterator`]: crate::LendingIterator
/// [`step_by`]: crate::LendingIterator::step_by
#[derive(Clone, Debug)]
#[must_use = "iterators are lazy and do nothing unless consumed"]
pub struct StepBy<I> {
    iter: I,
    step: usize,
    first_take: bool,
}

impl<I> StepBy<I> {
    #[inline]
    pub(crate) fn new(iter: I, step: usize) -> StepBy<I> {
        assert!(step != 0);
        StepBy {
            iter,
            step: step - 1,
            first_take: true,
        }
    }
}

impl<I> LendingIterator for StepBy<I>
where
    I: LendingIterator,
{
    type Item<'a>
        = I::Item<'a>
    where
        Self: 'a;

    #[inline]
    fn next(&mut self) -> Option<Self::Item<'_>> {
        if self.first_take {
            self.first_take = false;
            self.iter.next()
        } else {
            self.iter.nth(self.step)
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::{LendingIterator, ToLendingIterator};

    fn identity(x: i32) -> i32 {
        x
    }

    #[test]
    fn step_by_basic() {
        let result: Vec<_> = (0..10)
            .into_lending()
            .step_by(3)
            .map(identity)
            .into_iter()
            .collect();
        assert_eq!(result, vec![0, 3, 6, 9]);
    }
}