more_itertools/grouping/
split_at.rs

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