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

#[test]
fn test_take() {
    let v = vec![0isize, 1, 2, 3];
    let r: Vec<_> = v.into_iter().take(3).collect();
    assert_eq!(r, vec![0, 1, 2]);

    let v = vec![0isize, 1];
    let r: Vec<_> = v.into_iter().take(3).collect();
    assert_eq!(r, vec![0, 1]);
}

/**
(<em>a</em><sub>0</sub>, ..., <em>a</em><sub><em>i</em>-1</sub>, <em>a</em><sub><em>i</em></sub>, <em>a</em><sub><em>i</em>+1</sub>, ...), <em>i</em>
&nbsp;&rarr;&nbsp;
(<em>a</em><sub>0</sub>, ..., <em>a</em><sub><em>i</em>-1</sub>)

*/
pub trait TakeExactlyIterator<E>: Iterator<Item=E> + Sized {
    /**
Creates an iterator that yields *exactly* `n` elements from the subject iterator.

# Failure

The iterator will panic if there are less than `n` elements in the subject iterator.
    */
    fn take_exactly(self, n: usize) -> TakeExactly<Self> {
        TakeExactly {
            iter: self,
            left: n,
        }
    }
}

impl<It, E> TakeExactlyIterator<E> for It where It: Iterator<Item=E> {}

#[derive(Clone, Debug)]
pub struct TakeExactly<It> {
    iter: It,
    left: usize,
}

impl<It> TakeExactly<It> {
    /**
Unwraps the iterator, returning the underlying iterator.
    */
    pub fn unwrap(self) -> It {
        self.iter
    }
}

impl<It, E> Iterator for TakeExactly<It> where It: Iterator<Item=E> {
    type Item = E;

    fn next(&mut self) -> Option<E> {
        match self.left {
            0 => None,
            _ => match self.iter.next() {
                None => panic!("take_exactly expected {} more elements from iterator, but ran out", self.left),
                e @ _ => {
                    self.left -= 1;
                    e
                }
            }
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.left, Some(self.left))
    }
}

#[test]
fn test_take_exactly() {
    use std::thread;

    let v = vec![0isize, 1, 2, 3];
    let r: Vec<_> = v.into_iter().take_exactly(3).collect();
    assert_eq!(r, vec![0, 1, 2]);

    let v = vec![0isize, 1];
    let r = thread::spawn(move || { v.into_iter().take_exactly(3).collect::<Vec<_>>(); }).join();
    assert!(r.is_err());
}