more_itertools/grouping/
split_when.rs

1use std::collections::LinkedList;
2use crate::error::{self, Error};
3
4struct SplitWhenOutputItem<T> {
5    items: Vec<T>,
6    error: Option<Error>,
7    finished: bool
8}
9
10impl<T> SplitWhenOutputItem<T> {
11    pub fn new(items: Vec<T>, error: Option<Error>, finished: bool) -> Self {
12        return SplitWhenOutputItem {
13            items,
14            error,
15            finished
16        };
17    }
18}
19
20pub struct SplitWhen<T> {
21    ret_buf: LinkedList<SplitWhenOutputItem<T>>,
22    iter: Box<dyn Iterator<Item = Result<T,Error>>>,
23    pred: fn(&T, &T) -> Result<bool, Error>,
24    maxsplit: i128,
25    splited: i128,
26    iter_finished: bool
27}
28
29impl<T> SplitWhen<T> {
30    pub fn set_last_output_item_finished(&mut self) {
31        if self.ret_buf.len() > 0 {
32            self.ret_buf.back_mut().unwrap().finished = true;
33        }
34    }
35}
36
37impl<T> Iterator for SplitWhen<T>
38{
39    type Item = Result<Vec<T>, Error>;
40
41    fn next(&mut self) -> Option<Self::Item> {
42        loop {
43            // flush the cached data
44            while self.ret_buf.len() > 0 {
45                if self.ret_buf.front().unwrap().finished {
46                    let ret = self.ret_buf.pop_front();
47                    if ret.as_ref().unwrap().error.is_none() {
48                        return Some(Ok(ret.unwrap().items)); 
49                    } else {
50                        return Some(Err(ret.unwrap().error.unwrap()));
51                    }
52                } else {
53                    break; // jump to consume logic
54                }
55            }
56
57            // consume the iterator
58            if self.iter_finished {
59                return None;
60            }
61
62            if let Some(v) = self.iter.next() {
63                match v {
64                    Ok(ok_v) => {
65                        if self.ret_buf.back_mut().unwrap().items.len() < 1 {
66                            self.ret_buf.back_mut().unwrap().items.push(ok_v);
67                            continue;
68                        }
69
70                        // main logic
71                        let pred_ret = (self.pred)(self.ret_buf.back_mut().unwrap().items.last().unwrap(), &ok_v);
72                        match pred_ret {
73                            Ok(ok_pred_ret) => {
74                                if ok_pred_ret && (self.maxsplit < 0 || self.splited < self.maxsplit) {
75                                    self.splited += 1;
76                                    self.set_last_output_item_finished();
77                                    self.ret_buf.push_back(SplitWhenOutputItem::new(vec![ok_v], None, false));
78                                } else {
79                                    self.ret_buf.back_mut().unwrap().items.push(ok_v);
80                                }
81                            },
82                            Err(err_pred_ret) => {
83                                self.set_last_output_item_finished();
84                                self.ret_buf.push_back(SplitWhenOutputItem::new(Vec::new(), Some(error::any_error(err_pred_ret.kind(), "[split_when] ".to_string()+ err_pred_ret.message().unwrap())), true));
85                                self.iter_finished = true;
86                            }
87                        }
88                    },
89                    Err(err_v) => { // upstream error
90                        self.set_last_output_item_finished();
91                        self.ret_buf.push_back(SplitWhenOutputItem::new(Vec::new(), Some(err_v), true));
92                        self.iter_finished = true;
93                    }
94                }
95            } else {
96                self.set_last_output_item_finished();
97                self.iter_finished = true;
98                continue;
99            }
100        }
101    }
102}
103
104pub fn split_when<T>(iter: Box<dyn Iterator<Item = Result<T,Error>>>, pred: fn(&T, &T) -> Result<bool, Error>, maxsplit: i128) -> Box<dyn Iterator<Item = Result<Vec<T>,Error>>>
105where
106T: 'static
107{
108    let mut ret = SplitWhen {
109        ret_buf: LinkedList::new(),
110        iter,
111        pred,
112        maxsplit,
113        splited: 0,
114        iter_finished: false
115    };
116
117    ret.ret_buf.push_back(SplitWhenOutputItem {
118        items: Vec::new(),
119        error: None,
120        finished: false
121    });
122
123    return Box::new(ret);
124}
125
126#[cfg(test)]
127mod tests {
128    use crate::utils::{extract_value_from_result_vec, generate_okok_iterator, generate_okokerr_iterator};
129
130    use super::*;
131
132    #[test]
133    fn test1() {
134        let v = vec![1, 2, 3, 3, 2, 5, 2, 4, 2];
135        let sw = split_when(generate_okok_iterator(v), |x, y| { Ok(x < y) }, -1);
136        let ret = extract_value_from_result_vec(sw.collect::<Vec<_>>());
137        // println!("{:?}", ret);
138        assert_eq!(vec![vec![1], vec![2], vec![3, 3, 2], vec![5, 2], vec![4, 2]], ret.0);
139
140
141        let v = vec![1, 2, 3, 3, 2, 5, 2, 4, 2];
142        let sw = split_when(generate_okok_iterator(v), |x, y| { Ok(x < y) }, 2);
143        let ret = extract_value_from_result_vec(sw.collect::<Vec<_>>());
144        assert_eq!(vec![vec![1], vec![2], vec![3, 3, 2, 5, 2, 4, 2]], ret.0);
145
146        let v = vec![1, 2, 3, 3, 2, 5, 2, 4, 2];
147        let sw = split_when(generate_okok_iterator(v), |x, y| { Ok(x < y) }, 10);
148        let ret = extract_value_from_result_vec(sw.collect::<Vec<_>>());
149        // println!("{:?}", ret);
150        assert_eq!(vec![vec![1], vec![2], vec![3, 3, 2], vec![5, 2], vec![4, 2]], ret.0);
151
152        let v = vec![1, 2, 3, 3, 2, 5, 2, 4, 2];
153        let sw = split_when(generate_okok_iterator(v), |x, y| { Ok(x < y) }, 0);
154        let ret = extract_value_from_result_vec(sw.collect::<Vec<_>>());
155        // println!("{:?}", ret);
156        assert_eq!(vec![vec![1, 2, 3, 3, 2, 5, 2, 4, 2]], ret.0);
157    }
158
159    #[test]
160    fn test1_error() {
161        let v = vec![1, 2, 3, 3, 2, 5, 2, 4, 2];
162        let sw = split_when(generate_okokerr_iterator(v, error::overflow_error("[test]".to_string())), |x, y| { Ok(x < y) }, -1);
163        let ret = extract_value_from_result_vec(sw.collect::<Vec<_>>());
164        // println!("{:?}", ret);
165        assert_eq!(vec![vec![1], vec![2], vec![3, 3, 2], vec![5, 2], vec![4, 2]], ret.0);
166        assert_eq!(error::Kind::OverflowError, ret.1.unwrap().kind());
167
168
169        let v = vec![1, 2, 3, 3, 2, 5, 2, 4, 2];
170        let sw = split_when(generate_okokerr_iterator(v, error::overflow_error("[test]".to_string())), 
171                                                        |x, y| { if *y == 5 { Err(error::overflow_error("[test]".to_string())) } else { Ok(x < y)} }, -1);
172        let ret = extract_value_from_result_vec(sw.collect::<Vec<_>>());
173        // println!("{:?}", ret);
174        assert_eq!(vec![vec![1], vec![2], vec![3, 3, 2]], ret.0);
175        assert_eq!(error::Kind::OverflowError, ret.1.unwrap().kind());
176    }
177}