more_itertools/itertools/
tee.rs

1use std::{cell::RefCell, collections::VecDeque, rc::Rc};
2use crate::error::Error;
3
4#[allow(dead_code)]
5struct TeeInner<T> 
6where
7T: Clone
8{
9    buf1: VecDeque<T>,
10    buf2: VecDeque<T>,
11    iter: Box<dyn Iterator<Item = Result<T, Error>>>,
12    iter_finished: bool,
13    iter_error: Option<Error>
14}
15
16impl<T> TeeInner<T> 
17where
18T: Clone
19{
20    pub fn new(iter: Box<dyn Iterator<Item = Result<T, Error>>>) -> TeeInner<T> {
21        return TeeInner {
22            buf1: VecDeque::new(),
23            buf2: VecDeque::new(),
24            iter,
25            iter_finished: false,
26            iter_error: None
27        };
28    }
29}
30
31pub struct Tee<T> 
32where
33T: Clone 
34{
35    inner: Rc<RefCell<TeeInner<T>>>
36}
37
38struct TeeCursor<T> 
39where
40T: Clone
41{
42    no: usize,
43    inner: Rc<RefCell<TeeInner<T>>>
44}
45
46impl<T> Iterator for TeeCursor<T> 
47where
48T: Clone
49{
50    type Item = Result<T, Error>;
51
52    fn next(&mut self) -> Option<Self::Item> {
53        let _next = self.inner.borrow_mut().iter.next();
54        match _next {
55            None => {
56                self.inner.borrow_mut().iter_finished = true;
57            }, 
58            Some(v) => {
59                if v.is_ok() {
60                    self.inner.borrow_mut().buf1.push_back(v.as_ref().ok().unwrap().clone());
61                    self.inner.borrow_mut().buf2.push_back(v.as_ref().ok().unwrap().clone());
62                } else {
63                    self.inner.borrow_mut().iter_error = Some(v.err().unwrap());
64                }
65            }
66        }
67        if self.no == 1 {
68            if self.inner.borrow_mut().iter_finished && self.inner.borrow_mut().buf1.len() == 0 {
69                return None;
70            } else {
71                return Some(Ok(self.inner.borrow_mut().buf1.pop_front().unwrap()));
72            }
73        } else if self.no == 2 {
74            if self.inner.borrow_mut().iter_finished && self.inner.borrow_mut().buf2.len() == 0 {
75                return None;
76            } else {
77                return Some(Ok(self.inner.borrow_mut().buf2.pop_front().unwrap()));
78            }
79        } else {
80            return None;
81        }
82    }
83}
84
85impl<T> Tee<T> 
86where
87T: Clone + 'static
88{
89    pub fn new(iter: Box<dyn Iterator<Item=Result<T,Error>>>) -> Tee<T> {
90        let inner = TeeInner::new(iter);
91
92        let ret = Tee {
93            inner: Rc::new(RefCell::new(inner))
94        };
95
96        return ret;
97    }
98
99    pub fn iter(&self) -> (Box<dyn Iterator<Item=Result<T,Error>>>, Box<dyn Iterator<Item=Result<T,Error>>>) {
100        let ret0: Box<dyn Iterator<Item=Result<T,Error>>> = Box::new(TeeCursor {
101            no: 1,
102            inner: Rc::clone(&self.inner)
103        });
104
105        let ret1: Box<dyn Iterator<Item=Result<T,Error>>> = Box::new(TeeCursor {
106            no: 2,
107            inner: Rc::clone(&self.inner)
108        });
109
110        return (ret0, ret1);
111    }
112}
113
114
115
116pub fn tee<T>(iterator: Box<dyn Iterator<Item=Result<T,Error>>>) -> (Box<dyn Iterator<Item=Result<T,Error>>>, Box<dyn Iterator<Item=Result<T,Error>>>) 
117where
118T: Clone + 'static
119{
120    let t = Tee::new(iterator);
121    return t.iter();
122}
123
124#[cfg(test)]
125mod tests {
126    use std::vec;
127
128    use crate::utils::extract_value_from_result_vec;
129    use crate::utils::generate_okok_iterator;
130
131    use super::*;
132
133    #[test]
134    fn test1() {
135        let v = generate_okok_iterator(vec![1,2,3,4,5]);
136        let (t1,t2) = tee(v);
137        let ret1 = extract_value_from_result_vec(t1.collect::<Vec<_>>());
138        assert_eq!(vec![1, 2, 3, 4, 5], ret1.0);
139        assert_eq!(vec![1, 2, 3, 4, 5], extract_value_from_result_vec(t2.collect::<Vec<_>>()).0);
140
141        let v = generate_okok_iterator(vec![1,2,3,4,5]);
142        let (mut t1, mut t2) = tee(v);
143        assert_eq!(Some(1), t1.next().unwrap().ok());
144        assert_eq!(Some(1), t2.next().unwrap().ok());
145        assert_eq!(Some(2), t1.next().unwrap().ok());
146        assert_eq!(Some(2), t2.next().unwrap().ok());
147        assert_eq!(Some(3), t1.next().unwrap().ok());
148        assert_eq!(Some(3), t2.next().unwrap().ok());
149        assert_eq!(Some(4), t1.next().unwrap().ok());
150        assert_eq!(Some(4), t2.next().unwrap().ok());
151        assert_eq!(Some(5), t1.next().unwrap().ok());
152        assert_eq!(Some(5), t2.next().unwrap().ok());
153        assert_eq!(None, t1.next());
154        assert_eq!(None, t2.next());
155
156        let v = generate_okok_iterator(vec![1,2,3,4,5]);
157        let (mut t1, mut t2) = tee(v);
158        assert_eq!(Some(1), t1.next().unwrap().ok());
159        assert_eq!(Some(2), t1.next().unwrap().ok());
160        assert_eq!(Some(1), t2.next().unwrap().ok());
161        assert_eq!(Some(2), t2.next().unwrap().ok());
162
163        assert_eq!(Some(3), t1.next().unwrap().ok());
164        assert_eq!(Some(4), t1.next().unwrap().ok());
165        assert_eq!(Some(3), t2.next().unwrap().ok());
166        assert_eq!(Some(4), t2.next().unwrap().ok());
167
168        assert_eq!(Some(5), t1.next().unwrap().ok());
169        assert_eq!(Some(5), t2.next().unwrap().ok());
170
171        assert_eq!(None, t1.next());
172        assert_eq!(None, t2.next());
173
174        let v = generate_okok_iterator(vec![1,2,3,4]);
175        let (mut t1, mut t2) = tee(v);
176        assert_eq!(Some(1), t1.next().unwrap().ok());
177        assert_eq!(Some(2), t1.next().unwrap().ok());
178        assert_eq!(Some(3), t1.next().unwrap().ok());
179        assert_eq!(Some(1), t2.next().unwrap().ok());
180        assert_eq!(Some(2), t2.next().unwrap().ok());
181        assert_eq!(Some(3), t2.next().unwrap().ok());
182
183        assert_eq!(Some(4), t1.next().unwrap().ok());
184        assert_eq!(Some(4), t2.next().unwrap().ok());
185
186        assert_eq!(None, t1.next());
187        assert_eq!(None, t2.next());
188    }
189}
more_itertools/itertools/tee.rs

more_itertools/itertools/
tee.rs
