more_itertools/grouping/
distribute.rs

1use std::rc::Rc;
2
3use crate::error;
4use crate::error::Error;
5use crate::sequence::Sequence;
6
7#[allow(dead_code)]
8struct DistributeInner<T> 
9where
10T: Clone + PartialEq + 'static
11{
12    pub(crate) buf: Box<dyn Sequence<T>>,
13    pub(crate) bucket_count: usize,
14}
15
16pub struct Distribute<T> 
17where
18T: Clone + PartialEq + 'static
19{
20    inner: Rc<DistributeInner<T>>
21}
22
23#[allow(dead_code)]
24pub struct Cursor<T>
25where
26T: Clone + PartialEq + 'static
27{
28    dist_inner: Rc<DistributeInner<T>>,
29    cur: usize,
30    step: usize
31}
32
33impl<T> Iterator for Cursor<T>
34where
35T: Clone + PartialEq
36{
37    type Item = Result<T, Error>;
38
39    fn next(&mut self) -> Option<Self::Item> {
40        if self.cur >= self.dist_inner.buf.len() {
41            return None;
42        }
43
44        let ret = usize::overflowing_add(self.cur, self.step);
45        if ret.1 {
46            // cursor overflow
47            return Some(Err(error::overflow_error("cursor overflow".to_string())));
48        }
49
50        let real_ret = Some(Ok(self.dist_inner.buf.get(self.cur).unwrap().clone()));
51        
52        self.cur = ret.0;
53
54        return real_ret;
55    }
56}
57
58impl<T> Distribute<T> 
59where
60T: Clone + PartialEq + 'static
61{
62    pub fn new(buf: Box<dyn Sequence<T>>, bucket_count: usize) -> Self {
63        let inner = DistributeInner {
64            buf: buf,
65            bucket_count: bucket_count,
66        };
67
68        let ret = Distribute {
69            inner: Rc::new(inner)
70        };
71
72        return ret;
73    }
74
75    pub fn iter(&self, bucket_no: usize) -> Box<dyn Iterator<Item = Result<T, Error>>> {
76        
77        let iter: Box<dyn Iterator<Item = Result<T, Error>>> = Box::new(Cursor {
78                dist_inner: Rc::clone(&self.inner),
79                cur: bucket_no,
80                step: self.inner.bucket_count
81            });
82
83        return iter;
84    }
85}
86
87pub fn distribute<T>(buf: Box<dyn Sequence<T>>, bucket_cnt: usize) -> Distribute<T>
88where
89T: Clone + PartialEq + 'static
90{
91    let dist = Distribute::new(buf, bucket_cnt);
92    return dist;
93}
94
95
96#[cfg(test)]
97mod tests {
98    use crate::sequence::create_seq_from_vec;
99
100    use super::*;
101
102    #[test]
103    fn test1() {
104        let v = create_seq_from_vec(vec![1,2,3,4,5,6,7,8,9,10]);
105        let dist = distribute(v, 3);
106
107        let mut cur_0 = dist.iter(0);
108        assert_eq!(1, cur_0.next().unwrap().ok().unwrap());
109        assert_eq!(4, cur_0.next().unwrap().ok().unwrap());
110        assert_eq!(7, cur_0.next().unwrap().ok().unwrap());
111        assert_eq!(10, cur_0.next().unwrap().ok().unwrap());
112        assert_eq!(None, cur_0.next());
113    }
114
115    #[test]
116    fn test2() {
117        let v = create_seq_from_vec(vec![1,2,3,4,5,6,7,8,9,10]);
118        let dist = distribute(v, 3);
119
120        let mut cur_0: Box<dyn Iterator<Item = Result<i32, Error>>> = dist.iter(0);
121        let mut cur_1 = dist.iter(1);
122        let mut cur_2 = dist.iter(2);
123        assert_eq!(1, cur_0.next().unwrap().ok().unwrap());
124        assert_eq!(2, cur_1.next().unwrap().ok().unwrap());
125        assert_eq!(3, cur_2.next().unwrap().ok().unwrap());
126
127
128        assert_eq!(4, cur_0.next().unwrap().ok().unwrap());
129        assert_eq!(5, cur_1.next().unwrap().ok().unwrap());
130        assert_eq!(6, cur_2.next().unwrap().ok().unwrap());
131        
132        
133        assert_eq!(7, cur_0.next().unwrap().ok().unwrap());
134        assert_eq!(8, cur_1.next().unwrap().ok().unwrap());
135        assert_eq!(9, cur_2.next().unwrap().ok().unwrap());
136
137        assert_eq!(10, cur_0.next().unwrap().ok().unwrap());
138        assert_eq!(None, cur_1.next());
139        assert_eq!(None, cur_2.next());
140        
141        assert_eq!(None, cur_0.next());
142    }
143
144    #[test]
145    fn test3() {
146        let v = create_seq_from_vec(vec![1,2,3]);
147        let dist = distribute(v, 5);
148
149        let mut cur_0 = dist.iter(0);
150        let mut cur_1 = dist.iter(1);
151        let mut cur_2 = dist.iter(2);
152        let mut cur_3 = dist.iter(3);
153        let mut cur_4 = dist.iter(4);
154
155        assert_eq!(1, cur_0.next().unwrap().ok().unwrap());
156        assert_eq!(2, cur_1.next().unwrap().ok().unwrap());
157        assert_eq!(3, cur_2.next().unwrap().ok().unwrap());
158        assert_eq!(None, cur_3.next());
159        assert_eq!(None, cur_4.next());
160    }
161}