diff --git a/src/groupbylazy.rs b/src/groupbylazy.rs index 99778a87c..7e9ff3567 100644 --- a/src/groupbylazy.rs +++ b/src/groupbylazy.rs @@ -575,6 +575,17 @@ where fn drop_group(&self, client: usize) { self.inner.borrow_mut().drop_group(client); } + + fn complete(&self, client: usize) { + let mut inner = self.inner.borrow_mut(); + if client >= inner.oldest_buffered_group + && (client < inner.top_group + || (client == inner.top_group + && inner.buffer.len() > inner.top_group - inner.bottom_group)) + { + inner.lookup_buffer(client); + } + } } impl<'a, I> IntoIterator for &'a IntoChunks @@ -625,9 +636,11 @@ where let index = self.parent.index.get(); self.parent.index.set(index + 1); let inner = &mut *self.parent.inner.borrow_mut(); + let size = inner.key.size; inner.step(index).map(|elt| Chunk { parent: self.parent, index, + remaining: size - 1, first: Some(elt), }) } @@ -644,6 +657,7 @@ where { parent: &'a IntoChunks, index: usize, + remaining: usize, first: Option, } @@ -668,6 +682,11 @@ where if let elt @ Some(..) = self.first.take() { return elt; } + if self.remaining == 0 { + self.parent.complete(self.index); + return None; + } + self.remaining -= 1; self.parent.step(self.index) } } diff --git a/tests/test_std.rs b/tests/test_std.rs index 6800166f9..1127b7c2e 100644 --- a/tests/test_std.rs +++ b/tests/test_std.rs @@ -1075,6 +1075,102 @@ fn chunks() { } } +#[test] +fn chunks_preserve_source_position_and_live_groups() { + let mut source = 1..=10; + { + let chunks = source.by_ref().chunks(2); + let mut groups = chunks.into_iter(); + assert_eq!(groups.next().unwrap().collect_vec(), vec![1, 2]); + } + assert_eq!(source.next(), Some(3)); + + for size in [1, 2, 3] { + let mut source = 1..=20; + { + let chunks = source.by_ref().chunks(size); + let mut groups = chunks.into_iter(); + assert_eq!(groups.next().unwrap().count(), size); + assert_eq!(groups.next().unwrap().count(), size); + } + assert_eq!(source.next(), Some((2 * size + 1) as i32)); + } + + let chunks = (1..=8).chunks(2); + let mut groups = chunks.into_iter(); + let first = groups.next().unwrap(); + let second = groups.next().unwrap(); + assert_eq!(second.collect_vec(), vec![3, 4]); + assert_eq!(first.collect_vec(), vec![1, 2]); + + let chunks = (1..=8).chunks(2); + let mut groups = chunks.into_iter(); + let mut first = groups.next().unwrap(); + assert_eq!(first.next(), Some(1)); + drop(first); + assert_eq!(groups.next().unwrap().collect_vec(), vec![3, 4]); +} + +#[test] +fn chunks_release_buffered_slots() { + for pairs in [10, 100] { + let chunks = (0..(pairs * 4 + 2)).chunks(2); + let mut groups = chunks.into_iter(); + for _ in 0..pairs { + let oldest = groups.next().unwrap(); + let newest = groups.next().unwrap(); + oldest.for_each(drop); + newest.for_each(drop); + } + let debug = format!("{chunks:?}"); + assert!(debug.contains("buffer: [")); + assert!(debug.matches("IntoIter").count() <= 2); + } +} + +#[test] +fn chunks_stop_after_capacity_and_source_exhaustion() { + let chunks = (1..=5).chunks(2); + let mut groups = chunks.into_iter(); + let mut first = groups.next().unwrap(); + assert_eq!(first.next(), Some(1)); + assert_eq!(first.next(), Some(2)); + assert_eq!(first.next(), None); + assert_eq!(first.next(), None); + assert_eq!(groups.next().unwrap().collect_vec(), vec![3, 4]); + assert_eq!(groups.next().unwrap().collect_vec(), vec![5]); + assert!(groups.next().is_none()); +} + +#[test] +fn chunks_do_not_repoll_a_non_fused_source_after_none() { + struct PostNone { + state: usize, + } + + impl Iterator for PostNone { + type Item = usize; + + fn next(&mut self) -> Option { + let result = match self.state { + 0 => Some(1), + 1 => None, + _ => Some(2), + }; + self.state += 1; + result + } + } + + let chunks = PostNone { state: 0 }.chunks(3); + let mut groups = chunks.into_iter(); + let mut first = groups.next().unwrap(); + assert_eq!(first.next(), Some(1)); + assert_eq!(first.next(), None); + assert_eq!(first.next(), None); + assert!(groups.next().is_none()); +} + #[test] fn concat_empty() { let data: Vec> = Vec::new();