more_itertools/grouping/
distribute.rs1use 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 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}