• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 use std::cell::{Cell, RefCell};
2 use alloc::vec::{self, Vec};
3 
4 /// A trait to unify FnMut for GroupBy with the chunk key in IntoChunks
5 trait KeyFunction<A> {
6     type Key;
call_mut(&mut self, arg: A) -> Self::Key7     fn call_mut(&mut self, arg: A) -> Self::Key;
8 }
9 
10 impl<'a, A, K, F: ?Sized> KeyFunction<A> for F
11     where F: FnMut(A) -> K
12 {
13     type Key = K;
14     #[inline]
call_mut(&mut self, arg: A) -> Self::Key15     fn call_mut(&mut self, arg: A) -> Self::Key {
16         (*self)(arg)
17     }
18 }
19 
20 
21 /// ChunkIndex acts like the grouping key function for IntoChunks
22 #[derive(Debug)]
23 struct ChunkIndex {
24     size: usize,
25     index: usize,
26     key: usize,
27 }
28 
29 impl ChunkIndex {
30     #[inline(always)]
new(size: usize) -> Self31     fn new(size: usize) -> Self {
32         ChunkIndex {
33             size,
34             index: 0,
35             key: 0,
36         }
37     }
38 }
39 
40 impl<'a, A> KeyFunction<A> for ChunkIndex {
41     type Key = usize;
42     #[inline(always)]
call_mut(&mut self, _arg: A) -> Self::Key43     fn call_mut(&mut self, _arg: A) -> Self::Key {
44         if self.index == self.size {
45             self.key += 1;
46             self.index = 0;
47         }
48         self.index += 1;
49         self.key
50     }
51 }
52 
53 
54 struct GroupInner<K, I, F>
55     where I: Iterator
56 {
57     key: F,
58     iter: I,
59     current_key: Option<K>,
60     current_elt: Option<I::Item>,
61     /// flag set if iterator is exhausted
62     done: bool,
63     /// Index of group we are currently buffering or visiting
64     top_group: usize,
65     /// Least index for which we still have elements buffered
66     oldest_buffered_group: usize,
67     /// Group index for `buffer[0]` -- the slots
68     /// bottom_group..oldest_buffered_group are unused and will be erased when
69     /// that range is large enough.
70     bottom_group: usize,
71     /// Buffered groups, from `bottom_group` (index 0) to `top_group`.
72     buffer: Vec<vec::IntoIter<I::Item>>,
73     /// index of last group iter that was dropped, usize::MAX == none
74     dropped_group: usize,
75 }
76 
77 impl<K, I, F> GroupInner<K, I, F>
78     where I: Iterator,
79           F: for<'a> KeyFunction<&'a I::Item, Key=K>,
80           K: PartialEq,
81 {
82     /// `client`: Index of group that requests next element
83     #[inline(always)]
step(&mut self, client: usize) -> Option<I::Item>84     fn step(&mut self, client: usize) -> Option<I::Item> {
85         /*
86         println!("client={}, bottom_group={}, oldest_buffered_group={}, top_group={}, buffers=[{}]",
87                  client, self.bottom_group, self.oldest_buffered_group,
88                  self.top_group,
89                  self.buffer.iter().map(|elt| elt.len()).format(", "));
90         */
91         if client < self.oldest_buffered_group {
92             None
93         } else if client < self.top_group ||
94             (client == self.top_group &&
95              self.buffer.len() > self.top_group - self.bottom_group)
96         {
97             self.lookup_buffer(client)
98         } else if self.done {
99             None
100         } else if self.top_group == client {
101             self.step_current()
102         } else {
103             self.step_buffering(client)
104         }
105     }
106 
107     #[inline(never)]
lookup_buffer(&mut self, client: usize) -> Option<I::Item>108     fn lookup_buffer(&mut self, client: usize) -> Option<I::Item> {
109         // if `bufidx` doesn't exist in self.buffer, it might be empty
110         let bufidx = client - self.bottom_group;
111         if client < self.oldest_buffered_group {
112             return None;
113         }
114         let elt = self.buffer.get_mut(bufidx).and_then(|queue| queue.next());
115         if elt.is_none() && client == self.oldest_buffered_group {
116             // FIXME: VecDeque is unfortunately not zero allocation when empty,
117             // so we do this job manually.
118             // `bottom_group..oldest_buffered_group` is unused, and if it's large enough, erase it.
119             self.oldest_buffered_group += 1;
120             // skip forward further empty queues too
121             while self.buffer.get(self.oldest_buffered_group - self.bottom_group)
122                              .map_or(false, |buf| buf.len() == 0)
123             {
124                 self.oldest_buffered_group += 1;
125             }
126 
127             let nclear = self.oldest_buffered_group - self.bottom_group;
128             if nclear > 0 && nclear >= self.buffer.len() / 2 {
129                 let mut i = 0;
130                 self.buffer.retain(|buf| {
131                     i += 1;
132                     debug_assert!(buf.len() == 0 || i > nclear);
133                     i > nclear
134                 });
135                 self.bottom_group = self.oldest_buffered_group;
136             }
137         }
138         elt
139     }
140 
141     /// Take the next element from the iterator, and set the done
142     /// flag if exhausted. Must not be called after done.
143     #[inline(always)]
next_element(&mut self) -> Option<I::Item>144     fn next_element(&mut self) -> Option<I::Item> {
145         debug_assert!(!self.done);
146         match self.iter.next() {
147             None => { self.done = true; None }
148             otherwise => otherwise,
149         }
150     }
151 
152 
153     #[inline(never)]
step_buffering(&mut self, client: usize) -> Option<I::Item>154     fn step_buffering(&mut self, client: usize) -> Option<I::Item> {
155         // requested a later group -- walk through the current group up to
156         // the requested group index, and buffer the elements (unless
157         // the group is marked as dropped).
158         // Because the `Groups` iterator is always the first to request
159         // each group index, client is the next index efter top_group.
160         debug_assert!(self.top_group + 1 == client);
161         let mut group = Vec::new();
162 
163         if let Some(elt) = self.current_elt.take() {
164             if self.top_group != self.dropped_group {
165                 group.push(elt);
166             }
167         }
168         let mut first_elt = None; // first element of the next group
169 
170         while let Some(elt) = self.next_element() {
171             let key = self.key.call_mut(&elt);
172             match self.current_key.take() {
173                 None => {}
174                 Some(old_key) => if old_key != key {
175                     self.current_key = Some(key);
176                     first_elt = Some(elt);
177                     break;
178                 },
179             }
180             self.current_key = Some(key);
181             if self.top_group != self.dropped_group {
182                 group.push(elt);
183             }
184         }
185 
186         if self.top_group != self.dropped_group {
187             self.push_next_group(group);
188         }
189         if first_elt.is_some() {
190             self.top_group += 1;
191             debug_assert!(self.top_group == client);
192         }
193         first_elt
194     }
195 
push_next_group(&mut self, group: Vec<I::Item>)196     fn push_next_group(&mut self, group: Vec<I::Item>) {
197         // When we add a new buffered group, fill up slots between oldest_buffered_group and top_group
198         while self.top_group - self.bottom_group > self.buffer.len() {
199             if self.buffer.is_empty() {
200                 self.bottom_group += 1;
201                 self.oldest_buffered_group += 1;
202             } else {
203                 self.buffer.push(Vec::new().into_iter());
204             }
205         }
206         self.buffer.push(group.into_iter());
207         debug_assert!(self.top_group + 1 - self.bottom_group == self.buffer.len());
208     }
209 
210     /// This is the immediate case, where we use no buffering
211     #[inline]
step_current(&mut self) -> Option<I::Item>212     fn step_current(&mut self) -> Option<I::Item> {
213         debug_assert!(!self.done);
214         if let elt @ Some(..) = self.current_elt.take() {
215             return elt;
216         }
217         match self.next_element() {
218             None => None,
219             Some(elt) => {
220                 let key = self.key.call_mut(&elt);
221                 match self.current_key.take() {
222                     None => {}
223                     Some(old_key) => if old_key != key {
224                         self.current_key = Some(key);
225                         self.current_elt = Some(elt);
226                         self.top_group += 1;
227                         return None;
228                     },
229                 }
230                 self.current_key = Some(key);
231                 Some(elt)
232             }
233         }
234     }
235 
236     /// Request the just started groups' key.
237     ///
238     /// `client`: Index of group
239     ///
240     /// **Panics** if no group key is available.
group_key(&mut self, client: usize) -> K241     fn group_key(&mut self, client: usize) -> K {
242         // This can only be called after we have just returned the first
243         // element of a group.
244         // Perform this by simply buffering one more element, grabbing the
245         // next key.
246         debug_assert!(!self.done);
247         debug_assert!(client == self.top_group);
248         debug_assert!(self.current_key.is_some());
249         debug_assert!(self.current_elt.is_none());
250         let old_key = self.current_key.take().unwrap();
251         if let Some(elt) = self.next_element() {
252             let key = self.key.call_mut(&elt);
253             if old_key != key {
254                 self.top_group += 1;
255             }
256             self.current_key = Some(key);
257             self.current_elt = Some(elt);
258         }
259         old_key
260     }
261 }
262 
263 impl<K, I, F> GroupInner<K, I, F>
264     where I: Iterator,
265 {
266     /// Called when a group is dropped
drop_group(&mut self, client: usize)267     fn drop_group(&mut self, client: usize) {
268         // It's only useful to track the maximal index
269         if self.dropped_group == !0 || client > self.dropped_group {
270             self.dropped_group = client;
271         }
272     }
273 }
274 
275 /// `GroupBy` is the storage for the lazy grouping operation.
276 ///
277 /// If the groups are consumed in their original order, or if each
278 /// group is dropped without keeping it around, then `GroupBy` uses
279 /// no allocations. It needs allocations only if several group iterators
280 /// are alive at the same time.
281 ///
282 /// This type implements `IntoIterator` (it is **not** an iterator
283 /// itself), because the group iterators need to borrow from this
284 /// value. It should be stored in a local variable or temporary and
285 /// iterated.
286 ///
287 /// See [`.group_by()`](../trait.Itertools.html#method.group_by) for more information.
288 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
289 pub struct GroupBy<K, I, F>
290     where I: Iterator,
291 {
292     inner: RefCell<GroupInner<K, I, F>>,
293     // the group iterator's current index. Keep this in the main value
294     // so that simultaneous iterators all use the same state.
295     index: Cell<usize>,
296 }
297 
298 /// Create a new
new<K, J, F>(iter: J, f: F) -> GroupBy<K, J::IntoIter, F> where J: IntoIterator, F: FnMut(&J::Item) -> K,299 pub fn new<K, J, F>(iter: J, f: F) -> GroupBy<K, J::IntoIter, F>
300     where J: IntoIterator,
301           F: FnMut(&J::Item) -> K,
302 {
303     GroupBy {
304         inner: RefCell::new(GroupInner {
305             key: f,
306             iter: iter.into_iter(),
307             current_key: None,
308             current_elt: None,
309             done: false,
310             top_group: 0,
311             oldest_buffered_group: 0,
312             bottom_group: 0,
313             buffer: Vec::new(),
314             dropped_group: !0,
315         }),
316         index: Cell::new(0),
317     }
318 }
319 
320 impl<K, I, F> GroupBy<K, I, F>
321     where I: Iterator,
322 {
323     /// `client`: Index of group that requests next element
step(&self, client: usize) -> Option<I::Item> where F: FnMut(&I::Item) -> K, K: PartialEq,324     fn step(&self, client: usize) -> Option<I::Item>
325         where F: FnMut(&I::Item) -> K,
326               K: PartialEq,
327     {
328         self.inner.borrow_mut().step(client)
329     }
330 
331     /// `client`: Index of group
drop_group(&self, client: usize)332     fn drop_group(&self, client: usize) {
333         self.inner.borrow_mut().drop_group(client)
334     }
335 }
336 
337 impl<'a, K, I, F> IntoIterator for &'a GroupBy<K, I, F>
338     where I: Iterator,
339           I::Item: 'a,
340           F: FnMut(&I::Item) -> K,
341           K: PartialEq
342 {
343     type Item = (K, Group<'a, K, I, F>);
344     type IntoIter = Groups<'a, K, I, F>;
345 
into_iter(self) -> Self::IntoIter346     fn into_iter(self) -> Self::IntoIter {
347         Groups { parent: self }
348     }
349 }
350 
351 
352 /// An iterator that yields the Group iterators.
353 ///
354 /// Iterator element type is `(K, Group)`:
355 /// the group's key `K` and the group's iterator.
356 ///
357 /// See [`.group_by()`](../trait.Itertools.html#method.group_by) for more information.
358 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
359 pub struct Groups<'a, K: 'a, I: 'a, F: 'a>
360     where I: Iterator,
361           I::Item: 'a
362 {
363     parent: &'a GroupBy<K, I, F>,
364 }
365 
366 impl<'a, K, I, F> Iterator for Groups<'a, K, I, F>
367     where I: Iterator,
368           I::Item: 'a,
369           F: FnMut(&I::Item) -> K,
370           K: PartialEq
371 {
372     type Item = (K, Group<'a, K, I, F>);
373 
374     #[inline]
next(&mut self) -> Option<Self::Item>375     fn next(&mut self) -> Option<Self::Item> {
376         let index = self.parent.index.get();
377         self.parent.index.set(index + 1);
378         let inner = &mut *self.parent.inner.borrow_mut();
379         inner.step(index).map(|elt| {
380             let key = inner.group_key(index);
381             (key, Group {
382                 parent: self.parent,
383                 index,
384                 first: Some(elt),
385             })
386         })
387     }
388 }
389 
390 /// An iterator for the elements in a single group.
391 ///
392 /// Iterator element type is `I::Item`.
393 pub struct Group<'a, K: 'a, I: 'a, F: 'a>
394     where I: Iterator,
395           I::Item: 'a,
396 {
397     parent: &'a GroupBy<K, I, F>,
398     index: usize,
399     first: Option<I::Item>,
400 }
401 
402 impl<'a, K, I, F> Drop for Group<'a, K, I, F>
403     where I: Iterator,
404           I::Item: 'a,
405 {
drop(&mut self)406     fn drop(&mut self) {
407         self.parent.drop_group(self.index);
408     }
409 }
410 
411 impl<'a, K, I, F> Iterator for Group<'a, K, I, F>
412     where I: Iterator,
413           I::Item: 'a,
414           F: FnMut(&I::Item) -> K,
415           K: PartialEq,
416 {
417     type Item = I::Item;
418     #[inline]
next(&mut self) -> Option<Self::Item>419     fn next(&mut self) -> Option<Self::Item> {
420         if let elt @ Some(..) = self.first.take() {
421             return elt;
422         }
423         self.parent.step(self.index)
424     }
425 }
426 
427 ///// IntoChunks /////
428 
429 /// Create a new
new_chunks<J>(iter: J, size: usize) -> IntoChunks<J::IntoIter> where J: IntoIterator,430 pub fn new_chunks<J>(iter: J, size: usize) -> IntoChunks<J::IntoIter>
431     where J: IntoIterator,
432 {
433     IntoChunks {
434         inner: RefCell::new(GroupInner {
435             key: ChunkIndex::new(size),
436             iter: iter.into_iter(),
437             current_key: None,
438             current_elt: None,
439             done: false,
440             top_group: 0,
441             oldest_buffered_group: 0,
442             bottom_group: 0,
443             buffer: Vec::new(),
444             dropped_group: !0,
445         }),
446         index: Cell::new(0),
447     }
448 }
449 
450 
451 /// `ChunkLazy` is the storage for a lazy chunking operation.
452 ///
453 /// `IntoChunks` behaves just like `GroupBy`: it is iterable, and
454 /// it only buffers if several chunk iterators are alive at the same time.
455 ///
456 /// This type implements `IntoIterator` (it is **not** an iterator
457 /// itself), because the chunk iterators need to borrow from this
458 /// value. It should be stored in a local variable or temporary and
459 /// iterated.
460 ///
461 /// Iterator element type is `Chunk`, each chunk's iterator.
462 ///
463 /// See [`.chunks()`](../trait.Itertools.html#method.chunks) for more information.
464 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
465 pub struct IntoChunks<I>
466     where I: Iterator,
467 {
468     inner: RefCell<GroupInner<usize, I, ChunkIndex>>,
469     // the chunk iterator's current index. Keep this in the main value
470     // so that simultaneous iterators all use the same state.
471     index: Cell<usize>,
472 }
473 
474 
475 impl<I> IntoChunks<I>
476     where I: Iterator,
477 {
478     /// `client`: Index of chunk that requests next element
step(&self, client: usize) -> Option<I::Item>479     fn step(&self, client: usize) -> Option<I::Item> {
480         self.inner.borrow_mut().step(client)
481     }
482 
483     /// `client`: Index of chunk
drop_group(&self, client: usize)484     fn drop_group(&self, client: usize) {
485         self.inner.borrow_mut().drop_group(client)
486     }
487 }
488 
489 impl<'a, I> IntoIterator for &'a IntoChunks<I>
490     where I: Iterator,
491           I::Item: 'a,
492 {
493     type Item = Chunk<'a, I>;
494     type IntoIter = Chunks<'a, I>;
495 
into_iter(self) -> Self::IntoIter496     fn into_iter(self) -> Self::IntoIter {
497         Chunks {
498             parent: self,
499         }
500     }
501 }
502 
503 
504 /// An iterator that yields the Chunk iterators.
505 ///
506 /// Iterator element type is `Chunk`.
507 ///
508 /// See [`.chunks()`](../trait.Itertools.html#method.chunks) for more information.
509 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
510 pub struct Chunks<'a, I: 'a>
511     where I: Iterator,
512           I::Item: 'a,
513 {
514     parent: &'a IntoChunks<I>,
515 }
516 
517 impl<'a, I> Iterator for Chunks<'a, I>
518     where I: Iterator,
519           I::Item: 'a,
520 {
521     type Item = Chunk<'a, I>;
522 
523     #[inline]
next(&mut self) -> Option<Self::Item>524     fn next(&mut self) -> Option<Self::Item> {
525         let index = self.parent.index.get();
526         self.parent.index.set(index + 1);
527         let inner = &mut *self.parent.inner.borrow_mut();
528         inner.step(index).map(|elt| {
529             Chunk {
530                 parent: self.parent,
531                 index,
532                 first: Some(elt),
533             }
534         })
535     }
536 }
537 
538 /// An iterator for the elements in a single chunk.
539 ///
540 /// Iterator element type is `I::Item`.
541 pub struct Chunk<'a, I: 'a>
542     where I: Iterator,
543           I::Item: 'a,
544 {
545     parent: &'a IntoChunks<I>,
546     index: usize,
547     first: Option<I::Item>,
548 }
549 
550 impl<'a, I> Drop for Chunk<'a, I>
551     where I: Iterator,
552           I::Item: 'a,
553 {
drop(&mut self)554     fn drop(&mut self) {
555         self.parent.drop_group(self.index);
556     }
557 }
558 
559 impl<'a, I> Iterator for Chunk<'a, I>
560     where I: Iterator,
561           I::Item: 'a,
562 {
563     type Item = I::Item;
564     #[inline]
next(&mut self) -> Option<Self::Item>565     fn next(&mut self) -> Option<Self::Item> {
566         if let elt @ Some(..) = self.first.take() {
567             return elt;
568         }
569         self.parent.step(self.index)
570     }
571 }
572