more_itertools/grouping/
split_into.rs

1use std::collections::LinkedList;
2
3use crate::error::Error;
4
5struct SplitIntoOutputItem<T> {
6    items: Vec<T>,
7    size: usize,
8    finished: bool
9}
10
11pub struct SplitInto<T>
12{
13    ret_buf: LinkedList<SplitIntoOutputItem<T>>,
14    iter: Box<dyn Iterator<Item = Result<T,Error>>>,
15    sizes: Vec<usize>,
16    iter_finished: bool
17}
18
19impl<T> Iterator for SplitInto<T>
20{
21    type Item = Result<Vec<T>, Error>;
22
23    fn next(&mut self) -> Option<Self::Item> {
24        loop {
25            if self.iter_finished {
26                if self.ret_buf.len() > 0 {
27                    let front = self.ret_buf.pop_front();
28                    return Some(Ok(front.unwrap().items));
29                } else {
30                    return None;
31                }
32            }
33
34            if self.ret_buf.len() > 0 {
35                let front = self.ret_buf.front_mut().unwrap();
36                if front.items.len() == front.size {
37                    front.finished = true;
38                }
39
40                if front.finished {
41                    let front = self.ret_buf.pop_front();
42                    return Some(Ok(front.unwrap().items));
43                }
44            }
45
46            let _next = self.iter.next();
47            if let Some(v) = _next {
48                match v {
49                    Ok(ok_v) => {
50                        if self.ret_buf.len() == 0 {
51                            continue;
52                        }
53
54                        let front = self.ret_buf.front_mut();
55                        front.unwrap().items.push(ok_v);
56
57                        continue;
58                    },
59                    Err(err_v) => { // upstream error
60                        self.iter_finished = true;
61                        return Some(Err(err_v));
62                    }
63                }
64            } else {
65                self.iter_finished = true;
66                for item in self.ret_buf.iter_mut() {
67                    item.finished = true;
68                }
69                continue;
70            }
71        }
72    }
73}
74
75pub fn split_into<T>(iter: Box<dyn Iterator<Item = Result<T,Error>>>, sizes: Vec<usize>) -> Box<dyn Iterator<Item = Result<Vec<T>,Error>>>
76where
77T: 'static
78{
79    let mut ret = SplitInto {
80        ret_buf: LinkedList::new(),
81        iter,
82        sizes: sizes.clone(),
83        iter_finished: false
84    };
85
86    for size in ret.sizes.iter() {
87        ret.ret_buf.push_back(SplitIntoOutputItem {
88            finished: false,
89            size: *size,
90            items: Vec::new()
91        });
92    }
93
94    return Box::new(ret);
95}
96
97
98#[cfg(test)]
99mod tests {
100    use crate::{error, utils::{extract_value_from_result_vec, generate_okok_iterator, generate_okokerr_iterator}};
101
102    use super::*;
103
104    #[test]
105    fn test1() {
106        let v = vec![1,2,3,4,5,6];
107        let sizes = vec![1,2,3];
108        let si = split_into(generate_okok_iterator(v), sizes);
109        let ret = si.collect::<Vec<_>>();
110        let ret2 = extract_value_from_result_vec(ret);
111        assert_eq!(vec![vec![1], vec![2, 3], vec![4, 5, 6]], ret2.0);
112
113        let v = vec![1,2,3,4,5,6];
114        let sizes = vec![2,3];
115        let si = split_into(generate_okok_iterator(v), sizes);
116        let ret = si.collect::<Vec<_>>();
117        let ret2 = extract_value_from_result_vec(ret);
118        assert_eq!(vec![vec![1, 2], vec![3, 4, 5]], ret2.0);
119
120        let v = vec![1,2,3,4];
121        let sizes = vec![1,2,3,4];
122        let si = split_into(generate_okok_iterator(v), sizes);
123        let ret = si.collect::<Vec<_>>();
124        let ret2 = extract_value_from_result_vec(ret);
125        assert_eq!(vec![vec![1], vec![2, 3], vec![4], vec![]], ret2.0);
126
127        let v = vec![1,2,3,4];
128        let sizes = vec![1,2,0,3,4];
129        let si = split_into(generate_okok_iterator(v), sizes);
130        let ret = si.collect::<Vec<_>>();
131        let ret2 = extract_value_from_result_vec(ret);
132        assert_eq!(vec![vec![1], vec![2, 3], vec![], vec![4], vec![]], ret2.0);
133    }
134
135    #[test]
136    fn test1_error() {
137        let v = vec![1,2,3,4,5,6];
138        let sizes = vec![1,2,3];
139        let si = split_into(generate_okokerr_iterator(v, error::overflow_error("[test]".to_string())), sizes);
140        let ret = si.collect::<Vec<_>>();
141        let ret2 = extract_value_from_result_vec(ret);
142        assert_eq!(vec![vec![1], vec![2, 3], vec![4, 5, 6]], ret2.0);
143        assert_eq!(error::Kind::OverflowError, ret2.1.unwrap().kind());
144
145        let v = vec![1,2,3,4,5,6];
146        let sizes = vec![1,2,4];
147        let si = split_into(generate_okokerr_iterator(v, error::overflow_error("[test]".to_string())), sizes);
148        let ret = si.collect::<Vec<_>>();
149        let ret2 = extract_value_from_result_vec(ret);
150        assert_eq!(vec![vec![1], vec![2, 3]], ret2.0);
151        assert_eq!(error::Kind::OverflowError, ret2.1.unwrap().kind());
152    }
153
154}
more_itertools/grouping/split_into.rs

more_itertools/grouping/
split_into.rs
