more_itertools/itertools/
groupby.rs

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

more_itertools/itertools/
groupby.rs
