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))
        }
    }
}