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
extern crate smallvec;
use std::iter::{DoubleEndedIterator, ExactSizeIterator};
use smallvec::SmallVec;
pub trait IteratorExt: Iterator + Sized {
fn peekable_nth(self) -> PeekableNth<Self>;
}
#[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,
{
#[inline]
pub fn peek(&mut self) -> Option<&I::Item> {
self.peek_nth(0)
}
#[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))
}
}
}