peek-nth 0.2.0

An iterator adapter that allows you to efficiently peek the nth item of an iterator.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124

//! An iterator adapter that allows you to efficiently peek the nth item of an iterator.
//!
//! Itermediate values are memoized and heap allocations are avoided when possible.
//!
//! ## Usage
//!
//! ```rust
//! extern crate peek_nth;
//!
//! use peek_nth::IteratorExt;
//!
//! fn main() {
//!     let mut iter = "Hello, world!".chars().peekable_nth();
//!
//!     assert_eq!(iter.peek_nth(4), Some(&'o'));
//!     assert_eq!(iter.peek_nth(3), Some(&'l'));
//!     assert_eq!(iter.peek_nth(2), Some(&'l'));
//!     assert_eq!(iter.peek_nth(1), Some(&'e'));
//!     assert_eq!(iter.peek_nth(0), Some(&'H'));
//!     assert_eq!(iter.peek_nth(7), Some(&'w'));
//!     assert_eq!(iter.collect::<String>(), "Hello, world!");
//! }
//!```

extern crate smallvec;

use std::iter::{DoubleEndedIterator, ExactSizeIterator};

use smallvec::SmallVec;

/// Adds a peekable_nth() method to types that implement [`std::iter::Iterator`].
///
/// [`std::iter::Iterator`]: https://doc.rust-lang.org/std/iter/trait.Iterator.html
pub trait IteratorExt: Iterator + Sized {
    fn peekable_nth(self) -> PeekableNth<Self>;
}

/// An iterator with a peek_nth() method that returns an optional reference to the nth element.
#[derive(Clone, Debug)]
pub struct PeekableNth<I>
where
    I: Iterator,
{
    iter: I,
    next: SmallVec<[I::Item; 64]>,
}

impl<I> IteratorExt for I
where
    I: Iterator,
{
    #[inline]
    fn peekable_nth(self) -> PeekableNth<I> {
        PeekableNth {
            iter: self,
            next: SmallVec::new(),
        }
    }
}

impl<I> PeekableNth<I>
where
    I: Iterator,
{
    /// Returns a reference to the next value without advancing the iterator.
    #[inline]
    pub fn peek(&mut self) -> Option<&I::Item> {
        self.peek_nth(0)
    }

    /// Returns a reference to the nth(n) value without advancing the iterator.
    #[inline]
    pub fn peek_nth(&mut self, n: usize) -> Option<&I::Item> {
        let length = self.next.len();
        let offset = n + 1;

        if offset > length {
            for _ in length..offset {
                self.next.push(self.iter.next()?);
            }
        }

        self.next.get(n)
    }
}

impl<I> DoubleEndedIterator for PeekableNth<I>
where
    I: DoubleEndedIterator,
{
    #[inline]
    fn next_back(&mut self) -> Option<I::Item> {
        match self.iter.next_back() {
            None if !self.next.is_empty() => self.next.pop(),
            option => option,
        }
    }
}

impl<I> ExactSizeIterator for PeekableNth<I>
where
    I: ExactSizeIterator,
{
    #[inline]
    fn len(&self) -> usize {
        self.iter.len()
    }
}

impl<I> Iterator for PeekableNth<I>
where
    I: Iterator,
{
    type Item = I::Item;

    #[inline]
    fn next(&mut self) -> Option<I::Item> {
        if self.next.is_empty() {
            self.iter.next()
        } else {
            Some(self.next.remove(0))
        }
    }
}