gat_lending_iterator/adapters/
take.rs

1use crate::LendingIterator;
2
3/// A Lending iterator that only lends the first `n` iterations of `iter`.
4#[derive(Clone, Debug)]
5#[must_use = "iterators are lazy and do nothing unless consumed"]
6pub struct Take<I> {
7    iter: I,
8    n: usize,
9}
10
11impl<I> Take<I>
12where
13    I: LendingIterator,
14{
15    pub(crate) fn new(iter: I, n: usize) -> Take<I> {
16        Take { iter, n }
17    }
18}
19
20impl<I> LendingIterator for Take<I>
21where
22    I: LendingIterator,
23{
24    type Item<'a>
25        = I::Item<'a>
26    where
27        I: 'a;
28
29    #[inline]
30    #[allow(clippy::if_not_else)]
31    fn next(&mut self) -> Option<Self::Item<'_>> {
32        if self.n != 0 {
33            self.n -= 1;
34            self.iter.next()
35        } else {
36            None
37        }
38    }
39}
40
41#[cfg(test)]
42mod test {
43    use crate::{LendingIterator, ToLendingIterator};
44
45    fn identity(x: i32) -> i32 {
46        x
47    }
48
49    #[test]
50    fn test() {
51        assert_eq!(
52            std::iter::repeat(())
53                .into_lending()
54                .take(5)
55                .fold(0, |count, ()| { count + 1 }),
56            5
57        );
58    }
59
60    #[test]
61    fn take_basic() {
62        let result: Vec<_> = (0..10)
63            .into_lending()
64            .take(3)
65            .map(identity)
66            .into_iter()
67            .collect();
68        assert_eq!(result, vec![0, 1, 2]);
69    }
70
71    #[test]
72    fn take_more_than_available() {
73        let result: Vec<_> = (0..3)
74            .into_lending()
75            .take(10)
76            .map(identity)
77            .into_iter()
78            .collect();
79        assert_eq!(result, vec![0, 1, 2]);
80    }
81}