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
#[cfg(feature = "crossbeam")]
extern crate crossbeam_channel;
#[cfg(feature = "crossbeam")]
mod crossbeam;
use std::iter::{self, Iterator};
use std::vec;
use std::sync::mpsc;
pub trait TryIterRecv<'a> {
type Iter: Iterator + 'a;
type Error;
fn try_iter(&'a self) -> Self::Iter;
fn recv(&self) -> Result<<Self::Iter as Iterator>::Item, Self::Error>;
}
pub trait BatchRecv<'a>: TryIterRecv<'a> {
fn batch_recv(&'a self, n: usize) -> Result<BatchIter<Self::Iter>, Self::Error> {
let first = self.recv()?;
let rest = self.try_iter().take(n - 1);
let inner = vec![first].into_iter().chain(rest);
Ok(inner)
}
}
impl<'a, T> BatchRecv<'a> for T
where
T: TryIterRecv<'a>,
{
}
pub type BatchIter<I> = iter::Chain<vec::IntoIter<<I as Iterator>::Item>, iter::Take<I>>;
impl<'a, T: 'a> TryIterRecv<'a> for mpsc::Receiver<T> {
type Iter = mpsc::TryIter<'a, T>;
type Error = mpsc::RecvError;
fn try_iter(&'a self) -> mpsc::TryIter<'a, T> {
mpsc::Receiver::try_iter(&self)
}
fn recv(&self) -> Result<T, Self::Error> {
mpsc::Receiver::recv(&self)
}
}
#[cfg(test)]
mod tests {
use std::sync::mpsc;
use std::thread;
use super::BatchRecv;
#[test]
fn take3of4() {
let (tx, rx) = mpsc::sync_channel(10);
tx.send(1).unwrap();
tx.send(2).unwrap();
tx.send(3).unwrap();
tx.send(4).unwrap();
let first3: Vec<_> = rx.batch_recv(3).unwrap().collect();
assert_eq!(first3, vec![1, 2, 3]);
}
#[test]
fn take2of2() {
let (tx, rx) = mpsc::sync_channel(10);
tx.send(1).unwrap();
tx.send(2).unwrap();
let first2: Vec<_> = rx.batch_recv(3).unwrap().collect();
assert_eq!(first2, vec![1, 2]);
}
#[test]
fn take_first1() {
let (tx, rx) = mpsc::sync_channel(10);
thread::spawn(move || {
thread::yield_now();
tx.send(1).unwrap();
});
let first1: Vec<_> = rx.batch_recv(3).unwrap().collect();
assert_eq!(first1, vec![1]);
}
}